Default GPTQ to the fused NVFP4 update

CHANGELOG.rst

@@ -24,6 +24,7 @@ Changelog
 - Speed up NVFP4 QAT and QAD of Hugging Face fused MoE experts: all experts of a projection are now fake-quantized in one kernel launch instead of one per expert, about 2x faster per MoE layer on Qwen3-30B-A3B with FSDP2.
 - ``hf_ptq.py`` calibrates Qwen3.8 (``qwen4_exp``) without running its QSA indexer's per-query Python loop when every visible key fits the indexer's token budget, with bit-identical results. On Qwen3.8-Flash-Next this saves about two hours of layerwise calibration.
 - Speed up Hugging Face QAT and QAD with ``QATTrainer``, ``QADTrainer`` or ``--manual_gc``: the ``gc.collect()`` run before each training step now skips the model and optimizer state, saving about 1 s per step (30%) on Qwen3-30B-A3B.
+- GPTQ now defaults to ``fused=True``, running the fused Triton update for static NVFP4 weights on SM89+ (4/6 scaling and other formats keep the unfused loop). Set ``fused: false`` to restore the previous behavior.
 
 *Speculative Decoding*
 

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

@@ -15,22 +15,22 @@
 
 """Fused Triton kernels for GPTQ blockwise weight-update.
 
-A kernel for scalar (NVFP4) quantization with inline two-level scale computation.
-Fuses scale computation + quantization + per-column GPTQ error propagation into
-one launch per GPTQ block, avoiding the Python-level per-column loop.
+A kernel for scalar (NVFP4) quantization with precomputed FP8 block scales.
+Fuses quantization + per-column GPTQ error propagation into one launch per GPTQ
+block, avoiding the Python-level per-column loop.
 
 Architecture:
   - One Triton program per output row.
   - ``w_full [BLOCK_SIZE]`` register tensor holds working weights.
-  - Per-column: calls ``nvfp4_scalar_qdq()`` for FP4 QDQ with inline scale
-    computation, then propagates error via ``w_full -= err * h_inv_row``.
+  - Per-column: calls ``nvfp4_scalar_quant()`` for FP4 QDQ with the column's block
+    scale, then propagates error via ``w_full -= err * h_inv_row``.
 """
 
 import torch
 import triton
 import triton.language as tl
 
-from ..common.nvfp4_quant import nvfp4_scalar_qdq
+from ..common.nvfp4_quant import nvfp4_scalar_quant
 
 __all__ = ["gptq_fused_block_scalar"]
 
@@ -45,68 +45,62 @@ def _gptq_scalar_kernel(
     w_ptr,
     qw_ptr,
     err_ptr,
-    amax_ptr,
-    global_scale,
+    scale_ptr,
     hinv_ptr,
     num_rows,
-    n_amax_blocks,
+    n_scale_blocks,
     quant_block_size,
     block_start,
+    n_cols,
     BLOCK_SIZE: tl.constexpr,
 ):
     row = tl.program_id(0)
     if row >= num_rows:
         return
 
-    w_base = w_ptr + row * BLOCK_SIZE
-    qw_base = qw_ptr + row * BLOCK_SIZE
-    err_base = err_ptr + row * BLOCK_SIZE
-    amax_base = amax_ptr + row * n_amax_blocks
+    w_base = w_ptr + row * n_cols
+    qw_base = qw_ptr + row * n_cols
+    err_base = err_ptr + row * n_cols
+    scale_base = scale_ptr + row * n_scale_blocks
 
+    # BLOCK_SIZE is n_cols rounded up to a power of two, as tl.arange requires.
     j_range = tl.arange(0, BLOCK_SIZE)
-    w_full = tl.load(w_base + j_range)
+    in_block = j_range < n_cols
+    w_full = tl.load(w_base + j_range, mask=in_block, other=0.0)
 
-    for col in range(0, BLOCK_SIZE, 1):
-        block_amax = tl.load(amax_base + (block_start + col) // quant_block_size)
+    for col in range(0, n_cols, 1):
+        block_scale = tl.load(scale_base + (block_start + col) // quant_block_size)
 
         w_scalar = tl.sum(tl.where(j_range == col, w_full, 0.0))
         q_scalar = tl.sum(
-            nvfp4_scalar_qdq(
-                tl.full([1], w_scalar, dtype=tl.float32),
-                block_amax,
-                global_scale,
-                1,
-            )
+            nvfp4_scalar_quant(tl.full([1], w_scalar, dtype=tl.float32), block_scale, 1)
         )
 
-        d_val = tl.load(hinv_ptr + col * BLOCK_SIZE + col)
+        d_val = tl.load(hinv_ptr + col * n_cols + col)
         err_val = (w_scalar - q_scalar) / d_val
         tl.store(err_base + col, err_val)
         tl.store(qw_base + col, q_scalar)
 
-        remaining = j_range > col
-        hinv_row = tl.load(hinv_ptr + col * BLOCK_SIZE + j_range, mask=remaining, other=0.0)
+        remaining = (j_range > col) & in_block
+        hinv_row = tl.load(hinv_ptr + col * n_cols + j_range, mask=remaining, other=0.0)
         w_full = w_full - err_val * hinv_row
 
 
 def gptq_fused_block_scalar(
     w_block: torch.Tensor,
-    block_amax: torch.Tensor,
-    global_scale: float,
+    block_scale: torch.Tensor,
     h_inv_cho_blk: torch.Tensor,
     quant_block_size: int,
     block_start: int,
 ) -> tuple[torch.Tensor, torch.Tensor]:
     """Run scalar GPTQ (NVFP4) column loop for one block in a single Triton kernel launch.
 
-    Computes FP8-quantized scales from per-block amax inline via
-    :func:`nvfp4_scalar_qdq`, then performs NVFP4 fake quantization and
+    Performs NVFP4 fake quantization with the given FP8-quantized block scales and
     GPTQ error propagation per column.
 
     Args:
         w_block:         Working weights ``[num_rows, block_size]`` (float32).
-        block_amax:      Per-block amax values ``[num_rows, n_amax_blocks]`` (float32).
-        global_scale:    Pre-computed ``global_amax / (6.0 * 448.0)`` (scalar).
+        block_scale:     FP8-quantized per-block scales ``[num_rows, n_scale_blocks]`` (float32).
         h_inv_cho_blk:   Block of upper-Cholesky inverse Hessian ``[block_size, block_size]``.
         quant_block_size: Number of elements sharing one scale factor.
         block_start:     Column offset of this block in the full weight matrix.
@@ -126,14 +120,14 @@ def gptq_fused_block_scalar(
             w_block.contiguous(),
             qw_block,
             err_block,
-            block_amax.contiguous(),
-            global_scale,
+            block_scale.contiguous(),
             h_inv_cho_blk.contiguous(),
             num_rows,
-            block_amax.shape[1],
+            block_scale.shape[1],
             quant_block_size,
             block_start,
-            BLOCK_SIZE=block_size,
+            block_size,
+            BLOCK_SIZE=triton.next_power_of_2(block_size),
         )
 
     return qw_block, err_block

modelopt/torch/quantization/config.py

@@ -1266,10 +1266,11 @@ class GPTQCalibConfig(_HessianInputDtypeConfig, QuantizeAlgorithmConfig):
         of the group_size used in the quantization.""",
     )
     fused: bool = ModeloptField(
-        default=False,
+        default=True,
         title="Use fused Triton kernel for GPTQ.",
         description="""When True, use a fused Triton kernel that combines quantization and
-        per-column error propagation into one launch per GPTQ block.""",
+        per-column error propagation into one launch per GPTQ block. Applies to static NVFP4
+        weights without 4/6 scaling on SM89+; other cases use the unfused loop.""",
     )
     importance_weighted: bool = ModeloptField(
         default=False,

modelopt/torch/quantization/model_calib.py

@@ -2520,7 +2520,7 @@ def gptq(
     forward_loop: ForwardLoop,
     perc_damp: float = 0.01,
     block_size: int = 128,
-    fused: bool = False,
+    fused: bool = True,
     importance_weighted: bool = False,
     hessian_input_dtype: str = "float32",
 ):

modelopt/torch/quantization/utils/calib_utils.py

@@ -46,6 +46,7 @@
 import torch
 import torch.distributed as dist
 
+import modelopt.torch.kernels.quantization.gemm as triton_kernel
 from modelopt.torch.utils import (
     get_available_gpu_memory,
     get_process_gpus,
@@ -54,6 +55,8 @@
 )
 from modelopt.torch.utils.network import bind_forward_method, unpatch_forward_method
 
+from .numeric_utils import E4M3_MAX, fp8_max_for_normalization
+
 # A GPU takes a new Hessian only if enough memory stays available after it for the calibration
 # forward that follows. Until GPTQ has measured that forward on the GPU, keep this share free.
 _HESSIAN_UNMEASURED_RESERVE = 0.35
@@ -490,34 +493,52 @@ def _prepare_hessian_inverse(self, hessian, perc_damp):
     def _blockwise_update(self, block_size):
         """Column-wise GPTQ update.
 
-        When ``self.fused`` is True and the weight quantizer is an
-        ``NVFP4StaticQuantizer``, uses :func:`gptq_blockwise_update_fused_scalar`
-        (a fused Triton kernel).  Otherwise falls back to
+        When ``self.fused`` is True and :meth:`_fused_supported` holds, uses
+        :func:`gptq_blockwise_update_fused_scalar` (a fused Triton kernel).  Otherwise falls back to
         :func:`gptq_blockwise_update` (unfused column-by-column loop).
         """
         assert self.weight is not None and self.h_inv is not None, (
             "_blockwise_update called before _prepare_hessian_inverse()"
         )
         quantizer = self.weight_quantizer
 
-        if self.fused and getattr(quantizer, "_is_nvfp4_static_quantizer", False):
+        if self.fused and self._fused_supported(quantizer, self.weight):
             block_sizes = quantizer.block_sizes
+            # The kernel looks up each column's scale by absolute index, so GPTQ blocks need not
+            # align with quantizer groups.
             quant_block_size = block_sizes.get(-1) or block_sizes.get(1)
-            if quant_block_size is not None and block_size % quant_block_size != 0:
-                raise ValueError(
-                    f"GPTQ block_size ({block_size}) must be divisible by the quantizer"
-                    f" group_size ({quant_block_size})"
-                )
             out_features, num_cols = self.weight.shape
-            n_blocks = num_cols // quant_block_size
-            block_amax = quantizer.amax.reshape(out_features, n_blocks).float()
-            global_scale = quantizer.global_amax.float().item() / (6.0 * 448.0)
+            # A width the group size does not divide carries a scale for its padded last group.
+            n_blocks = -(-num_cols // quant_block_size)
+            # The quantizer's own FP8 scales, so the fused update rounds onto the same grid.
+            block_scale = triton_kernel.compute_fp4_scales(quantizer.amax, quantizer.global_amax)
             gptq_blockwise_update_fused_scalar(
-                self.weight, block_amax, global_scale, self.h_inv, block_size, quant_block_size
+                self.weight,
+                block_scale.reshape(out_features, n_blocks),
+                self.h_inv,
+                block_size,
+                quant_block_size,
             )
         else:
             gptq_blockwise_update(self.weight, self.h_inv, block_size, quantizer)
 
+    def _fused_supported(self, quantizer, weight):
+        # The fused kernel stands in for the quantizer's forward with a plain static QDQ over
+        # last-dim groups, so anything else that forward does (pre-quant scale, rotation, LSQ,
+        # 2-D blocks) takes the unfused loop. It also needs FP8 Triton casts (SM89+) and the
+        # default FP8 scale range (no 4/6).
+        return (
+            getattr(quantizer, "_is_nvfp4_static_quantizer", False)
+            and {k for k in quantizer.block_sizes if isinstance(k, int)} in ({-1}, {1})
+            and quantizer.pre_quant_scale is None
+            and not quantizer.rotate_is_enabled
+            and not quantizer._lsq
+            and fp8_max_for_normalization(quantizer) == E4M3_MAX
+            and triton_kernel.IS_AVAILABLE
+            and weight.is_cuda
+            and torch.cuda.get_device_capability(weight.device) >= (8, 9)
+        )
+
     def _print_mse_error(self, hessian, weight):
         """Log Hessian-weighted relative MSE between ``self.weight`` and original ``weight``."""
         w_orig = weight.float()
@@ -565,20 +586,19 @@ def gptq_blockwise_update(weight, h_inv, block_size, quantize_fn):
         weight[:, block_end:].addmm_(errs, h_inv[block_start:block_end, block_end:], alpha=-1)
 
 
-def gptq_blockwise_update_fused_scalar(
-    weight, block_amax, global_scale, h_inv, block_size, quant_block_size
-):
+def gptq_blockwise_update_fused_scalar(weight, block_scale, h_inv, block_size, quant_block_size):
     """Fused GPTQ blockwise update for NVFP4 scalar quantization.
 
-    Uses a fused Triton kernel that combines scale computation, quantization,
-    and per-column error propagation into one launch per GPTQ block, avoiding
-    the Python-level per-column loop in :func:`gptq_blockwise_update`.
+    Uses a fused Triton kernel that combines quantization and per-column error propagation
+    into one launch per GPTQ block, avoiding the Python-level per-column loop in
+    :func:`gptq_blockwise_update`.
 
     Args:
         weight: Weight tensor ``[out_features, in_features]``, modified **in-place**
             with fake-quantized values.
-        block_amax: Per-block amax values ``[out_features, n_amax_blocks]``.
-        global_scale: Pre-computed ``global_amax / (6.0 * 448.0)`` (scalar).
+        block_scale: FP8-quantized per-block scales ``[out_features, n_scale_blocks]``, as
+            :func:`compute_fp4_scales <modelopt.torch.kernels.quantization.gemm.compute_fp4_scales>`
+            returns them.
         h_inv: Upper-triangular Cholesky factor of the damped inverse Hessian
             ``[in_features, in_features]``.
         block_size: Number of columns to process per GPTQ block.
@@ -591,8 +611,7 @@ def gptq_blockwise_update_fused_scalar(
         be = min(bs + block_size, num_cols)
         qw, err = gptq_fused_block_scalar(
             weight[:, bs:be].clone().contiguous(),
-            block_amax,
-            global_scale,
+            block_scale,
             h_inv[bs:be, bs:be].contiguous(),
             quant_block_size,
             bs,