[Quantization] Tighten FP8 sweep input contracts and add dispatch test

Address realAsma's review feedback on the NVFP4 FP8 sweep kernel:

- TritonNVFP4MSECalibrator.collect: replace `assert x.ndim == 2` with
  ValueError so the contract still holds under `python -O`, validate
  block_size > 0 before use, and derive n_blocks from x.shape[0] so a
  zero last-dim cannot trigger division before the shape check.
- nvfp4_fp8_scale_sweep: drop the public `candidates` parameter. The
  candidate set is fixed (FP8 E4M3 valid values / 448) and a wrong
  length would silently inflate `tl.static_range` codegen, while
  nonpositive/nonfinite entries violate the kernel's scale assumptions.
  No internal caller used the override.
- Add test_mse_calibrate_dispatch covering the public default + opt-out
  wiring: confirms `mse_calibrate(fp8_scale_sweep=True)` installs
  TritonNVFP4MSECalibrator by default and falls back to NVFP4MSECalibrator
  when MODELOPT_NVFP4_TRITON_SWEEP=0.

Signed-off-by: Chenjie Luo <chenjiel@nvidia.com>

Author: cjluo-nv

modelopt/torch/kernels/quantization/gemm/nvfp4_fp8_sweep.py

@@ -113,7 +113,6 @@ def nvfp4_fp8_scale_sweep(
     x: torch.Tensor,
     global_amax: torch.Tensor,
     block_size: int = 16,
-    candidates: torch.Tensor | None = None,
 ) -> torch.Tensor:
     """Find the per-block FP8 scale that minimizes NVFP4 quantization MSE.
 
@@ -126,8 +125,6 @@ def nvfp4_fp8_scale_sweep(
             ``block_size``; layout is treated as a flat ``[N_BLOCKS, BLOCK_SIZE]``.
         global_amax: Scalar FP32 global amax (``= reduce_amax(per_block_amax)``).
         block_size: NVFP4 block size (typically 16).
-        candidates: Optional precomputed candidate tensor of shape ``[126]`` (must
-            be the FP8 E4M3 valid values divided by 448). Built lazily if omitted.
 
     Returns:
         ``best_amax`` of shape ``[N_BLOCKS]``, fp32, on the same device as ``x``.
@@ -139,13 +136,7 @@ def nvfp4_fp8_scale_sweep(
     if x.numel() % block_size != 0:
         raise ValueError(f"x.numel() ({x.numel()}) is not divisible by block_size ({block_size}).")
 
-    if candidates is None:
-        candidates = fp8_scale_candidates(x.device)
-    candidates = candidates.contiguous().to(device=x.device, dtype=torch.float32)
-    if candidates.ndim != 1 or candidates.numel() == 0:
-        raise ValueError(
-            f"candidates must be a non-empty 1-D tensor; got shape {tuple(candidates.shape)}."
-        )
+    candidates = fp8_scale_candidates(x.device).to(dtype=torch.float32)
 
     n_blocks = x.numel() // block_size
     x_flat = x.contiguous().view(-1)

modelopt/torch/quantization/calib/mse.py

@@ -259,10 +259,16 @@ def collect(self, x: torch.Tensor):
 
         x = x.detach()
         # The weight quantizer reshapes its input to [n_blocks, block_size] before
-        # calling collect (see TensorQuantizer._process_for_blockquant).
-        assert x.ndim == 2, f"Expected x to be [n_blocks, block_size]; got shape {tuple(x.shape)}."
+        # calling collect (see TensorQuantizer._process_for_blockquant). Validate
+        # via ValueError so the contract still holds under ``python -O``.
+        if x.ndim != 2:
+            raise ValueError(
+                f"Expected x to be [n_blocks, block_size]; got shape {tuple(x.shape)}."
+            )
         block_size = x.shape[-1]
-        n_blocks = x.numel() // block_size
+        if block_size <= 0:
+            raise ValueError(f"x.shape[-1] must be positive; got {block_size}.")
+        n_blocks = x.shape[0]
         if n_blocks != self._n_blocks:
             raise ValueError(
                 f"initial amax.numel() ({self._n_blocks}) does not match the number "

tests/gpu/torch/quantization/test_nvfp4_fp8_sweep_kernel.py

@@ -153,7 +153,7 @@ def test_reset_allows_recollect():
 @requires_triton
 def test_input_validation():
     """``nvfp4_fp8_scale_sweep`` should reject malformed inputs cleanly."""
-    from modelopt.torch.kernels.quantization.gemm import fp8_scale_candidates, nvfp4_fp8_scale_sweep
+    from modelopt.torch.kernels.quantization.gemm import nvfp4_fp8_scale_sweep
 
     device = "cuda"
     x = torch.randn(64, BLOCK_SIZE, device=device)
@@ -173,11 +173,66 @@ def test_input_validation():
     with pytest.raises(ValueError, match="not divisible"):
         nvfp4_fp8_scale_sweep(x, g, block_size=15)
 
-    # Empty / wrong-rank candidates.
-    with pytest.raises(ValueError, match="non-empty 1-D"):
-        nvfp4_fp8_scale_sweep(x, g, candidates=torch.empty(0, device=device))
-    with pytest.raises(ValueError, match="non-empty 1-D"):
-        nvfp4_fp8_scale_sweep(x, g, candidates=fp8_scale_candidates(device).reshape(2, -1))
+
+@requires_triton
+def test_mse_calibrate_dispatch(monkeypatch):
+    """``mse_calibrate(fp8_scale_sweep=True)`` must install the right calibrator class.
+
+    Default path: ``TritonNVFP4MSECalibrator``.
+    With ``MODELOPT_NVFP4_TRITON_SWEEP=0``: ``NVFP4MSECalibrator`` (and not its subclass).
+    """
+    from _test_utils.torch.quantization.models import SimpleLinear
+
+    import modelopt.torch.quantization as mtq
+    from modelopt.torch.quantization.extensions import get_cuda_ext_mx
+    from modelopt.torch.quantization.nn import TensorQuantizer
+
+    if get_cuda_ext_mx() is None:
+        pytest.skip("cuda_ext_mx is not available")
+
+    cfg = {
+        "quant_cfg": [
+            {
+                "quantizer_name": "*weight_quantizer",
+                "cfg": {
+                    "num_bits": (2, 1),
+                    "block_sizes": {-1: 16, "type": "static", "scale_bits": (4, 3)},
+                    "axis": None,
+                },
+                "enable": True,
+            },
+            {"quantizer_name": "*input_quantizer", "enable": False},
+        ],
+        "algorithm": {"method": "mse", "fp8_scale_sweep": True},
+    }
+
+    def _quantize_and_get_weight_calibrators(model):
+        calib_data = [model.get_input().cuda() for _ in range(2)]
+
+        def forward_loop(m):
+            for batch in calib_data:
+                m(batch)
+
+        mtq.quantize(model, cfg, forward_loop=forward_loop)
+        return [
+            type(m._calibrator)
+            for name, m in model.named_modules()
+            if isinstance(m, TensorQuantizer)
+            and name.endswith("weight_quantizer")
+            and getattr(m, "_calibrator", None) is not None
+        ]
+
+    # Default: triton path.
+    monkeypatch.delenv("MODELOPT_NVFP4_TRITON_SWEEP", raising=False)
+    types_default = _quantize_and_get_weight_calibrators(SimpleLinear().cuda())
+    assert types_default, "expected at least one weight quantizer with a calibrator"
+    assert all(t is TritonNVFP4MSECalibrator for t in types_default), types_default
+
+    # Opt-out: reference path, exact class match (TritonNVFP4MSECalibrator is a subclass).
+    monkeypatch.setenv("MODELOPT_NVFP4_TRITON_SWEEP", "0")
+    types_optout = _quantize_and_get_weight_calibrators(SimpleLinear().cuda())
+    assert types_optout, "expected at least one weight quantizer with a calibrator"
+    assert all(t is NVFP4MSECalibrator for t in types_optout), types_optout
 
 
 def _bench(fn, warmup=2, iters=5):