[Quantization] Fold Triton FP8 sweep into NVFP4MSECalibrator

Per realAsma's review, collapse TritonNVFP4MSECalibrator into NVFP4MSECalibrator
as an internal fast path rather than a separately-exported subclass:

- mse.py: NVFP4MSECalibrator.collect() picks the fused Triton kernel via a
  predicate _can_use_triton_fast_path(x) that requires error_func is None,
  CUDA input, blocked layout matching the per-block amax, the kernel package
  importable, and MODELOPT_NVFP4_TRITON_SWEEP \!= "0". Otherwise falls back
  to the parent's reference 126-step sweep. Override reset() to clear only
  per-cycle state and keep _initial_amax (shape [num_blocks], small) so the
  calibrator is reusable; the multi-collect-after-fast-path case raises a
  RuntimeError with a clear message. TritonNVFP4MSECalibrator class deleted.
- model_calib.py: always instantiate NVFP4MSECalibrator; drop the
  TritonNVFP4MSECalibrator import and the env-var dispatch (now internal).
- tests: drop the TritonNVFP4MSECalibrator references. Force the requested
  path via a _force_sweep_path() context manager around the env var. New
  dispatch tests assert the predicate's behavior for the env opt-out, custom
  error_func, and CPU input cases. test_mse_calibrate_end_to_end exercises
  the full mtq.quantize wiring with default and MODELOPT_NVFP4_TRITON_SWEEP=0
  and asserts bitwise-identical model outputs.

This fixes a latent correctness issue: the previous TritonNVFP4MSECalibrator
silently ignored a custom error_func, so a caller passing a Hessian-weighted
loss (e.g. local-Hessian calibration) would have gotten plain squared-error
results from the kernel. The new predicate routes any non-None error_func to
the reference path so the user's metric is honored.

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

Author: cjluo-nv

modelopt/torch/quantization/calib/mse.py

@@ -16,6 +16,7 @@
 """Calibrator that returns the MSE amax of all collected tensors."""
 
 import math
+import os
 from collections.abc import Callable
 
 import torch
@@ -24,7 +25,7 @@
 from .. import utils as quant_utils
 from .calibrator import _Calibrator
 
-__all__ = ["MseCalibrator", "NVFP4MSECalibrator", "TritonNVFP4MSECalibrator"]
+__all__ = ["MseCalibrator", "NVFP4MSECalibrator"]
 
 
 class MseCalibrator(_Calibrator):
@@ -172,7 +173,15 @@ def compute_amax(self, verbose: bool = False):
 
 
 class NVFP4MSECalibrator(MseCalibrator):
-    """Per-block FP8 scale sweep calibrator for NVFP4 static quantization."""
+    """Per-block FP8 scale sweep calibrator for NVFP4 static quantization.
+
+    Uses a fused Triton kernel as an internal fast path on the first ``collect`` call
+    when (a) ``error_func is None``, (b) the input tensor is on CUDA in the standard
+    blocked ``[n_blocks, block_size]`` layout, and (c) Triton + the kernel package are
+    importable. Falls back to the reference 126-step Python sweep otherwise (custom
+    ``error_func`` users, multi-``collect`` activation flows, CPU inputs, or when the
+    fast path is disabled via ``MODELOPT_NVFP4_TRITON_SWEEP=0``).
+    """
 
     def __init__(
         self,
@@ -185,6 +194,8 @@ def __init__(
         """Initialize NVFP4 MSE calibrator with per-block and global amax."""
         super().__init__(amax=amax, axis=axis, quant_func=quant_func, error_func=error_func)
         self._global_amax = global_amax
+        # Set by the Triton fast path on its (one-shot) collect; consumed by compute_amax.
+        self._best_amax_fast: torch.Tensor | None = None
 
     def _compute_candidate_amax(self, candidates: torch.Tensor) -> torch.Tensor:
         if candidates.ndim != 0:  # Called during final compute amax
@@ -204,94 +215,70 @@ def _generate_candidates(self, device: torch.device) -> torch.Tensor:
         fp8_values = fp8_values[valid_mask]
         return fp8_values / 448.0
 
+    def _can_use_triton_fast_path(self, x: torch.Tensor) -> bool:
+        """Whether the Triton fast path is usable for this ``collect`` input.
 
-class TritonNVFP4MSECalibrator(NVFP4MSECalibrator):
-    """Triton-fused FP8 scale sweep calibrator for NVFP4 weight MSE.
-
-    Numerically equivalent to :class:`NVFP4MSECalibrator` but evaluates all 126
-    candidates in a single fused Triton kernel — one weight read instead of 126.
-
-    Limitation: a single ``collect()`` call is supported per ``compute_amax`` cycle.
-    This matches the static weight-MSE flow (``mse_calibrate``'s weight loop), where
-    the calibrator is collected once per weight and immediately consumed. For
-    activation calibration (multiple ``collect`` calls), use :class:`NVFP4MSECalibrator`.
-    Call :meth:`reset` to free internal state and re-enable :meth:`collect`.
-    """
-
-    def __init__(
-        self,
-        amax: torch.Tensor,
-        global_amax: torch.Tensor,
-        axis: int | tuple | list | None = None,
-        quant_func: Callable | None = None,
-        error_func: Callable | None = None,
-    ):
-        """Initialize the Triton-fused NVFP4 MSE calibrator.
-
-        See :class:`NVFP4MSECalibrator`. ``quant_func``/``error_func`` are unused by
-        the kernel path but accepted for API parity. Tile shape and ``num_warps`` are
-        autotuned by the kernel per ``N_BLOCKS``.
+        The kernel produces the final per-block amax in one shot, so it's only usable
+        when the caller wants the standard squared-error sweep on a single CUDA tensor
+        whose layout already matches the per-block amax.
         """
-        super().__init__(
-            amax=amax,
-            global_amax=global_amax,
-            axis=axis,
-            quant_func=quant_func,
-            error_func=error_func,
-        )
-        # Stash shape metadata so collect() can keep working after reset() releases
-        # the (potentially large) _initial_amax buffer.
-        self._initial_amax_shape = tuple(amax.shape)
-        self._initial_amax_dtype = amax.dtype
-        self._n_blocks = int(amax.numel())
-        self._best_amax: torch.Tensor | None = None
+        if self._error_func is not None:
+            return False
+        if not x.is_cuda:
+            return False
+        if os.environ.get("MODELOPT_NVFP4_TRITON_SWEEP", "1") == "0":
+            return False
+        if self._initial_amax is None:
+            return False
+        if x.ndim != 2 or x.shape[0] != int(self._initial_amax.numel()):
+            return False
+        try:
+            from modelopt.torch.kernels.quantization.gemm import nvfp4_fp8_scale_sweep  # noqa: F401
+        except ImportError:
+            return False
+        return True
 
     @torch.no_grad()
     def collect(self, x: torch.Tensor):
-        """Run the fused FP8 sweep kernel and store the resulting per-block amax."""
-        from modelopt.torch.kernels.quantization.gemm import nvfp4_fp8_scale_sweep
-
-        if self._best_amax is not None:
+        """Collect input statistics. Uses the Triton fast path when eligible."""
+        if self._best_amax_fast is not None:
             raise RuntimeError(
-                "TritonNVFP4MSECalibrator.collect() is one-shot; call reset() to "
-                "discard the previous result before collecting again."
-            )
-
-        x = x.detach()
-        # The weight quantizer reshapes its input to [n_blocks, block_size] before
-        # 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)}."
+                "NVFP4MSECalibrator: the Triton fast path produced a final amax on a "
+                "previous collect() call; multi-collect after the fast path is not "
+                "supported. Call reset() to start a fresh cycle, set "
+                "MODELOPT_NVFP4_TRITON_SWEEP=0, or pass a non-None error_func to force "
+                "the reference path for activation-style accumulation."
             )
-        block_size = x.shape[-1]
-        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 "
-                f"of NVFP4 blocks in x ({n_blocks})."
+        # Fast path is eligible only on the first call, before the reference accumulator
+        # has produced any state.
+        if self._losses_sum is None and self._can_use_triton_fast_path(x):
+            from modelopt.torch.kernels.quantization.gemm import nvfp4_fp8_scale_sweep
+
+            best_flat = nvfp4_fp8_scale_sweep(x.detach(), self._global_amax, block_size=x.shape[-1])
+            # Match the original shape/dtype of the initial amax so downstream
+            # load_calib_amax behaves identically to the reference path.
+            self._best_amax_fast = best_flat.reshape(self._initial_amax.shape).to(
+                self._initial_amax.dtype
             )
-
-        best_amax_flat = nvfp4_fp8_scale_sweep(
-            x,
-            self._global_amax,
-            block_size=block_size,
-        )
-        # Match the original shape/dtype of the initial amax so downstream
-        # load_calib_amax behaves identically to the reference path.
-        self._best_amax = best_amax_flat.reshape(self._initial_amax_shape).to(
-            self._initial_amax_dtype
-        )
+            return
+        super().collect(x)
 
     @torch.no_grad()
     def compute_amax(self, verbose: bool = False):
-        """Return the per-block amax computed during ``collect``."""
-        return self._best_amax
+        """Return the per-block amax — from the fast path if it ran, else from the reference sweep."""
+        if self._best_amax_fast is not None:
+            return self._best_amax_fast
+        return super().compute_amax(verbose=verbose)
 
     def reset(self):
-        """Reset the stored best amax. Subsequent ``collect`` calls are allowed."""
-        self._best_amax = None
-        super().reset()
+        """Reset per-cycle state. Keep ``_initial_amax`` so the calibrator stays reusable.
+
+        ``MseCalibrator.reset()`` intentionally drops ``_initial_amax`` to free memory in
+        the multi-step search, but the NVFP4 per-block amax is shape ``[num_blocks]`` —
+        small enough to keep so a follow-up ``collect()`` can run again on the same
+        calibrator instance.
+        """
+        self._best_amax_fast = None
+        self._losses_sum = None
+        self._candidates = None
+        self._amax = None

modelopt/torch/quantization/model_calib.py

@@ -16,7 +16,6 @@
 """Calibration utilities."""
 
 import math
-import os
 import time
 import warnings
 from collections.abc import Callable
@@ -38,7 +37,7 @@
 from modelopt.torch.utils.distributed import DistributedProcessGroup, ParallelState
 from modelopt.torch.utils.network import bind_forward_method, unpatch_forward_method
 
-from .calib import MseCalibrator, NVFP4MSECalibrator, TritonNVFP4MSECalibrator, _Calibrator
+from .calib import MseCalibrator, NVFP4MSECalibrator, _Calibrator
 from .conversion import create_and_replace_svdquant_linear_on_the_fly, set_quantizer_by_cfg_context
 from .nn import NVFP4StaticQuantizer, QuantModule, SequentialQuantizer, TensorQuantizer
 from .utils import (
@@ -355,11 +354,6 @@ def mse_calibrate(
     weight_quantizers = []
     seen_modules = set()
 
-    # Triton-fused FP8 sweep is on by default for NVFP4 static quant; set
-    # MODELOPT_NVFP4_TRITON_SWEEP=0 to fall back to the reference for debugging.
-    use_triton_fp8_sweep = os.environ.get("MODELOPT_NVFP4_TRITON_SWEEP", "1") != "0"
-    nvfp4_calibrator_cls = TritonNVFP4MSECalibrator if use_triton_fp8_sweep else NVFP4MSECalibrator
-
     for name, module in list(model.named_modules()):
         if isinstance(module, TensorQuantizer) and not module._disabled:
             if module._calibrator is not None and not module._dynamic and hasattr(module, "_amax"):
@@ -397,7 +391,10 @@ def mse_calibrate(
                         continue
 
                 if fp8_scale_sweep and is_nvfp4_static:
-                    module._calibrator = nvfp4_calibrator_cls(
+                    # NVFP4MSECalibrator internally selects a fused Triton kernel for
+                    # the standard squared-error sweep; set MODELOPT_NVFP4_TRITON_SWEEP=0
+                    # to force the reference Python sweep for debugging.
+                    module._calibrator = NVFP4MSECalibrator(
                         amax=initial_amax,
                         axis=module._calibrator._axis,
                         global_amax=module.global_amax,