add gptq fused kernel

sychen52 · sychen52 · commit fc14872b357c · 2026-04-17T16:03:25.000-07:00
Signed-off-by: Shiyang Chen &lt;shiychen@nvidia.com&gt;
diff --git a/modelopt/torch/quantization/triton/__init__.py b/modelopt/torch/quantization/triton/__init__.py
@@ -33,6 +33,7 @@
         # fp4_kernel works on any CUDA GPU with triton
         from .fp4_kernel import *
         from .fp8_kernel import *
+        from .nvfp4_quant import *
 
         # fp4_kernel_hopper requires compute >= 8.9 (uses tl.float8e4nv)
         if torch.cuda.get_device_capability() >= (8, 9):
diff --git a/modelopt/torch/quantization/triton/fp4_kernel.py b/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)
diff --git a/modelopt/torch/quantization/triton/fp4_kernel_hopper.py b/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
 
 __all__ = ["fp4_fake_quant_block"]
 
@@ -90,31 +91,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)
diff --git a/modelopt/torch/quantization/triton/gptq_fused_kernel.py b/modelopt/torch/quantization/triton/gptq_fused_kernel.py
@@ -0,0 +1,133 @@
+# 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.
+Each kernel fuses quantization + per-column GPTQ error propagation into
+one launch per GPTQ block, avoiding the Python-level per-column loop.
+
+Architecture (both kernels):
+  - One Triton program per output row.
+  - ``w_full [BLOCK_SIZE]`` register tensor holds working weights.
+  - Per-column error propagation: ``w_full -= err * h_inv_row``.
+
+Scalar kernel (``_gptq_scalar_kernel``):
+  - Calls ``nvfp4_scalar_quant()`` from ``nvfp4_quant.py`` per column.
+"""
+
+import torch
+import triton
+import triton.language as tl
+
+from .nvfp4_quant import nvfp4_scalar_quant
+
+__all__ = ["gptq_fused_block_scalar"]
+
+
+# ---------------------------------------------------------------------------
+# Scalar kernel — NVFP4 quantization + error propagation
+# ---------------------------------------------------------------------------
+
+
+@triton.jit
+def _gptq_scalar_kernel(
+    w_ptr,
+    qw_ptr,
+    err_ptr,
+    scales_ptr,
+    hinv_ptr,
+    num_rows,
+    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
+    scales_base = scales_ptr + row * n_scale_blocks
+
+    j_range = tl.arange(0, BLOCK_SIZE)
+    w_full = tl.load(w_base + j_range, mask=j_range < n_cols, other=0.0)
+
+    for col in range(0, BLOCK_SIZE, 1):
+        scale = tl.load(scales_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_quant(
+                tl.full([1], w_scalar, dtype=tl.float32),
+                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) & (j_range < n_cols)
+        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,
+    scales_2d: torch.Tensor,
+    h_inv_cho_blk: torch.Tensor,
+    quant_block_size: int,
+    block_start: int,
+    n_cols: int,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Run scalar GPTQ (NVFP4) column loop for one block in a single Triton kernel launch.
+
+    Args:
+        w_block:         Working weights ``[num_rows, block_size]`` (float32).
+        scales_2d:       Pre-computed 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.
+        n_cols:          Number of active columns in this block.
+
+    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,
+        scales_2d.contiguous(),
+        h_inv_cho_blk.contiguous(),
+        num_rows,
+        scales_2d.shape[1],
+        quant_block_size,
+        block_start,
+        n_cols,
+        BLOCK_SIZE=block_size,
+    )
+
+    return qw_block, err_block
diff --git a/modelopt/torch/quantization/triton/nvfp4_quant.py b/modelopt/torch/quantization/triton/nvfp4_quant.py
@@ -0,0 +1,95 @@
+# 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.
+
+"""Composable Triton JIT functions for NVFP4 (E2M1) fake quantization.
+
+Single source of truth for FP4 decision-boundary rounding.  Used by:
+  - ``fp4_kernel.py``         (standalone blockwise fake quant)
+  - ``fp4_kernel_hopper.py``  (Hopper block-pointer variant)
+  - ``gptq_fused_kernel.py``  (fused GPTQ scalar path)
+
+FP4 (E2M1) representable magnitudes: {0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0}
+"""
+
+import triton
+import triton.language as tl
+from triton.language.extra.cuda import libdevice
+
+
+@triton.jit
+def fp4_round_magnitude(abs_scaled):
+    """Round ``|x| / scale`` to the nearest FP4 (E2M1) magnitude.
+
+    Works with any tensor shape — the caller is responsible for computing
+    ``abs_scaled = |x| / scale`` beforehand.
+
+    Returns:
+        Tensor of same shape as *abs_scaled* with values in
+        {0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0}.
+    """
+    return 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)),
+                    ),
+                ),
+            ),
+        ),
+    )
+
+
+@triton.jit
+def nvfp4_scalar_quant(
+    x,  # [N] float32, already loaded
+    scale,  # float32 scalar: pre-computed block scale (amax / 6.0)
+    N: tl.constexpr,
+):
+    """NVFP4 scalar fake quantization for a group of elements sharing one scale.
+
+    Quantizes each element independently: divide by scale, round to nearest
+    FP4 (E2M1) value via ``fp4_round_magnitude``, multiply by scale.
+
+    Args:
+        x:     [N] float32 tensor of values to quantize (already in registers).
+        scale: float32 scalar block scale.
+        N:     Compile-time number of elements.
+
+    Returns:
+        x_quant: [N] float32, fake-quantized values.
+    """
+    x_abs = tl.abs(x)
+    # Guard against degenerate scale (matching CUDA kernel behavior)
+    scale_safe = tl.where(
+        (scale == 0.0) | libdevice.isnan(scale) | (tl.abs(scale) == float("inf")),
+        1.0,
+        scale,
+    )
+    abs_scaled = x_abs / scale_safe
+    q_val = fp4_round_magnitude(abs_scaled)
+    x_rescaled = q_val * scale_safe
+    x_quant = tl.where(x >= 0, x_rescaled, -x_rescaled)
+    return x_quant
diff --git a/tests/gpu/torch/quantization/conftest.py b/tests/gpu/torch/quantization/conftest.py
diff --git a/tests/gpu/torch/quantization/test_gptq.py b/tests/gpu/torch/quantization/test_gptq.py
