From cb32c7895c9a50cff09eb5ca5be85080cb72543b Mon Sep 17 00:00:00 2001
From: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
Date: Tue, 9 Jun 2026 00:27:29 +0000
Subject: [PATCH 1/4] Add Hessian-weighted NVFP4 FP8 scale-sweep Triton kernel
 for local_hessian

Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
---
 .../quantization/gemm/nvfp4_fp8_sweep.py      | 179 ++++++++++++++-
 modelopt/torch/quantization/calib/mse.py      |  37 +++-
 modelopt/torch/quantization/model_calib.py    |  53 ++++-
 .../test_nvfp4_fp8_sweep_kernel.py            | 209 +++++++++++++++++-
 4 files changed, 446 insertions(+), 32 deletions(-)

diff --git a/modelopt/torch/kernels/quantization/gemm/nvfp4_fp8_sweep.py b/modelopt/torch/kernels/quantization/gemm/nvfp4_fp8_sweep.py
index e15eab328ab..2f6c99ade02 100644
--- a/modelopt/torch/kernels/quantization/gemm/nvfp4_fp8_sweep.py
+++ b/modelopt/torch/kernels/quantization/gemm/nvfp4_fp8_sweep.py
@@ -33,9 +33,14 @@
 import triton.language as tl
 
 from ._fp8_scale_candidates import fp8_scale_candidates
+from .fp4_kernel import compute_fp4_scales
 from .nvfp4_quant import fp4_round_magnitude
 
-__all__ = ["fp8_scale_candidates", "nvfp4_fp8_scale_sweep"]
+__all__ = [
+    "fp8_scale_candidates",
+    "nvfp4_fp8_scale_sweep",
+    "nvfp4_fp8_scale_sweep_hessian",
+]
 
 
 # Selected from a (BLOCKS_PER_PROGRAM, num_warps) sweep on B300:
@@ -102,6 +107,23 @@ def _fp8_scale_sweep_kernel(
     tl.store(best_amax_ptr + block_idx, best_amax, mask=block_mask)
 
 
+def _prepare_block_sweep(x: torch.Tensor, block_size: int):
+    """Validate a block-sweep weight tensor; return ``(n_blocks, x_flat, best_amax)``.
+
+    Shared by both FP8 scale-sweep entry points so their input contracts stay in sync.
+    """
+    if not x.is_cuda:
+        raise ValueError("nvfp4 FP8 scale sweep requires a CUDA tensor.")
+    if not isinstance(block_size, int) or block_size <= 0:
+        raise ValueError(f"block_size must be a positive int, got {block_size!r}.")
+    if x.numel() % block_size != 0:
+        raise ValueError(f"x.numel() ({x.numel()}) is not divisible by block_size ({block_size}).")
+    n_blocks = x.numel() // block_size
+    x_flat = x.contiguous().view(-1)
+    best_amax = torch.empty(n_blocks, dtype=torch.float32, device=x.device)
+    return n_blocks, x_flat, best_amax
+
+
 def nvfp4_fp8_scale_sweep(
     x: torch.Tensor,
     global_amax: torch.Tensor,
@@ -122,19 +144,9 @@ def nvfp4_fp8_scale_sweep(
     Returns:
         ``best_amax`` of shape ``[N_BLOCKS]``, fp32, on the same device as ``x``.
     """
-    if not x.is_cuda:
-        raise ValueError("nvfp4_fp8_scale_sweep requires a CUDA tensor.")
-    if not isinstance(block_size, int) or block_size <= 0:
-        raise ValueError(f"block_size must be a positive int, got {block_size!r}.")
-    if x.numel() % block_size != 0:
-        raise ValueError(f"x.numel() ({x.numel()}) is not divisible by block_size ({block_size}).")
-
+    n_blocks, x_flat, best_amax = _prepare_block_sweep(x, block_size)
     candidates = fp8_scale_candidates(x.device).to(dtype=torch.float32)
-
-    n_blocks = x.numel() // block_size
-    x_flat = x.contiguous().view(-1)
     global_amax_f32 = global_amax.detach().to(device=x.device, dtype=torch.float32).reshape(1)
-    best_amax = torch.empty(n_blocks, dtype=torch.float32, device=x.device)
 
     grid = lambda meta: (triton.cdiv(n_blocks, meta["BLOCKS_PER_PROGRAM"]),)
     with torch.cuda.device(x.device):
@@ -148,3 +160,146 @@ def nvfp4_fp8_scale_sweep(
             NUM_CANDIDATES=int(candidates.numel()),
         )
     return best_amax
+
+
+# One program sweeps ROWS_PER_PROGRAM output rows of a single cin-block, so that block's
+# [BLOCK_SIZE, BLOCK_SIZE] Hessian is loaded once and reused (it is shared across all output
+# rows). The per-candidate quadratic form dwᵀ H dw is a ``[ROWS, BS] x [BS, BS]`` matmul run
+# on tensor cores via ``tl.dot`` — far faster than a per-block broadcast (which also spills
+# the gathered Hessian tile). A single config keeps the (already heavy) 126-way-unrolled
+# kernel to one compilation per weight shape; ROWS=32/num_warps=4 was fastest across shapes
+# on the target GPU (B300), and ROWS=32 covers any COUT via row masking.
+_FP8_SWEEP_HESSIAN_AUTOTUNE_CONFIGS = [
+    triton.Config({"ROWS_PER_PROGRAM": 32}, num_warps=4),
+]
+
+
+@triton.autotune(configs=_FP8_SWEEP_HESSIAN_AUTOTUNE_CONFIGS, key=["COUT", "N_CIN_BLOCKS"])
+@triton.jit
+def _fp8_scale_sweep_hessian_kernel(
+    x_ptr,  # [COUT * N_CIN_BLOCKS * BLOCK_SIZE], any float dtype (loaded as fp32)
+    hessian_ptr,  # [N_CIN_BLOCKS * BLOCK_SIZE * BLOCK_SIZE] fp32
+    candidate_scales_ptr,  # [NUM_CANDIDATES] fp32: per-candidate FP8-quantized block scale
+    candidate_amaxes_ptr,  # [NUM_CANDIDATES] fp32: per-candidate block amax (kernel output value)
+    best_amax_ptr,  # [COUT * N_CIN_BLOCKS] fp32 output
+    COUT,
+    N_CIN_BLOCKS,
+    BLOCK_SIZE: tl.constexpr,
+    NUM_CANDIDATES: tl.constexpr,
+    ROWS_PER_PROGRAM: tl.constexpr,
+):
+    pid = tl.program_id(axis=0)
+    # One cin-block per program; a tile of output rows shares its Hessian.
+    cin_block = pid % N_CIN_BLOCKS
+    rows = (pid // N_CIN_BLOCKS) * ROWS_PER_PROGRAM + tl.arange(0, ROWS_PER_PROGRAM)
+    row_mask = rows < COUT
+    # NVFP4 block layout is row-major over (cout, cin // block_size): block = row*N_CIN + n.
+    block_idx = rows * N_CIN_BLOCKS + cin_block
+
+    # Signed weights are required: the Hessian quadratic form has cross-terms, so unlike
+    # the plain squared-error sweep the residual sign does not cancel.
+    elem = tl.arange(0, BLOCK_SIZE)
+    w = tl.load(
+        x_ptr + block_idx[:, None] * BLOCK_SIZE + elem[None, :], mask=row_mask[:, None], other=0.0
+    ).to(tl.float32)  # [ROWS, BS]
+    w_abs = tl.abs(w)
+    w_sign = tl.where(w >= 0, 1.0, -1.0)
+
+    # This cin-block's Hessian, loaded once and reused across rows and candidates.
+    p = tl.arange(0, BLOCK_SIZE)
+    q = tl.arange(0, BLOCK_SIZE)
+    hessian = tl.load(
+        hessian_ptr + cin_block * (BLOCK_SIZE * BLOCK_SIZE) + p[:, None] * BLOCK_SIZE + q[None, :]
+    ).to(tl.float32)  # [BS, BS]
+
+    best_loss = tl.full([ROWS_PER_PROGRAM], float("inf"), dtype=tl.float32)
+    best_idx = tl.zeros([ROWS_PER_PROGRAM], dtype=tl.int32)
+
+    # Real hardware loop (not unrolled): with a tl.dot in the body, unrolling all 126
+    # candidates explodes compile time for no runtime gain (the matmul dominates).
+    # Scales are precomputed on the host with the SAME FP8-E4M3 round-trip the reference
+    # fake-quant uses, so the kernel's per-block residual is bit-identical to the reference.
+    for k in tl.range(NUM_CANDIDATES):
+        scale = tl.load(candidate_scales_ptr + k).to(tl.float32)
+        # scale == 0 only when global_amax == 0 (all weights zero -> dw == 0 -> loss 0).
+        scale_safe = tl.where(scale == 0.0, 1.0, scale)
+        q_mag = fp4_round_magnitude(w_abs / scale_safe)
+        dw = w_sign * (w_abs - q_mag * scale_safe)  # = w - quant(w), [ROWS, BS]
+        # Per-row quadratic form dwᵀ H dw (H symmetric): (dw @ H) then row-wise dot with dw.
+        # allow_tf32=False keeps the matmul in true fp32 to match the reference sweep.
+        hdw = tl.dot(dw, hessian, allow_tf32=False)  # [ROWS, BS]
+        loss = tl.sum(hdw * dw, axis=1)  # [ROWS]
+        is_better = loss < best_loss
+        best_loss = tl.where(is_better, loss, best_loss)
+        best_idx = tl.where(is_better, k, best_idx)
+
+    best_amax = tl.load(candidate_amaxes_ptr + best_idx, mask=row_mask, other=0.0).to(tl.float32)
+    tl.store(best_amax_ptr + block_idx, best_amax, mask=row_mask)
+
+
+def nvfp4_fp8_scale_sweep_hessian(
+    x: torch.Tensor,
+    global_amax: torch.Tensor,
+    hessian: torch.Tensor,
+    block_size: int = 16,
+) -> torch.Tensor:
+    """Find the per-block FP8 scale minimizing the Hessian-weighted NVFP4 quant error.
+
+    Hessian-weighted counterpart of :func:`nvfp4_fp8_scale_sweep`: for each NVFP4 block
+    it minimizes ``dwᵀ H dw`` (``dw = w - quant(w)``) over the 126 FP8 E4M3 candidates,
+    where ``H`` is the per-cin-block local Hessian shared across all output rows. Used by
+    :class:`NVFP4MSECalibrator` for ``local_hessian`` calibration.
+
+    Args:
+        x: Weight tensor on CUDA in the blocked ``[N_BLOCKS, block_size]`` layout, row-major
+            over ``(cout, cin // block_size)`` so flat block ``b`` has cin-block
+            ``b % (cin // block_size)``.
+        global_amax: Scalar FP32 global amax (``= reduce_amax(per_block_amax)``).
+        hessian: Per-cin-block Hessian of shape ``[cin // block_size, block_size, block_size]``,
+            fp32 (typically normalized by sample count).
+        block_size: NVFP4 block size (typically 16).
+
+    Returns:
+        ``best_amax`` of shape ``[N_BLOCKS]``, fp32, on the same device as ``x``.
+    """
+    n_blocks, x_flat, best_amax = _prepare_block_sweep(x, block_size)
+    if hessian.dim() != 3 or hessian.shape[1] != block_size or hessian.shape[2] != block_size:
+        raise ValueError(
+            f"hessian must have shape [n_cin_blocks, {block_size}, {block_size}], "
+            f"got {tuple(hessian.shape)}."
+        )
+    n_cin_blocks = hessian.shape[0]
+    if n_blocks % n_cin_blocks != 0:
+        raise ValueError(
+            f"n_blocks ({n_blocks}) is not divisible by n_cin_blocks ({n_cin_blocks})."
+        )
+
+    cout = n_blocks // n_cin_blocks
+    global_amax_f32 = global_amax.detach().to(device=x.device, dtype=torch.float32).reshape(())
+
+    # Per-candidate block amax (= global_amax * fp8/448) and the matching FP8-E4M3-quantized
+    # block scale. Computing the scale here with the reference's ``compute_fp4_scales`` (the
+    # exact path the fake-quant uses) makes the kernel's quantization bit-identical to the
+    # reference sweep, so the only residual divergence is fp32 reduction order in tl.dot.
+    candidate_amaxes = (fp8_scale_candidates(x.device).to(dtype=torch.float32)) * global_amax_f32
+    candidate_scales = compute_fp4_scales(
+        candidate_amaxes, global_amax_f32, quantize_block_scales=True
+    ).to(dtype=torch.float32)
+
+    hessian_flat = hessian.contiguous().to(device=x.device, dtype=torch.float32).view(-1)
+
+    # One program per (row-tile, cin-block): grid = cdiv(cout, ROWS) * n_cin_blocks.
+    grid = lambda meta: (triton.cdiv(cout, meta["ROWS_PER_PROGRAM"]) * n_cin_blocks,)
+    with torch.cuda.device(x.device):
+        _fp8_scale_sweep_hessian_kernel[grid](
+            x_flat,
+            hessian_flat,
+            candidate_scales,
+            candidate_amaxes,
+            best_amax,
+            cout,
+            n_cin_blocks,
+            BLOCK_SIZE=block_size,
+            NUM_CANDIDATES=int(candidate_amaxes.numel()),
+        )
+    return best_amax
diff --git a/modelopt/torch/quantization/calib/mse.py b/modelopt/torch/quantization/calib/mse.py
index e19b83c81a8..e28b73522d7 100644
--- a/modelopt/torch/quantization/calib/mse.py
+++ b/modelopt/torch/quantization/calib/mse.py
@@ -189,10 +189,17 @@ def __init__(
         axis: int | tuple | list | None = None,
         quant_func: Callable | None = None,
         error_func: Callable | None = None,
+        hessian: torch.Tensor | None = None,
     ):
-        """Initialize NVFP4 MSE calibrator with per-block and global amax."""
+        """Initialize NVFP4 MSE calibrator with per-block and global amax.
+
+        ``hessian`` (per-cin-block ``[cin // block_size, block_size, block_size]``) enables
+        the Hessian-weighted Triton fast path (local_hessian); ``error_func`` carries the
+        same metric for the reference fallback when the fast path is unavailable.
+        """
         super().__init__(amax=amax, axis=axis, quant_func=quant_func, error_func=error_func)
         self._global_amax = global_amax.to(dtype=torch.float32)
+        self._hessian = hessian
         # Set by collect() after either sweep path; consumed by compute_amax.
         self._best_amax: torch.Tensor | None = None
 
@@ -211,15 +218,13 @@ def _generate_candidates(self, device: torch.device) -> torch.Tensor:
 
         return fp8_scale_candidates(device)
 
-    def _can_use_triton_fast_path(self, x: torch.Tensor) -> bool:
-        """Whether the Triton fast path is usable for this ``collect`` input.
+    def _triton_sweep_eligible(self, x: torch.Tensor) -> bool:
+        """Shared prerequisites for either Triton sweep kernel on this ``collect`` input.
 
-        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.
+        The kernels produce the final per-block amax in one shot, so they're only usable
+        on a single CUDA tensor whose blocked ``[n_blocks, block_size]`` layout already
+        matches the per-block amax.
         """
-        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":
@@ -234,6 +239,14 @@ def _can_use_triton_fast_path(self, x: torch.Tensor) -> bool:
             return False
         return True
 
+    def _can_use_triton_fast_path(self, x: torch.Tensor) -> bool:
+        """Whether the plain squared-error Triton fast path is usable (no custom metric)."""
+        return self._error_func is None and self._triton_sweep_eligible(x)
+
+    def _can_use_hessian_fast_path(self, x: torch.Tensor) -> bool:
+        """Whether the Hessian-weighted Triton fast path is usable (local_hessian)."""
+        return self._hessian is not None and self._triton_sweep_eligible(x)
+
     @torch.no_grad()
     def collect(self, x: torch.Tensor):
         """Collect input statistics and cache the final per-block amax."""
@@ -242,6 +255,14 @@ def collect(self, x: torch.Tensor):
                 "NVFP4MSECalibrator: a previous collect() call produced a final amax; "
                 "multi-collect is not supported. Call reset() to start a fresh cycle."
             )
+        if self._can_use_hessian_fast_path(x):
+            from modelopt.torch.kernels.quantization.gemm import nvfp4_fp8_scale_sweep_hessian
+
+            best_flat = nvfp4_fp8_scale_sweep_hessian(
+                x.detach(), self._global_amax, self._hessian, block_size=x.shape[-1]
+            )
+            self._best_amax = best_flat.reshape(self._initial_amax.shape).to(dtype=torch.float32)
+            return
         if self._can_use_triton_fast_path(x):
             from modelopt.torch.kernels.quantization.gemm import nvfp4_fp8_scale_sweep
 
diff --git a/modelopt/torch/quantization/model_calib.py b/modelopt/torch/quantization/model_calib.py
index b64d8f6a600..54838c66c12 100644
--- a/modelopt/torch/quantization/model_calib.py
+++ b/modelopt/torch/quantization/model_calib.py
@@ -463,11 +463,13 @@ def _make_weight_mse_calibrator(
     stop_multiplier: float,
     fp8_scale_sweep: bool,
     error_func: Callable[[torch.Tensor, torch.Tensor], torch.Tensor] | None = None,
+    hessian: torch.Tensor | None = None,
 ) -> _Calibrator | None:
     """Create the MSE calibrator for one eligible weight quantizer (``None`` if ineligible).
 
-    ``error_func`` overrides the squared-error metric (local-Hessian's per-block weighting);
-    when set, NVFP4's Triton fast path is bypassed for the reference sweep.
+    ``error_func`` overrides the squared-error metric (local-Hessian's per-block weighting).
+    ``hessian`` (the same per-cin-block metric as a raw tensor) enables NVFP4's Hessian-weighted
+    Triton fast path; ``error_func`` then serves only as the reference fallback.
     """
     if (
         not isinstance(weight_quantizer, TensorQuantizer)
@@ -503,6 +505,7 @@ def _make_weight_mse_calibrator(
                 global_amax=weight_quantizer.global_amax,
                 quant_func=quant_func,
                 error_func=error_func,
+                hessian=hessian,
             )
         # fp8_scale_sweep applies only to registered backends and static NVFP4; skip others.
         return None
@@ -575,11 +578,14 @@ def _mse_calibrate_weights(
     stop_multiplier: float,
     fp8_scale_sweep: bool,
     error_func_for: Callable[[TensorQuantizer], Callable | None] | None = None,
+    hessian_for: Callable[[TensorQuantizer], torch.Tensor | None] | None = None,
 ):
     """Run MSE weight calibration over all eligible quantizers (shared by mse / local-Hessian).
 
     ``error_func_for`` maps a weight quantizer to an optional per-weight error function
-    (local-Hessian's Hessian metric); ``None`` means plain squared error.
+    (local-Hessian's Hessian metric); ``None`` means plain squared error. ``hessian_for``
+    maps a weight quantizer to the same metric as a raw per-cin-block Hessian tensor,
+    enabling the Hessian-weighted Triton fast path.
     """
     seen_modules: set[int] = set()
     pbar = tqdm(desc="MSE weight calibration")
@@ -590,6 +596,7 @@ def _mse_calibrate_weights(
         with enable_weight_access_and_writeback(parent_module, model, name_to_module):
             for weight, weight_quantizer in parent_module.iter_weights_for_calibration():
                 error_func = error_func_for(weight_quantizer) if error_func_for else None
+                hessian = hessian_for(weight_quantizer) if hessian_for else None
                 cal = _make_weight_mse_calibrator(
                     weight_quantizer,
                     step_size,
@@ -597,6 +604,7 @@ def _mse_calibrate_weights(
                     stop_multiplier,
                     fp8_scale_sweep,
                     error_func=error_func,
+                    hessian=hessian,
                 )
                 if cal is None:
                     continue
@@ -626,6 +634,7 @@ def __init__(self, cout: int, cin: int, block_size: int):
         # Not block-divisible -> no Hessian (falls back to plain MSE).
         self.is_enabled = cin % block_size == 0
         self.hessian_per_block: torch.Tensor | None = None
+        self._normalized_hessian: torch.Tensor | None = None
         self.num_samples = 0
 
     @torch.no_grad()
@@ -643,6 +652,20 @@ def accumulate(self, input_tensor: torch.Tensor) -> None:
             self.hessian_per_block += hessian_batch
         self.num_samples += input_tensor.numel() // self.cin
 
+    def normalized_hessian(self) -> torch.Tensor | None:
+        """Per-cin-block Hessian ``H / num_samples`` (``None`` if no samples).
+
+        Shared by both the Triton fast path and the reference ``error_func`` so the two
+        consume one tensor; cached because the accumulated buffer may be freed afterwards.
+        """
+        if (
+            self._normalized_hessian is None
+            and self.hessian_per_block is not None
+            and self.num_samples
+        ):
+            self._normalized_hessian = self.hessian_per_block / self.num_samples
+        return self._normalized_hessian
+
     def build_error_func(
         self, keep_buffer: bool = False
     ) -> Callable[[torch.Tensor, torch.Tensor], torch.Tensor] | None:
@@ -650,11 +673,11 @@ def build_error_func(
 
         Frees the raw Hessian buffer unless ``keep_buffer`` (kept for debug inspection).
         """
-        if self.hessian_per_block is None or self.num_samples == 0:
+        hessian = self.normalized_hessian()
+        if hessian is None:
             return None
         cout = self.cout
         bs = self.block_size
-        hessian = self.hessian_per_block / self.num_samples
         if not keep_buffer:
             self.hessian_per_block = None
 
@@ -855,10 +878,12 @@ def capture(weight_quantizer, weight, input_tensor):
             "diverge under tensor/data parallelism. Treat local_hessian as single-rank for now."
         )
 
-    # Phase 3: build error funcs and run the shared MSE weight loop.
+    # Phase 3: build the Hessian metric (raw tensor for the Triton fast path, error_func for
+    # the reference fallback) and run the shared MSE weight loop. The two share one tensor.
     error_funcs = {
         qid: acc.build_error_func(keep_buffer=debug) for qid, acc in accumulators.items()
     }
+    hessians = {qid: acc.normalized_hessian() for qid, acc in accumulators.items()}
     print_rank_0("local_hessian: Running MSE calibration with local Hessian loss...")
     _mse_calibrate_weights(
         model,
@@ -868,19 +893,25 @@ def capture(weight_quantizer, weight, input_tensor):
         stop_multiplier=stop_multiplier,
         fp8_scale_sweep=fp8_scale_sweep,
         error_func_for=lambda q: error_funcs.get(id(q)),
+        hessian_for=lambda q: hessians.get(id(q)),
     )
 
-    # Free the per-block Hessians (pinned by error_func closures) and the sweep's cached
-    # allocations so export starts from a defragmented allocator.
+    # Free the per-block Hessians (pinned by error_func closures, calibrator references, and
+    # the accumulators' own normalized-Hessian cache) and the sweep's cached allocations so
+    # export starts from a defragmented allocator. Keep the accumulators only for debug.
     error_funcs.clear()
+    hessians.clear()
     for module in name_to_module.values():
         if isinstance(module, TensorQuantizer) and isinstance(module._calibrator, MseCalibrator):
             module._calibrator._error_func = None
-    if torch.cuda.is_available():
-        torch.cuda.empty_cache()
-
+            if isinstance(module._calibrator, NVFP4MSECalibrator):
+                module._calibrator._hessian = None
     if debug:
         model._local_hessian_accumulators = accumulators
+    else:
+        accumulators.clear()
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
 
     print_rank_0("local_hessian: Calibration complete.")
 
diff --git a/tests/gpu/torch/quantization/test_nvfp4_fp8_sweep_kernel.py b/tests/gpu/torch/quantization/test_nvfp4_fp8_sweep_kernel.py
index d1eba4987d3..05eb53b324c 100644
--- a/tests/gpu/torch/quantization/test_nvfp4_fp8_sweep_kernel.py
+++ b/tests/gpu/torch/quantization/test_nvfp4_fp8_sweep_kernel.py
@@ -33,9 +33,13 @@
 
 import modelopt.torch.opt as mto
 import modelopt.torch.quantization as mtq
-from modelopt.torch.kernels.quantization.gemm import nvfp4_fp8_scale_sweep
+from modelopt.torch.kernels.quantization.gemm import (
+    nvfp4_fp8_scale_sweep,
+    nvfp4_fp8_scale_sweep_hessian,
+)
 from modelopt.torch.quantization.calib import NVFP4MSECalibrator
 from modelopt.torch.quantization.extensions import get_cuda_ext_mx
+from modelopt.torch.quantization.model_calib import _LocalHessianAccumulator
 from modelopt.torch.quantization.nn import TensorQuantizer
 from modelopt.torch.quantization.tensor_quant import static_blockwise_fp4_fake_quant
 
@@ -371,6 +375,209 @@ def forward_loop(m):
     assert torch.equal(y_default, y_optout)
 
 
+# --------------------------------------------------------------------------------------
+# Hessian-weighted sweep (local_hessian fast path)
+# --------------------------------------------------------------------------------------
+
+
+def _build_hessian_accumulator(cout, cin, hessian_input, block_size=BLOCK_SIZE):
+    """Real ``_LocalHessianAccumulator`` so the test exercises the production metric."""
+    acc = _LocalHessianAccumulator(cout, cin, block_size)
+    acc.accumulate(hessian_input)
+    return acc
+
+
+def _run_hessian_reference(x_blocks, per_block_amax, global_amax, acc):
+    """Reference 126-step sweep using the Hessian-weighted ``error_func`` (Triton off)."""
+    with _force_sweep_path(triton_enabled=False):
+        cal = NVFP4MSECalibrator(
+            amax=per_block_amax,
+            axis=0,
+            global_amax=global_amax,
+            quant_func=_reference_quant_func(global_amax),
+            error_func=acc.build_error_func(keep_buffer=True),
+        )
+        cal.collect(x_blocks)
+        return cal.compute_amax()
+
+
+def _run_hessian_triton(x_blocks, per_block_amax, global_amax, acc):
+    """Hessian-weighted Triton fast path (same metric as a raw per-cin-block tensor)."""
+    with _force_sweep_path(triton_enabled=True):
+        cal = NVFP4MSECalibrator(
+            amax=per_block_amax,
+            axis=0,
+            global_amax=global_amax,
+            quant_func=_reference_quant_func(global_amax),
+            hessian=acc.normalized_hessian(),
+        )
+        cal.collect(x_blocks)
+        return cal.compute_amax()
+
+
+def _total_hessian_loss(x_blocks, per_block_amax, global_amax, hessian):
+    """Total Hessian-weighted quantization error ``Σ dwᵀ H dw`` under the production
+    (CUDA ``static_blockwise_fp4_fake_quant``) rounding used at deployment — the objective
+    the sweep minimizes, summed over all blocks."""
+    n_blocks = x_blocks.shape[0]
+    n_cin = hessian.shape[0]
+    h_per_block = hessian[torch.arange(n_blocks, device=x_blocks.device) % n_cin]
+    xq = static_blockwise_fp4_fake_quant(x_blocks.float(), per_block_amax, global_amax)
+    dw = x_blocks.float() - xq
+    return (torch.einsum("nij,nj->ni", h_per_block, dw) * dw).sum()
+
+
+@requires_triton
+@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
+@pytest.mark.parametrize(
+    ("cout", "cin"),
+    [(8, 64), (1, 256), (256, 2048)],  # multiple rows share a cin-block Hessian; MoE-expert-sized
+)
+@pytest.mark.parametrize("seed", [0, 1])
+def test_hessian_parity_random_weights(seed, cout, cin, dtype):
+    """The Hessian Triton sweep must select the same per-block scales as the reference
+    Hessian-weighted 126-step sweep, exercising shapes where many output rows share one
+    per-cin-block Hessian (the ``b % n_cin_blocks`` mapping).
+
+    The kernel quantizes candidates with the SAME FP8-E4M3-quantized block scale the
+    reference uses (precomputed via ``compute_fp4_scales``), so the per-block residual is
+    bit-identical. fp32/fp16 are therefore bit-exact here; for bf16 only the occasional
+    block whose two best candidates are within fp32 noise may flip (``tl.dot`` vs ``einsum``
+    accumulation order), so we cap the mismatch fraction tightly and require the total
+    objective (Hessian loss) to be essentially unchanged.
+    """
+    torch.manual_seed(seed)
+    device = "cuda"
+    weight = torch.randn(cout, cin, device=device, dtype=dtype)
+    # Well-conditioned PSD Hessian from many tokens so the argmin is unambiguous.
+    hessian_input = torch.randn(512, cin, device=device, dtype=torch.float32)
+    acc = _build_hessian_accumulator(cout, cin, hessian_input)
+
+    x_blocks = weight.reshape(-1, BLOCK_SIZE)
+    per_block_amax = x_blocks.float().abs().amax(dim=-1)
+    global_amax = per_block_amax.max()
+    hessian = acc.normalized_hessian()
+
+    ref = _run_hessian_reference(x_blocks, per_block_amax, global_amax, acc)
+    tri = _run_hessian_triton(x_blocks, per_block_amax, global_amax, acc)
+
+    assert ref.shape == tri.shape
+    n_blocks = ref.numel()
+    n_diff = int((ref != tri).sum())
+
+    if dtype != torch.bfloat16:
+        # Matching scales + matching FP4 rounding => bit-exact for fp32/fp16.
+        assert torch.equal(ref, tri), (
+            f"{dtype} must be bit-exact: {n_diff}/{n_blocks} blocks differ"
+        )
+        return
+
+    # bf16: only rare fp32 reduction-order ties may flip, and the total objective the kernel
+    # achieves must match the reference's to within fp32 noise.
+    assert n_diff / n_blocks < 1e-3, f"{n_diff}/{n_blocks} blocks differ (>0.1%)"
+    loss_ref = _total_hessian_loss(x_blocks, ref, global_amax, hessian)
+    loss_tri = _total_hessian_loss(x_blocks, tri, global_amax, hessian)
+    rel_gap = ((loss_tri - loss_ref) / loss_ref.abs().clamp_min(1e-12)).abs().item()
+    assert rel_gap < 1e-3, (
+        f"aggregate Hessian-loss gap {rel_gap:.3e} too large "
+        f"({n_diff}/{n_blocks} boundary blocks flipped, dtype={dtype})"
+    )
+
+
+@requires_triton
+def test_hessian_sweep_input_validation():
+    """``nvfp4_fp8_scale_sweep_hessian`` should reject malformed inputs cleanly."""
+    device = "cuda"
+    cout, cin = 4, 64
+    x = torch.randn(cout, cin, device=device).reshape(-1, BLOCK_SIZE)
+    g = x.abs().amax()
+    h = torch.randn(cin // BLOCK_SIZE, BLOCK_SIZE, BLOCK_SIZE, device=device)
+
+    with pytest.raises(ValueError, match="CUDA"):
+        nvfp4_fp8_scale_sweep_hessian(x.cpu(), g.cpu(), h.cpu())
+    with pytest.raises(ValueError, match="block_size"):
+        nvfp4_fp8_scale_sweep_hessian(x, g, h, block_size=0)
+    # Wrong Hessian block dims.
+    with pytest.raises(ValueError, match="hessian must have shape"):
+        nvfp4_fp8_scale_sweep_hessian(x, g, torch.randn(4, 8, 8, device=device))
+
+
+@requires_triton
+def test_hessian_fast_path_dispatch():
+    """A non-None ``hessian`` takes the Hessian fast path even though ``error_func`` is set,
+    and ``MODELOPT_NVFP4_TRITON_SWEEP=0`` forces the reference path back on."""
+    torch.manual_seed(0)
+    cout, cin = 8, 64
+    weight = torch.randn(cout, cin, device="cuda", dtype=torch.float32)
+    acc = _build_hessian_accumulator(cout, cin, torch.randn(256, cin, device="cuda"))
+    x_blocks = weight.reshape(-1, BLOCK_SIZE)
+    per_block_amax = x_blocks.abs().amax(dim=-1)
+    global_amax = per_block_amax.max()
+
+    cal = NVFP4MSECalibrator(
+        amax=per_block_amax,
+        axis=0,
+        global_amax=global_amax,
+        quant_func=_reference_quant_func(global_amax),
+        error_func=acc.build_error_func(keep_buffer=True),
+        hessian=acc.normalized_hessian(),
+    )
+    with _force_sweep_path(triton_enabled=True):
+        assert cal._can_use_hessian_fast_path(x_blocks) is True
+    with _force_sweep_path(triton_enabled=False):
+        assert cal._can_use_hessian_fast_path(x_blocks) is False
+
+
+@requires_triton
+def test_hessian_speedup_report(capsys):
+    """Hessian Triton fast path must be >=30x faster than the reference Hessian sweep on a
+    representative 8192x4096 weight (~2M NVFP4 blocks), with parity on the chosen amax."""
+    torch.manual_seed(123)
+    device = "cuda"
+    cout, cin = 8192, 4096
+    weight = torch.randn(cout, cin, device=device, dtype=torch.float32)
+    hessian_input = torch.randn(512, cin, device=device, dtype=torch.float32)
+    acc = _build_hessian_accumulator(cout, cin, hessian_input)
+
+    x_blocks = weight.reshape(-1, BLOCK_SIZE)
+    per_block_amax = x_blocks.abs().amax(dim=-1)
+    global_amax = per_block_amax.max()
+
+    ref_amax = _run_hessian_reference(x_blocks, per_block_amax, global_amax, acc)
+    tri_amax = _run_hessian_triton(x_blocks, per_block_amax, global_amax, acc)
+    n_blocks = ref_amax.numel()
+    n_diff = int((ref_amax != tri_amax).sum())
+    # fp32 weights are bit-exact; any divergence (only on FP4-boundary blocks for lower
+    # precision) must leave the total objective essentially unchanged.
+    hessian = acc.normalized_hessian()
+    loss_ref = _total_hessian_loss(x_blocks, ref_amax, global_amax, hessian)
+    loss_tri = _total_hessian_loss(x_blocks, tri_amax, global_amax, hessian)
+    rel_gap = ((loss_tri - loss_ref) / loss_ref.abs().clamp_min(1e-12)).abs().item()
+    assert rel_gap < 1e-3, (
+        f"{n_diff}/{n_blocks} blocks disagree, aggregate Hessian-loss gap {rel_gap:.3e}"
+    )
+
+    # Reference Hessian sweep is seconds-slow (126 einsums over a 2M-block weight), so use
+    # fewer iters; the Triton path is sub-ms and gets the default count.
+    ref_t = _bench(
+        lambda: _run_hessian_reference(x_blocks, per_block_amax, global_amax, acc),
+        warmup=1,
+        iters=2,
+    )
+    tri_t = _bench(lambda: _run_hessian_triton(x_blocks, per_block_amax, global_amax, acc))
+    speedup = ref_t / tri_t
+
+    with capsys.disabled():
+        print(
+            f"\n[NVFP4 Hessian FP8 sweep] weight=({cout},{cin}) "
+            f"n_blocks={n_blocks} block_size={BLOCK_SIZE} mismatched_blocks={n_diff}\n"
+            f"  reference path: {ref_t * 1e3:8.2f} ms\n"
+            f"  triton fast path: {tri_t * 1e3:8.2f} ms\n"
+            f"  speedup: {speedup:.1f}x"
+        )
+    assert speedup >= 30.0, f"Hessian fast path speedup {speedup:.1f}x below 30x target"
+
+
 def _bench(fn, warmup=2, iters=5):
     for _ in range(warmup):
         fn()

From 0a6e8818e026db61d1fa9b755adf5f65a5b94bc6 Mon Sep 17 00:00:00 2001
From: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
Date: Tue, 9 Jun 2026 23:04:58 +0000
Subject: [PATCH 2/4] minor: changelog (0.46), address PR review, curate tests,
 trim comments

Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
---
 CHANGELOG.rst                                 |  1 +
 .../quantization/gemm/nvfp4_fp8_sweep.py      | 58 +++++++------------
 modelopt/torch/quantization/model_calib.py    | 10 ++--
 .../test_nvfp4_fp8_sweep_kernel.py            | 31 +---------
 4 files changed, 29 insertions(+), 71 deletions(-)

diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index 49c58586674..d3d0ec160ec 100755
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -8,6 +8,7 @@ Changelog
 
 - Add the ``day0-release`` agent skill (``.agents/skills/day0-release/``), a deterministic end-to-end driver that chains the PTQ → evaluation → comparison skills (the evaluation stage deploys the checkpoint itself) with an enforced gate after each stage and returns a publish decision (ACCEPT / REGRESSION / ANOMALOUS / INFEASIBLE). Ships three GPU-free, unit-tested gate scripts (``gate_ptq.py``, ``gate_run.py``, ``gate_compare.py``) that validate checkpoint coverage, evaluation-run completeness, and baseline-vs-candidate accuracy threshold. v1 reports and stops on regression; the recipe-search loop is deferred.
 - Add **streaming** speculative-decoding training (EAGLE3 / DFlash): the draft trains on base-model hidden states produced on the fly by a co-located ``vllm serve`` (no disk dump), moved trainer-side over NIXL RDMA, scaling to multi-node (dedicated serve replicas + DDP trainers). New launcher examples for NVFP4 Kimi-K2.5 / K2.6 on GB200/aarch64 under ``tools/launcher/examples/moonshotai/``.
+- Add a fused Triton fast path for ``local_hessian`` NVFP4 weight-scale search (the Hessian-weighted FP8-E4M3 scale sweep). For each NVFP4 block it minimizes ``dwᵀ H dw`` over the 126 candidate scales using the per-cin-block local Hessian on tensor cores, replacing the per-weight Python reference sweep — roughly **34x** faster on a single 8192x4096 weight and bit-exact with the reference for fp32/fp16 weights. Used automatically during ``local_hessian`` calibration for both dense and fused-MoE expert weights; falls back to the reference sweep on CPU, when Triton is unavailable, or via ``MODELOPT_NVFP4_TRITON_SWEEP=0``.
 
 0.45 (2026-06-xx)
 ^^^^^^^^^^^^^^^^^
diff --git a/modelopt/torch/kernels/quantization/gemm/nvfp4_fp8_sweep.py b/modelopt/torch/kernels/quantization/gemm/nvfp4_fp8_sweep.py
index 2f6c99ade02..8af4026f07b 100644
--- a/modelopt/torch/kernels/quantization/gemm/nvfp4_fp8_sweep.py
+++ b/modelopt/torch/kernels/quantization/gemm/nvfp4_fp8_sweep.py
@@ -162,13 +162,9 @@ def nvfp4_fp8_scale_sweep(
     return best_amax
 
 
-# One program sweeps ROWS_PER_PROGRAM output rows of a single cin-block, so that block's
-# [BLOCK_SIZE, BLOCK_SIZE] Hessian is loaded once and reused (it is shared across all output
-# rows). The per-candidate quadratic form dwᵀ H dw is a ``[ROWS, BS] x [BS, BS]`` matmul run
-# on tensor cores via ``tl.dot`` — far faster than a per-block broadcast (which also spills
-# the gathered Hessian tile). A single config keeps the (already heavy) 126-way-unrolled
-# kernel to one compilation per weight shape; ROWS=32/num_warps=4 was fastest across shapes
-# on the target GPU (B300), and ROWS=32 covers any COUT via row masking.
+# Each program sweeps a tile of output rows of one cin-block, so the block's [BS, BS] Hessian
+# loads once and dwᵀ H dw runs as a [ROWS, BS] x [BS, BS] tl.dot on tensor cores. ROWS=32 was
+# fastest in a shape sweep; a single config keeps compile time to one build per weight shape.
 _FP8_SWEEP_HESSIAN_AUTOTUNE_CONFIGS = [
     triton.Config({"ROWS_PER_PROGRAM": 32}, num_warps=4),
 ]
@@ -189,15 +185,13 @@ def _fp8_scale_sweep_hessian_kernel(
     ROWS_PER_PROGRAM: tl.constexpr,
 ):
     pid = tl.program_id(axis=0)
-    # One cin-block per program; a tile of output rows shares its Hessian.
     cin_block = pid % N_CIN_BLOCKS
     rows = (pid // N_CIN_BLOCKS) * ROWS_PER_PROGRAM + tl.arange(0, ROWS_PER_PROGRAM)
     row_mask = rows < COUT
-    # NVFP4 block layout is row-major over (cout, cin // block_size): block = row*N_CIN + n.
+    # Block layout is row-major over (cout, cin // block_size): block = row*N_CIN + n.
     block_idx = rows * N_CIN_BLOCKS + cin_block
 
-    # Signed weights are required: the Hessian quadratic form has cross-terms, so unlike
-    # the plain squared-error sweep the residual sign does not cancel.
+    # Signed residual: the Hessian cross-terms mean the sign does not cancel (unlike plain MSE).
     elem = tl.arange(0, BLOCK_SIZE)
     w = tl.load(
         x_ptr + block_idx[:, None] * BLOCK_SIZE + elem[None, :], mask=row_mask[:, None], other=0.0
@@ -205,28 +199,24 @@ def _fp8_scale_sweep_hessian_kernel(
     w_abs = tl.abs(w)
     w_sign = tl.where(w >= 0, 1.0, -1.0)
 
-    # This cin-block's Hessian, loaded once and reused across rows and candidates.
-    p = tl.arange(0, BLOCK_SIZE)
-    q = tl.arange(0, BLOCK_SIZE)
+    idx = tl.arange(0, BLOCK_SIZE)
     hessian = tl.load(
-        hessian_ptr + cin_block * (BLOCK_SIZE * BLOCK_SIZE) + p[:, None] * BLOCK_SIZE + q[None, :]
+        hessian_ptr
+        + cin_block * (BLOCK_SIZE * BLOCK_SIZE)
+        + idx[:, None] * BLOCK_SIZE
+        + idx[None, :]
     ).to(tl.float32)  # [BS, BS]
 
     best_loss = tl.full([ROWS_PER_PROGRAM], float("inf"), dtype=tl.float32)
     best_idx = tl.zeros([ROWS_PER_PROGRAM], dtype=tl.int32)
 
-    # Real hardware loop (not unrolled): with a tl.dot in the body, unrolling all 126
-    # candidates explodes compile time for no runtime gain (the matmul dominates).
-    # Scales are precomputed on the host with the SAME FP8-E4M3 round-trip the reference
-    # fake-quant uses, so the kernel's per-block residual is bit-identical to the reference.
+    # Non-unrolled loop: unrolling 126 tl.dot bodies explodes compile time for no runtime gain.
     for k in tl.range(NUM_CANDIDATES):
         scale = tl.load(candidate_scales_ptr + k).to(tl.float32)
-        # scale == 0 only when global_amax == 0 (all weights zero -> dw == 0 -> loss 0).
-        scale_safe = tl.where(scale == 0.0, 1.0, scale)
+        scale_safe = tl.where(scale == 0.0, 1.0, scale)  # scale == 0 only if global_amax == 0
         q_mag = fp4_round_magnitude(w_abs / scale_safe)
         dw = w_sign * (w_abs - q_mag * scale_safe)  # = w - quant(w), [ROWS, BS]
-        # Per-row quadratic form dwᵀ H dw (H symmetric): (dw @ H) then row-wise dot with dw.
-        # allow_tf32=False keeps the matmul in true fp32 to match the reference sweep.
+        # dwᵀ H dw per row (H symmetric); allow_tf32=False keeps it true fp32 vs the reference.
         hdw = tl.dot(dw, hessian, allow_tf32=False)  # [ROWS, BS]
         loss = tl.sum(hdw * dw, axis=1)  # [ROWS]
         is_better = loss < best_loss
@@ -275,22 +265,16 @@ def nvfp4_fp8_scale_sweep_hessian(
         )
 
     cout = n_blocks // n_cin_blocks
-    global_amax_f32 = global_amax.detach().to(device=x.device, dtype=torch.float32).reshape(())
-
-    # Per-candidate block amax (= global_amax * fp8/448) and the matching FP8-E4M3-quantized
-    # block scale. Computing the scale here with the reference's ``compute_fp4_scales`` (the
-    # exact path the fake-quant uses) makes the kernel's quantization bit-identical to the
-    # reference sweep, so the only residual divergence is fp32 reduction order in tl.dot.
-    candidate_amaxes = (fp8_scale_candidates(x.device).to(dtype=torch.float32)) * global_amax_f32
-    candidate_scales = compute_fp4_scales(
-        candidate_amaxes, global_amax_f32, quantize_block_scales=True
-    ).to(dtype=torch.float32)
-
-    hessian_flat = hessian.contiguous().to(device=x.device, dtype=torch.float32).view(-1)
-
-    # One program per (row-tile, cin-block): grid = cdiv(cout, ROWS) * n_cin_blocks.
     grid = lambda meta: (triton.cdiv(cout, meta["ROWS_PER_PROGRAM"]) * n_cin_blocks,)
     with torch.cuda.device(x.device):
+        global_amax_f32 = global_amax.detach().to(device=x.device, dtype=torch.float32).reshape(())
+        # Candidate scales via the reference ``compute_fp4_scales`` (the exact fake-quant path)
+        # keep the kernel's residual bit-identical to the reference sweep.
+        candidate_amaxes = fp8_scale_candidates(x.device).to(dtype=torch.float32) * global_amax_f32
+        candidate_scales = compute_fp4_scales(
+            candidate_amaxes, global_amax_f32, quantize_block_scales=True
+        ).to(dtype=torch.float32)
+        hessian_flat = hessian.contiguous().to(device=x.device, dtype=torch.float32).view(-1)
         _fp8_scale_sweep_hessian_kernel[grid](
             x_flat,
             hessian_flat,
diff --git a/modelopt/torch/quantization/model_calib.py b/modelopt/torch/quantization/model_calib.py
index 54838c66c12..53d497d43c9 100644
--- a/modelopt/torch/quantization/model_calib.py
+++ b/modelopt/torch/quantization/model_calib.py
@@ -878,8 +878,9 @@ def capture(weight_quantizer, weight, input_tensor):
             "diverge under tensor/data parallelism. Treat local_hessian as single-rank for now."
         )
 
-    # Phase 3: build the Hessian metric (raw tensor for the Triton fast path, error_func for
-    # the reference fallback) and run the shared MSE weight loop. The two share one tensor.
+    # Phase 3: weight search. Build error_funcs first so build_error_func caches the normalized
+    # Hessian (freeing the raw buffer) before normalized_hessian() reuses it; the fast path
+    # (tensor) and reference fallback (error_func) then share that one tensor.
     error_funcs = {
         qid: acc.build_error_func(keep_buffer=debug) for qid, acc in accumulators.items()
     }
@@ -896,9 +897,8 @@ def capture(weight_quantizer, weight, input_tensor):
         hessian_for=lambda q: hessians.get(id(q)),
     )
 
-    # Free the per-block Hessians (pinned by error_func closures, calibrator references, and
-    # the accumulators' own normalized-Hessian cache) and the sweep's cached allocations so
-    # export starts from a defragmented allocator. Keep the accumulators only for debug.
+    # Release the per-block Hessians (held by the error_func closures, calibrators, and the
+    # accumulators' cache) before empty_cache so export starts defragmented; keep only for debug.
     error_funcs.clear()
     hessians.clear()
     for module in name_to_module.values():
diff --git a/tests/gpu/torch/quantization/test_nvfp4_fp8_sweep_kernel.py b/tests/gpu/torch/quantization/test_nvfp4_fp8_sweep_kernel.py
index 05eb53b324c..cc7f1a3c2cb 100644
--- a/tests/gpu/torch/quantization/test_nvfp4_fp8_sweep_kernel.py
+++ b/tests/gpu/torch/quantization/test_nvfp4_fp8_sweep_kernel.py
@@ -502,36 +502,10 @@ def test_hessian_sweep_input_validation():
         nvfp4_fp8_scale_sweep_hessian(x, g, torch.randn(4, 8, 8, device=device))
 
 
-@requires_triton
-def test_hessian_fast_path_dispatch():
-    """A non-None ``hessian`` takes the Hessian fast path even though ``error_func`` is set,
-    and ``MODELOPT_NVFP4_TRITON_SWEEP=0`` forces the reference path back on."""
-    torch.manual_seed(0)
-    cout, cin = 8, 64
-    weight = torch.randn(cout, cin, device="cuda", dtype=torch.float32)
-    acc = _build_hessian_accumulator(cout, cin, torch.randn(256, cin, device="cuda"))
-    x_blocks = weight.reshape(-1, BLOCK_SIZE)
-    per_block_amax = x_blocks.abs().amax(dim=-1)
-    global_amax = per_block_amax.max()
-
-    cal = NVFP4MSECalibrator(
-        amax=per_block_amax,
-        axis=0,
-        global_amax=global_amax,
-        quant_func=_reference_quant_func(global_amax),
-        error_func=acc.build_error_func(keep_buffer=True),
-        hessian=acc.normalized_hessian(),
-    )
-    with _force_sweep_path(triton_enabled=True):
-        assert cal._can_use_hessian_fast_path(x_blocks) is True
-    with _force_sweep_path(triton_enabled=False):
-        assert cal._can_use_hessian_fast_path(x_blocks) is False
-
-
 @requires_triton
 def test_hessian_speedup_report(capsys):
-    """Hessian Triton fast path must be >=30x faster than the reference Hessian sweep on a
-    representative 8192x4096 weight (~2M NVFP4 blocks), with parity on the chosen amax."""
+    """Report the Hessian fast-path speedup on a representative 8192x4096 weight (~2M NVFP4
+    blocks); expect ~30x on A6000, higher on datacenter GPUs. Mirrors ``test_speedup_report`"""
     torch.manual_seed(123)
     device = "cuda"
     cout, cin = 8192, 4096
@@ -575,7 +549,6 @@ def test_hessian_speedup_report(capsys):
             f"  triton fast path: {tri_t * 1e3:8.2f} ms\n"
             f"  speedup: {speedup:.1f}x"
         )
-    assert speedup >= 30.0, f"Hessian fast path speedup {speedup:.1f}x below 30x target"
 
 
 def _bench(fn, warmup=2, iters=5):

From 0be2f65b82b207271e080897e9ad0c0afd0e3cb2 Mon Sep 17 00:00:00 2001
From: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
Date: Wed, 10 Jun 2026 00:10:49 +0000
Subject: [PATCH 3/4] minor

Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
---
 .../kernels/quantization/gemm/nvfp4_fp8_sweep.py  | 15 ++++++++-------
 1 file changed, 8 insertions(+), 7 deletions(-)

diff --git a/modelopt/torch/kernels/quantization/gemm/nvfp4_fp8_sweep.py b/modelopt/torch/kernels/quantization/gemm/nvfp4_fp8_sweep.py
index 8af4026f07b..641c7df71ee 100644
--- a/modelopt/torch/kernels/quantization/gemm/nvfp4_fp8_sweep.py
+++ b/modelopt/torch/kernels/quantization/gemm/nvfp4_fp8_sweep.py
@@ -163,14 +163,13 @@ def nvfp4_fp8_scale_sweep(
 
 
 # Each program sweeps a tile of output rows of one cin-block, so the block's [BS, BS] Hessian
-# loads once and dwᵀ H dw runs as a [ROWS, BS] x [BS, BS] tl.dot on tensor cores. ROWS=32 was
-# fastest in a shape sweep; a single config keeps compile time to one build per weight shape.
-_FP8_SWEEP_HESSIAN_AUTOTUNE_CONFIGS = [
-    triton.Config({"ROWS_PER_PROGRAM": 32}, num_warps=4),
-]
+# loads once and dwᵀ H dw runs as a [ROWS, BS] x [BS, BS] tl.dot on tensor cores.
+# ROWS_PER_PROGRAM=32 / num_warps=4 was fastest in a shape sweep; hard-coded (not autotuned)
+# since there is no second config to tune over.
+_HESSIAN_ROWS_PER_PROGRAM = 32
+_HESSIAN_NUM_WARPS = 4
 
 
-@triton.autotune(configs=_FP8_SWEEP_HESSIAN_AUTOTUNE_CONFIGS, key=["COUT", "N_CIN_BLOCKS"])
 @triton.jit
 def _fp8_scale_sweep_hessian_kernel(
     x_ptr,  # [COUT * N_CIN_BLOCKS * BLOCK_SIZE], any float dtype (loaded as fp32)
@@ -265,7 +264,7 @@ def nvfp4_fp8_scale_sweep_hessian(
         )
 
     cout = n_blocks // n_cin_blocks
-    grid = lambda meta: (triton.cdiv(cout, meta["ROWS_PER_PROGRAM"]) * n_cin_blocks,)
+    grid = (triton.cdiv(cout, _HESSIAN_ROWS_PER_PROGRAM) * n_cin_blocks,)
     with torch.cuda.device(x.device):
         global_amax_f32 = global_amax.detach().to(device=x.device, dtype=torch.float32).reshape(())
         # Candidate scales via the reference ``compute_fp4_scales`` (the exact fake-quant path)
@@ -285,5 +284,7 @@ def nvfp4_fp8_scale_sweep_hessian(
             n_cin_blocks,
             BLOCK_SIZE=block_size,
             NUM_CANDIDATES=int(candidate_amaxes.numel()),
+            ROWS_PER_PROGRAM=_HESSIAN_ROWS_PER_PROGRAM,
+            num_warps=_HESSIAN_NUM_WARPS,
         )
     return best_amax

From 4534bd589ae4e69349833e9faea14b957aaed7c3 Mon Sep 17 00:00:00 2001
From: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
Date: Mon, 15 Jun 2026 20:51:35 +0000
Subject: [PATCH 4/4] doc update

Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
---
 modelopt/torch/quantization/calib/mse.py | 16 +++++++++++-----
 1 file changed, 11 insertions(+), 5 deletions(-)

diff --git a/modelopt/torch/quantization/calib/mse.py b/modelopt/torch/quantization/calib/mse.py
index e28b73522d7..7d6583aacd2 100644
--- a/modelopt/torch/quantization/calib/mse.py
+++ b/modelopt/torch/quantization/calib/mse.py
@@ -175,11 +175,17 @@ def compute_amax(self, verbose: bool = False):
 class NVFP4MSECalibrator(MseCalibrator):
     """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 and caches
-    the final amax immediately, so this calibrator is one-shot between resets.
+    ``collect`` dispatches to one of two fused Triton fast paths, else the reference 126-step
+    Python sweep. Both fast paths require the input on CUDA in the blocked
+    ``[n_blocks, block_size]`` layout with Triton + the kernel package importable:
+
+    - **Hessian-weighted** (local_hessian): taken when ``hessian is not None`` — minimizes
+      ``dwᵀ H dw``. Wins over the plain path, so it fires even when ``error_func`` is also set.
+    - **plain squared-error**: taken when ``hessian is None and error_func is None``.
+
+    Otherwise (CPU, non-blocked layout, Triton unavailable, or an ``error_func`` with no
+    ``hessian``) it runs the reference sweep, using ``error_func`` as the metric when set.
+    The final amax is cached immediately, so this calibrator is one-shot between resets.
     """
 
     def __init__(