Add fused Triton kernel for local-Hessian NVFP4 weight-scale search

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)
 ^^^^^^^^^^^^^^^^^

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,131 @@ def nvfp4_fp8_scale_sweep(
             NUM_CANDIDATES=int(candidates.numel()),
         )
     return best_amax
+
+
+# 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_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.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)
+    cin_block = pid % N_CIN_BLOCKS
+    rows = (pid // N_CIN_BLOCKS) * ROWS_PER_PROGRAM + tl.arange(0, ROWS_PER_PROGRAM)
+    row_mask = rows < COUT
+    # Block layout is row-major over (cout, cin // block_size): block = row*N_CIN + n.
+    block_idx = rows * N_CIN_BLOCKS + cin_block
+
+    # 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
+    ).to(tl.float32)  # [ROWS, BS]
+    w_abs = tl.abs(w)
+    w_sign = tl.where(w >= 0, 1.0, -1.0)
+
+    idx = tl.arange(0, BLOCK_SIZE)
+    hessian = tl.load(
+        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)
+
+    # 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_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]
+        # 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
+        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
+    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)
+        # 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,
+            candidate_scales,
+            candidate_amaxes,
+            best_amax,
+            cout,
+            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

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__(
@@ -189,10 +195,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 +224,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 +245,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 +261,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