add gptq fused kernel

modelopt/torch/quantization/config.py

@@ -1549,6 +1549,12 @@ class GPTQCalibConfig(QuantizeAlgorithmConfig):
         description="""The block size for GPTQ weight update, which must be a multiple of the
         group_size used in the quantization.""",
     )
+    fused: bool = ModeloptField(
+        default=False,
+        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.""",
+    )
 
 
 QuantizeQuantCfgType = list[QuantizerCfgEntry]

modelopt/torch/quantization/model_calib.py

@@ -1698,6 +1698,7 @@ def gptq(
     forward_loop: ForwardLoop,
     perc_damp: float = 0.01,
     block_size: int = 128,
+    fused: bool = False,
 ):
     """GPTQ quantization.
 
@@ -1723,6 +1724,7 @@ def gptq(
         forward_loop: Callable that replays calibration inputs through *model*.
         perc_damp: Percentage of avg Hessian diagonal for damping (default: 0.01).
         block_size: Block size for GPTQ weight update.
+        fused: If True, use fused Triton kernel for NVFP4 static quantization.
     """
     total_start = time.time()
 
@@ -1745,7 +1747,7 @@ def _make_gptq_handle(name, m):
             cls = GPTQHelper
         else:
             cls = _GPTQ_HELPER_REGISTRY.get(backend, GPTQHelper)
-        return cls(m, name, offload_to_cpu=True)
+        return cls(m, name, offload_to_cpu=True, fused=fused)
 
     gptq_handles = {name: _make_gptq_handle(name, m) for name, m in quantized_layers}
     for handle in gptq_handles.values():

modelopt/torch/quantization/triton/fp4_kernel.py

@@ -24,7 +24,9 @@
 import triton
 import triton.language as tl
 
-__all__ = ["fp4_dequantize", "static_blockwise_fp4_fake_quant"]
+from .nvfp4_quant import nvfp4_scalar_quant
+
+__all__ = ["compute_fp4_scales", "fp4_dequantize", "static_blockwise_fp4_fake_quant"]
 
 
 _TORCH_TO_TL_DTYPE = {
@@ -198,52 +200,47 @@ def static_blockwise_fp4_fake_quant_kernel(
     idx = block_offset + tl.arange(0, BLOCK_SIZE)
 
     scale = tl.load(scale_ptr + pid).to(tl.float32)
-
     x = tl.load(x_ptr + idx).to(tl.float32)
 
-    x_abs = tl.abs(x)
-    # If scale is 0, inf, or nan, use 1.0 (matching CUDA kernel behavior)
-    # Note: (x != x) checks if x is NaN per IEEE 754
-    scale_safe = tl.where(
-        (scale == 0) | (scale != scale) | (tl.abs(scale) == float("inf")),  # noqa: PLR0124
-        1.0,
-        scale,
-    )
-    abs_scaled = x_abs / scale_safe
-
-    # FP4 values: 0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0
-    q_val = tl.where(
-        abs_scaled <= 0.25,
-        0.0,
-        tl.where(
-            abs_scaled < 0.75,
-            0.5,
-            tl.where(
-                abs_scaled <= 1.25,
-                1.0,
-                tl.where(
-                    abs_scaled < 1.75,
-                    1.5,
-                    tl.where(
-                        abs_scaled <= 2.5,
-                        2.0,
-                        tl.where(
-                            abs_scaled < 3.5,
-                            3.0,
-                            tl.where(abs_scaled <= 5.0, 4.0, 6.0),
-                        ),
-                    ),
-                ),
-            ),
-        ),
-    )
-
-    x_rescaled = q_val * scale_safe
-    x_quant = tl.where(x >= 0, x_rescaled, -x_rescaled)
+    x_quant = nvfp4_scalar_quant(x, scale, BLOCK_SIZE)
 
     tl.store(y_ptr + idx, x_quant.to(OUT_DTYPE))
 
 
+def compute_fp4_scales(
+    amax: torch.Tensor,
+    global_amax: torch.Tensor | None = None,
+    quantize_block_scales: bool = True,
+) -> torch.Tensor:
+    """Compute per-block FP4 scales from amax values.
+
+    ``scale = amax / 6.0``, optionally quantized to FP8 E4M3.
+
+    Args:
+        amax: Per-block amax values (any shape).
+        global_amax: Global amax for FP8 two-level scaling. Computed from *amax* if None.
+        quantize_block_scales: If True, quantize scales to FP8 E4M3.
+
+    Returns:
+        Per-block scales (same shape as *amax*), float32.
+    """
+    amax = amax.float()
+    scale = amax / 6.0  # FP4 max representable value is 6.0
+
+    if quantize_block_scales:
+        from modelopt.torch.quantization.tensor_quant import scaled_e4m3_impl
+        from modelopt.torch.quantization.utils import reduce_amax
+
+        if global_amax is None:
+            global_amax = reduce_amax(amax, axis=None, keepdims=False, squeeze_scalar=True)
+
+        global_amax = global_amax.float()
+        scale_fp8_quant_amax = global_amax / 6.0
+        scale = scaled_e4m3_impl(scale, scale_fp8_quant_amax)
+
+    return scale
+
+
 def static_blockwise_fp4_fake_quant(
     x: torch.Tensor,
     amax: torch.Tensor,
@@ -266,19 +263,7 @@ def static_blockwise_fp4_fake_quant(
     if out_dtype is None:
         out_dtype = x.dtype
 
-    amax = amax.float()  # Requires to be in float32
-    scale = amax / 6.0  # FP4 max representable value is 6.0
-
-    if quantize_block_scales:
-        from modelopt.torch.quantization.tensor_quant import scaled_e4m3_impl
-        from modelopt.torch.quantization.utils import reduce_amax
-
-        if global_amax is None:
-            global_amax = reduce_amax(amax, axis=None, keepdims=False, squeeze_scalar=True)
-
-        global_amax = global_amax.float()
-        scale_fp8_quant_amax = global_amax / 6.0
-        scale = scaled_e4m3_impl(scale, scale_fp8_quant_amax)
+    scale = compute_fp4_scales(amax, global_amax, quantize_block_scales)
 
     x_flat = x.contiguous().view(-1)
     y_flat = torch.empty_like(x_flat, dtype=out_dtype)

modelopt/torch/quantization/triton/fp4_kernel_hopper.py

@@ -24,6 +24,7 @@
 import triton.language as tl
 
 from .fp4_kernel import _torch_dtype_to_tl
+from .nvfp4_quant import fp4_round_magnitude, fp8_quantize_scale
 
 __all__ = ["fp4_fake_quant_block"]
 
@@ -79,9 +80,7 @@ def fp4_fake_quant_kernel(
 
     block_max = tl.max(x_abs, axis=2, keep_dims=True)
 
-    block_max_scaled = block_max / (6.0 * global_scale_safe)
-    block_max_scaled = tl.minimum(block_max_scaled, 448.0)
-    block_max_quant = block_max_scaled.to(tl.float8e4nv).to(tl.float32) * global_scale
+    block_max_quant = fp8_quantize_scale(block_max, global_scale_safe)
     block_max_quant = tl.where(block_max_quant >= 1e-5, block_max_quant, 1.0)
 
     block_max_quant_broadcast = tl.broadcast_to(
@@ -90,31 +89,7 @@ def fp4_fake_quant_kernel(
 
     abs_scaled = x_abs / block_max_quant_broadcast
 
-    q_val = tl.where(
-        abs_scaled <= 0.25,
-        0.0,
-        tl.where(
-            abs_scaled < 0.75,
-            0.5,
-            tl.where(
-                abs_scaled <= 1.25,
-                1.0,
-                tl.where(
-                    abs_scaled < 1.75,
-                    1.5,
-                    tl.where(
-                        abs_scaled <= 2.5,
-                        2.0,
-                        tl.where(
-                            abs_scaled < 3.5,
-                            3.0,
-                            tl.where(abs_scaled <= 5.0, 4.0, 6.0),
-                        ),
-                    ),
-                ),
-            ),
-        ),
-    )
+    q_val = fp4_round_magnitude(abs_scaled)
 
     x_rescaled = q_val * block_max_quant_broadcast
     x_rescaled = tl.where(tile_reshaped >= 0, x_rescaled, -x_rescaled)

modelopt/torch/quantization/triton/gptq_fused_kernel.py

@@ -0,0 +1,136 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""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.
+
+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``.
+"""
+
+import torch
+import triton
+import triton.language as tl
+
+from .nvfp4_quant import nvfp4_scalar_qdq
+
+__all__ = ["gptq_fused_block_scalar"]
+
+
+# ---------------------------------------------------------------------------
+# Scalar kernel — NVFP4 QDQ + error propagation
+# ---------------------------------------------------------------------------
+
+
+@triton.jit
+def _gptq_scalar_kernel(
+    w_ptr,
+    qw_ptr,
+    err_ptr,
+    amax_ptr,
+    global_scale,
+    hinv_ptr,
+    num_rows,
+    n_amax_blocks,
+    quant_block_size,
+    block_start,
+    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
+
+    j_range = tl.arange(0, BLOCK_SIZE)
+    w_full = tl.load(w_base + j_range)
+
+    for col in range(0, BLOCK_SIZE, 1):
+        block_amax = tl.load(amax_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,
+            )
+        )
+
+        d_val = tl.load(hinv_ptr + col * BLOCK_SIZE + 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)
+        w_full = w_full - err_val * hinv_row
+
+
+def gptq_fused_block_scalar(
+    w_block: torch.Tensor,
+    block_amax: torch.Tensor,
+    global_scale: float,
+    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
+    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).
+        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.
+
+    Returns:
+        ``(qw_block, err_block)`` each ``[num_rows, block_size]``.
+    """
+    num_rows, block_size = w_block.shape
+
+    qw_block = torch.empty_like(w_block)
+    err_block = torch.empty_like(w_block)
+
+    _gptq_scalar_kernel[(num_rows,)](
+        w_block.contiguous(),
+        qw_block,
+        err_block,
+        block_amax.contiguous(),
+        global_scale,
+        h_inv_cho_blk.contiguous(),
+        num_rows,
+        block_amax.shape[1],
+        quant_block_size,
+        block_start,
+        BLOCK_SIZE=block_size,
+    )
+
+    return qw_block, err_block