From 41b179e635c5b0bca9b2beb89d87247e9ea678c7 Mon Sep 17 00:00:00 2001 From: Shiyang Chen Date: Fri, 17 Apr 2026 15:56:13 -0700 Subject: [PATCH 1/5] add gptq fused kernel Signed-off-by: Shiyang Chen --- .../torch/quantization/triton/__init__.py | 1 + .../torch/quantization/triton/fp4_kernel.py | 93 ++++---- .../quantization/triton/fp4_kernel_hopper.py | 27 +-- .../quantization/triton/gptq_fused_kernel.py | 133 +++++++++++ .../torch/quantization/triton/nvfp4_quant.py | 95 ++++++++ tests/gpu/torch/quantization/conftest.py | 9 + tests/gpu/torch/quantization/test_gptq.py | 221 ++++++++++++++++++ 7 files changed, 500 insertions(+), 79 deletions(-) create mode 100644 modelopt/torch/quantization/triton/gptq_fused_kernel.py create mode 100644 modelopt/torch/quantization/triton/nvfp4_quant.py diff --git a/modelopt/torch/quantization/triton/__init__.py b/modelopt/torch/quantization/triton/__init__.py index def70e5914f..8934e4d7513 100644 --- 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 index 63a8b3dcb72..9eb6b2d49f5 100644 --- 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 index 2ec31863efc..c46f9c69a91 100644 --- 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 new file mode 100644 index 00000000000..53e9ab88862 --- /dev/null +++ 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 new file mode 100644 index 00000000000..63bf56b4b6c --- /dev/null +++ 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 index 95f2667bee6..9e34e5ef680 100644 --- a/tests/gpu/torch/quantization/conftest.py +++ b/tests/gpu/torch/quantization/conftest.py @@ -13,9 +13,18 @@ # See the License for the specific language governing permissions and # limitations under the License. + import pytest +from modelopt.torch.quantization import triton as triton_kernel + @pytest.fixture(autouse=True) def env_setup(monkeypatch): monkeypatch.setenv("NVIDIA_TF32_OVERRIDE", "0") + + +requires_triton = pytest.mark.skipif( + not triton_kernel.IS_AVAILABLE, + reason="Triton not available", +) diff --git a/tests/gpu/torch/quantization/test_gptq.py b/tests/gpu/torch/quantization/test_gptq.py index 2d5f9d6d707..5399fefe7da 100644 --- a/tests/gpu/torch/quantization/test_gptq.py +++ b/tests/gpu/torch/quantization/test_gptq.py @@ -14,15 +14,18 @@ # limitations under the License. import copy +import time import pytest import torch from _test_utils.torch.transformers_models import get_tiny_llama, get_tiny_tokenizer +from conftest import requires_triton import modelopt.torch.quantization as mtq from modelopt.torch.export.unified_export_hf import _export_quantized_weight from modelopt.torch.quantization.model_calib import gptq from modelopt.torch.quantization.qtensor.nvfp4_tensor import NVFP4QTensor +from modelopt.torch.quantization.triton.fp4_kernel import compute_fp4_scales from modelopt.torch.quantization.utils.calib_utils import update_hessian from modelopt.torch.utils.dataset_utils import create_forward_loop, get_dataset_dataloader @@ -231,3 +234,221 @@ def test_gptq_e2e_flow(quant_cfg): calibrate_loop = create_forward_loop(dataloader=calib_dataloader) model = mtq.quantize(model, quant_cfg, forward_loop=calibrate_loop) + + +# --------------------------------------------------------------------------- +# Fused Triton GPTQ kernel tests for NVFP4 scalar quantization +# --------------------------------------------------------------------------- + + +def _compute_h_inv(hessian, weight, percdamp=0.01): + """Compute damped upper-Cholesky inverse Hessian.""" + h = hessian.clone() + zero_cols = torch.nonzero(weight.eq(0).all(dim=0)).unsqueeze(-1) + h[zero_cols, :] = 0 + h[:, zero_cols] = 0 + h[zero_cols, zero_cols] = 1 + damp = percdamp * torch.mean(torch.diag(h)) + diag_idx = torch.arange(h.shape[0], device=h.device) + h[diag_idx, diag_idx] += damp + h = torch.cholesky_inverse(torch.linalg.cholesky(h)) + return torch.linalg.cholesky(h, upper=True) + + +def _make_nvfp4_test_data(quant_block_size, out_features, dim): + """Create weight, h_inv, and scales_2d for NVFP4 GPTQ tests.""" + weight = torch.randn(out_features, dim, device="cuda", dtype=torch.float32) + n_blocks = dim // quant_block_size + amax = weight.reshape(out_features, n_blocks, quant_block_size).abs().amax(dim=-1) + scales_2d = compute_fp4_scales(amax) + + inp = torch.randn(4, 32, dim, device="cuda") + hessian = torch.zeros(dim, dim, dtype=torch.float32) + hessian, _ = update_hessian(inp, hessian, 0) + hessian = hessian.to("cuda") + h_inv = _compute_h_inv(hessian, weight) + + return weight, scales_2d, h_inv + + +def _run_unfused_gptq_nvfp4(weight, scales_2d, h_inv, gptq_block_size, quant_block_size): + """Unfused NVFP4 GPTQ using the production Triton FP4 kernel per column. + + Both fused and unfused use the same frozen pre-computed scales so the + test verifies the fused kernel's correctness (not scale computation). + """ + from modelopt.torch.quantization.triton.fp4_kernel import static_blockwise_fp4_fake_quant + + out_features, num_cols = weight.shape + n_blocks = num_cols // quant_block_size + w = weight.float().clone() + q = torch.zeros_like(w) + # Recover amax from scales (scales = amax / 6.0, already FP8-quantized) + amax_flat = (scales_2d * 6.0).reshape(out_features * n_blocks) + + for i in range(0, num_cols, gptq_block_size): + j_end = min(i + gptq_block_size, num_cols) + e = torch.zeros(out_features, j_end - i, dtype=w.dtype, device=w.device) + + for j in range(i, j_end): + # Quantize full weight using production Triton FP4 kernel + w_blocked = w.reshape(out_features * n_blocks, quant_block_size) + qdq = static_blockwise_fp4_fake_quant( + w_blocked, + amax_flat, + quantize_block_scales=False, + ).reshape(out_features, num_cols) + q[:, j] = qdq[:, j] + + err = (w[:, j] - q[:, j]) / h_inv[j, j] + e[:, j - i] = err + w[:, j:j_end] -= err.unsqueeze(1) * h_inv[j, j:j_end].unsqueeze(0) + + if j_end < num_cols: + w[:, j_end:] -= e @ h_inv[i:j_end, j_end:] + + return q + + +def _run_fused_gptq_nvfp4(weight, scales_2d, h_inv, gptq_block_size, quant_block_size): + """Fused Triton GPTQ for NVFP4.""" + from modelopt.torch.quantization.triton.gptq_fused_kernel import gptq_fused_block_scalar + + dim = weight.shape[1] + w = weight.float().clone() + for bs in range(0, dim, gptq_block_size): + be = min(bs + gptq_block_size, dim) + nc = be - bs + qw, err = gptq_fused_block_scalar( + w[:, bs:be].clone().contiguous(), + scales_2d, + h_inv[bs:be, bs:be].contiguous(), + quant_block_size, + bs, + nc, + ) + w[:, bs:be] = qw + if be < dim: + w[:, be:].addmm_(err[:, :nc], h_inv[bs:be, be:], alpha=-1) + return w + + +_NVFP4_QUANT_BLOCK_SIZES = [16, 128] +_NVFP4_GPTQ_BLOCK_SIZES = [16, 128] + + +@requires_triton +@pytest.mark.parametrize("quant_block_size", _NVFP4_QUANT_BLOCK_SIZES) +@pytest.mark.parametrize("gptq_block_size", _NVFP4_GPTQ_BLOCK_SIZES) +def test_fused_vs_unfused_nvfp4(quant_block_size, gptq_block_size): + """Fused Triton NVFP4 GPTQ must match unfused production reference.""" + torch.manual_seed(42) + dim = max(256, quant_block_size * 4) + out_features = 64 + + weight, scales_2d, h_inv = _make_nvfp4_test_data( + quant_block_size, + out_features, + dim, + ) + + weight_fused = _run_fused_gptq_nvfp4( + weight, + scales_2d, + h_inv, + gptq_block_size, + quant_block_size, + ) + weight_unfused = _run_unfused_gptq_nvfp4( + weight, + scales_2d, + h_inv, + gptq_block_size, + quant_block_size, + ) + + assert not torch.equal(weight_fused, weight.float()), "Fused did not update weights" + assert not torch.equal(weight_unfused, weight.float()), "Unfused did not update weights" + + diff = (weight_fused - weight_unfused).abs() + max_abs = diff.max().item() + mean_abs = diff.mean().item() + denom = weight_unfused.abs().max().item() + rel_max = max_abs / denom if denom > 0 else 0.0 + + print( + f"\n[nvfp4] gptq_bs={gptq_block_size} quant_bs={quant_block_size}: " + f"max_abs={max_abs:.2e} mean_abs={mean_abs:.2e} rel_max={rel_max:.2e}" + ) + + torch.testing.assert_close(weight_fused, weight_unfused, atol=1e-4, rtol=1e-4) + + +_NVFP4_BENCH_CONFIGS = [ + (16, 128, 256, 512), + (16, 128, 256, 2048), + (16, 128, 256, 4096), + (128, 128, 256, 512), + (128, 128, 256, 2048), + (128, 128, 256, 4096), +] + +_NVFP4_BENCH_IDS = [f"qbs{qbs}_gbs{gbs}_{of}x{d}" for qbs, gbs, of, d in _NVFP4_BENCH_CONFIGS] + + +@requires_triton +@pytest.mark.parametrize( + ("quant_block_size", "gptq_block_size", "out_features", "dim"), + _NVFP4_BENCH_CONFIGS, + ids=_NVFP4_BENCH_IDS, +) +def test_fused_nvfp4_benchmark(quant_block_size, gptq_block_size, out_features, dim): + """Benchmark fused Triton NVFP4 GPTQ vs unfused production loop.""" + torch.manual_seed(42) + + weight, scales_2d, h_inv = _make_nvfp4_test_data( + quant_block_size, + out_features, + dim, + ) + + def _bench(fn, n_warmup=2, n_iters=5): + for _ in range(n_warmup): + fn() + torch.cuda.synchronize() + total = 0.0 + for _ in range(n_iters): + torch.cuda.synchronize() + t0 = time.perf_counter() + fn() + torch.cuda.synchronize() + total += time.perf_counter() - t0 + return total / n_iters + + def run_fused(): + return _run_fused_gptq_nvfp4( + weight, + scales_2d, + h_inv, + gptq_block_size, + quant_block_size, + ) + + def run_unfused(): + return _run_unfused_gptq_nvfp4( + weight, + scales_2d, + h_inv, + gptq_block_size, + quant_block_size, + ) + + t_fused = _bench(run_fused) + t_unfused = _bench(run_unfused) + speedup = t_unfused / t_fused if t_fused > 0 else float("inf") + + tag = f"qbs{quant_block_size}_gbs{gptq_block_size}_{out_features}x{dim}" + print( + f"\n[{tag}] Fused: {t_fused * 1e3:8.2f} ms | " + f"Unfused: {t_unfused * 1e3:8.2f} ms | Speedup: {speedup:.1f}x" + ) From 9452469114d5e7983c56333454eaf038c10a1d9b Mon Sep 17 00:00:00 2001 From: Shiyang Chen Date: Fri, 17 Apr 2026 17:21:25 -0700 Subject: [PATCH 2/5] address reviewer comments Signed-off-by: Shiyang Chen --- .../torch/quantization/triton/__init__.py | 1 - tests/gpu/torch/quantization/test_gptq.py | 69 ++++++++----------- 2 files changed, 29 insertions(+), 41 deletions(-) diff --git a/modelopt/torch/quantization/triton/__init__.py b/modelopt/torch/quantization/triton/__init__.py index 8934e4d7513..def70e5914f 100644 --- a/modelopt/torch/quantization/triton/__init__.py +++ b/modelopt/torch/quantization/triton/__init__.py @@ -33,7 +33,6 @@ # 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/tests/gpu/torch/quantization/test_gptq.py b/tests/gpu/torch/quantization/test_gptq.py index 5399fefe7da..1001ca57567 100644 --- a/tests/gpu/torch/quantization/test_gptq.py +++ b/tests/gpu/torch/quantization/test_gptq.py @@ -25,7 +25,6 @@ from modelopt.torch.export.unified_export_hf import _export_quantized_weight from modelopt.torch.quantization.model_calib import gptq from modelopt.torch.quantization.qtensor.nvfp4_tensor import NVFP4QTensor -from modelopt.torch.quantization.triton.fp4_kernel import compute_fp4_scales from modelopt.torch.quantization.utils.calib_utils import update_hessian from modelopt.torch.utils.dataset_utils import create_forward_loop, get_dataset_dataloader @@ -257,6 +256,8 @@ def _compute_h_inv(hessian, weight, percdamp=0.01): def _make_nvfp4_test_data(quant_block_size, out_features, dim): """Create weight, h_inv, and scales_2d for NVFP4 GPTQ tests.""" + from modelopt.torch.quantization.triton.fp4_kernel import compute_fp4_scales + weight = torch.randn(out_features, dim, device="cuda", dtype=torch.float32) n_blocks = dim // quant_block_size amax = weight.reshape(out_features, n_blocks, quant_block_size).abs().amax(dim=-1) @@ -393,24 +394,12 @@ def test_fused_vs_unfused_nvfp4(quant_block_size, gptq_block_size): (128, 128, 256, 4096), ] -_NVFP4_BENCH_IDS = [f"qbs{qbs}_gbs{gbs}_{of}x{d}" for qbs, gbs, of, d in _NVFP4_BENCH_CONFIGS] - -@requires_triton -@pytest.mark.parametrize( - ("quant_block_size", "gptq_block_size", "out_features", "dim"), - _NVFP4_BENCH_CONFIGS, - ids=_NVFP4_BENCH_IDS, -) -def test_fused_nvfp4_benchmark(quant_block_size, gptq_block_size, out_features, dim): - """Benchmark fused Triton NVFP4 GPTQ vs unfused production loop.""" - torch.manual_seed(42) +def bench_fused_nvfp4(): + """Benchmark fused Triton NVFP4 GPTQ vs unfused production loop (informational-only). - weight, scales_2d, h_inv = _make_nvfp4_test_data( - quant_block_size, - out_features, - dim, - ) + Not collected by pytest. Run directly: ``python tests/gpu/torch/quantization/test_gptq.py`` + """ def _bench(fn, n_warmup=2, n_iters=5): for _ in range(n_warmup): @@ -425,30 +414,30 @@ def _bench(fn, n_warmup=2, n_iters=5): total += time.perf_counter() - t0 return total / n_iters - def run_fused(): - return _run_fused_gptq_nvfp4( - weight, - scales_2d, - h_inv, - gptq_block_size, - quant_block_size, - ) + for quant_block_size, gptq_block_size, out_features, dim in _NVFP4_BENCH_CONFIGS: + torch.manual_seed(42) + weight, scales_2d, h_inv = _make_nvfp4_test_data(quant_block_size, out_features, dim) - def run_unfused(): - return _run_unfused_gptq_nvfp4( - weight, - scales_2d, - h_inv, - gptq_block_size, - quant_block_size, + def run_fused(): + return _run_fused_gptq_nvfp4( + weight, scales_2d, h_inv, gptq_block_size, quant_block_size + ) + + def run_unfused(): + return _run_unfused_gptq_nvfp4( + weight, scales_2d, h_inv, gptq_block_size, quant_block_size + ) + + t_fused = _bench(run_fused) + t_unfused = _bench(run_unfused) + speedup = t_unfused / t_fused if t_fused > 0 else float("inf") + + tag = f"qbs{quant_block_size}_gbs{gptq_block_size}_{out_features}x{dim}" + print( + f"[{tag}] Fused: {t_fused * 1e3:8.2f} ms | " + f"Unfused: {t_unfused * 1e3:8.2f} ms | Speedup: {speedup:.1f}x" ) - t_fused = _bench(run_fused) - t_unfused = _bench(run_unfused) - speedup = t_unfused / t_fused if t_fused > 0 else float("inf") - tag = f"qbs{quant_block_size}_gbs{gptq_block_size}_{out_features}x{dim}" - print( - f"\n[{tag}] Fused: {t_fused * 1e3:8.2f} ms | " - f"Unfused: {t_unfused * 1e3:8.2f} ms | Speedup: {speedup:.1f}x" - ) +if __name__ == "__main__": + bench_fused_nvfp4() From d2c32fcbce3deb3a7758c798b7682ba0784fe5b5 Mon Sep 17 00:00:00 2001 From: Shiyang Chen Date: Sun, 19 Apr 2026 13:52:38 -0700 Subject: [PATCH 3/5] comment 2nd round Signed-off-by: Shiyang Chen --- modelopt/torch/quantization/triton/gptq_fused_kernel.py | 8 ++------ tests/gpu/torch/quantization/test_gptq.py | 6 +++--- 2 files changed, 5 insertions(+), 9 deletions(-) diff --git a/modelopt/torch/quantization/triton/gptq_fused_kernel.py b/modelopt/torch/quantization/triton/gptq_fused_kernel.py index 53e9ab88862..24aedb5a7fc 100644 --- a/modelopt/torch/quantization/triton/gptq_fused_kernel.py +++ b/modelopt/torch/quantization/triton/gptq_fused_kernel.py @@ -53,7 +53,6 @@ def _gptq_scalar_kernel( n_scale_blocks, quant_block_size, block_start, - n_cols, BLOCK_SIZE: tl.constexpr, ): row = tl.program_id(0) @@ -66,7 +65,7 @@ def _gptq_scalar_kernel( 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) + w_full = tl.load(w_base + j_range) for col in range(0, BLOCK_SIZE, 1): scale = tl.load(scales_base + (block_start + col) // quant_block_size) @@ -85,7 +84,7 @@ def _gptq_scalar_kernel( tl.store(err_base + col, err_val) tl.store(qw_base + col, q_scalar) - remaining = (j_range > col) & (j_range < n_cols) + 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 @@ -96,7 +95,6 @@ def gptq_fused_block_scalar( 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. @@ -106,7 +104,6 @@ def gptq_fused_block_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. - n_cols: Number of active columns in this block. Returns: ``(qw_block, err_block)`` each ``[num_rows, block_size]``. @@ -126,7 +123,6 @@ def gptq_fused_block_scalar( scales_2d.shape[1], quant_block_size, block_start, - n_cols, BLOCK_SIZE=block_size, ) diff --git a/tests/gpu/torch/quantization/test_gptq.py b/tests/gpu/torch/quantization/test_gptq.py index 1001ca57567..7f011bb96ed 100644 --- a/tests/gpu/torch/quantization/test_gptq.py +++ b/tests/gpu/torch/quantization/test_gptq.py @@ -240,6 +240,7 @@ def test_gptq_e2e_flow(quant_cfg): # --------------------------------------------------------------------------- +# TODO(shiychen): This should be extracted out from production code path def _compute_h_inv(hessian, weight, percdamp=0.01): """Compute damped upper-Cholesky inverse Hessian.""" h = hessian.clone() @@ -272,6 +273,7 @@ def _make_nvfp4_test_data(quant_block_size, out_features, dim): return weight, scales_2d, h_inv +# TODO(shiychen): This should be extracted out from production code path def _run_unfused_gptq_nvfp4(weight, scales_2d, h_inv, gptq_block_size, quant_block_size): """Unfused NVFP4 GPTQ using the production Triton FP4 kernel per column. @@ -319,18 +321,16 @@ def _run_fused_gptq_nvfp4(weight, scales_2d, h_inv, gptq_block_size, quant_block w = weight.float().clone() for bs in range(0, dim, gptq_block_size): be = min(bs + gptq_block_size, dim) - nc = be - bs qw, err = gptq_fused_block_scalar( w[:, bs:be].clone().contiguous(), scales_2d, h_inv[bs:be, bs:be].contiguous(), quant_block_size, bs, - nc, ) w[:, bs:be] = qw if be < dim: - w[:, be:].addmm_(err[:, :nc], h_inv[bs:be, be:], alpha=-1) + w[:, be:].addmm_(err, h_inv[bs:be, be:], alpha=-1) return w From 08a7392b0c0151dfc0d6bc2b19786eab6e1f1dc9 Mon Sep 17 00:00:00 2001 From: Shiyang Chen Date: Mon, 20 Apr 2026 14:40:32 -0700 Subject: [PATCH 4/5] integration Signed-off-by: Shiyang Chen --- .../torch/quantization/utils/calib_utils.py | 202 +++++++++++++----- tests/gpu/torch/quantization/test_gptq.py | 74 ++----- 2 files changed, 170 insertions(+), 106 deletions(-) diff --git a/modelopt/torch/quantization/utils/calib_utils.py b/modelopt/torch/quantization/utils/calib_utils.py index 252f0af6fc8..5f511ea9f83 100644 --- a/modelopt/torch/quantization/utils/calib_utils.py +++ b/modelopt/torch/quantization/utils/calib_utils.py @@ -74,6 +74,42 @@ def update_hessian(input, hessian, n_samples): return hessian, n_samples +def compute_hessian_inverse(hessian, weight, perc_damp): + """Compute damped upper-Cholesky inverse Hessian. + + Dead-neuron columns (all-zero in ``weight``) are zeroed in the + Hessian before inversion, matching the FP-Quant reference: + https://github.com/IST-DASLab/FP-Quant/blob/d2e3092f968262c4de5fb050e1aef568a280dadd/src/quantization/gptq.py#L200 + + Args: + hessian: Hessian matrix ``[in_features, in_features]``. + weight: Weight matrix ``[out_features, in_features]`` for dead-neuron detection. + perc_damp: Percentage of average Hessian diagonal for damping. + + Returns: + Upper-triangular Cholesky factor of the damped inverse Hessian + ``[in_features, in_features]``. Falls back to the identity matrix + when the Hessian is not positive definite. + """ + h = hessian.clone() + zero_cols = torch.nonzero(weight.eq(0).all(dim=0)).unsqueeze(-1) + + h[zero_cols, :] = 0 + h[:, zero_cols] = 0 + h[zero_cols, zero_cols] = 1 + + damp = perc_damp * torch.mean(torch.diag(h)) + diag_indices = torch.arange(h.shape[0], device=h.device) + h[diag_indices, diag_indices] += damp + + try: + h = torch.cholesky_inverse(torch.linalg.cholesky(h)) + return torch.linalg.cholesky(h, upper=True) + except (RuntimeError, torch.linalg.LinAlgError): + print_rank_0("Warning: Hessian is not positive definite, using identity matrix") + return torch.eye(h.shape[0], device=h.device, dtype=h.dtype) + + class GPTQHelper: """Encapsulates per-module GPTQ state and operations. @@ -154,38 +190,14 @@ def update_weights(self, block_size, perc_damp): # ------------------------------------------------------------------ def _prepare_hessian_inverse(self, hessian, perc_damp): - """Compute damped inverse Hessian and store as ``self.h_inv``. - - Dead-neuron columns (all-zero in ``self.weight``) are zeroed in the - Hessian before inversion, matching the FP-Quant reference: - https://github.com/IST-DASLab/FP-Quant/blob/d2e3092f968262c4de5fb050e1aef568a280dadd/src/quantization/gptq.py#L200 - """ + """Compute damped inverse Hessian and store as ``self.h_inv``.""" assert self.weight is not None, "_prepare_hessian_inverse called before update_weights()" - h = hessian.clone() - zero_cols = torch.nonzero(self.weight.eq(0).all(dim=0)).unsqueeze(-1) - - h[zero_cols, :] = 0 - h[:, zero_cols] = 0 - h[zero_cols, zero_cols] = 1 - - damp = perc_damp * torch.mean(torch.diag(h)) - diag_indices = torch.arange(h.shape[0], device=h.device) - h[diag_indices, diag_indices] += damp - - try: - h = torch.cholesky_inverse(torch.linalg.cholesky(h)) - self.h_inv = torch.linalg.cholesky(h, upper=True) - except (RuntimeError, torch.linalg.LinAlgError): - print_rank_0("Warning: Hessian is not positive definite, using identity matrix") - self.h_inv = torch.eye(h.shape[0], device=h.device, dtype=h.dtype) + self.h_inv = compute_hessian_inverse(hessian, self.weight, perc_damp) def _blockwise_update(self, block_size): """Column-wise GPTQ update using full-matrix QDQ. - For each column, quantizes the full weight matrix via the quantizer and - extracts the quantized column. This is the standard GPTQ approach. - - Reads/writes ``self.weight`` and ``self.h_inv`` in-place. + Delegates to :func:`gptq_blockwise_update` with the module's weight quantizer. """ assert self.weight is not None and self.h_inv is not None, ( "_blockwise_update called before _prepare_hessian_inverse()" @@ -199,28 +211,7 @@ def _blockwise_update(self, block_size): f"GPTQ block_size ({block_size}) must be divisible by the quantizer" f" group_size ({group_size})" ) - num_cols = self.weight.shape[1] - - for block_start in range(0, num_cols, block_size): - block_end = min(block_start + block_size, num_cols) - n_cols_blk = block_end - block_start - h_inv_cho_blk = self.h_inv[block_start:block_end, block_start:block_end] - - wblk = self.weight.clone() - errs = torch.zeros_like(wblk[:, block_start:block_end]) - - for i in range(n_cols_blk): - w_ci = wblk[:, block_start + i] - d = h_inv_cho_blk[i, i] - qdq = quantizer(wblk) - self.weight[:, block_start + i] = qdq[:, block_start + i] - err = (w_ci - qdq[:, block_start + i]) / d - wblk[:, block_start + i : block_end].addr_(err, h_inv_cho_blk[i, i:], alpha=-1) - errs[:, i] = err - - self.weight[:, block_end:].addmm_( - errs, self.h_inv[block_start:block_end, block_end:], alpha=-1 - ) + gptq_blockwise_update(self.weight, self.h_inv, block_size, quantizer) def _print_mse_error(self, hessian): """Log Hessian-weighted relative MSE between ``self.weight`` and original weights.""" @@ -231,6 +222,115 @@ def _print_mse_error(self, hessian): print_rank_0(f"[{self.name}] Relative MSE error: {mse.item():.2e}{suffix}") +def gptq_blockwise_update(weight, h_inv, block_size, quantize_fn): + """Column-wise GPTQ update using full-matrix fake quantization. + + For each column, quantizes the full weight matrix via ``quantize_fn`` and + extracts the quantized column. Error is propagated to remaining columns + within the block and then to all subsequent columns via the inverse Hessian. + + Args: + weight: Weight tensor ``[out_features, in_features]``, modified **in-place** + with fake-quantized values. + 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. + quantize_fn: Callable ``(weight) -> qdq_weight`` that fake-quantizes + the full weight matrix. + """ + num_cols = weight.shape[1] + + for block_start in range(0, num_cols, block_size): + block_end = min(block_start + block_size, num_cols) + n_cols_blk = block_end - block_start + h_inv_cho_blk = h_inv[block_start:block_end, block_start:block_end] + + wblk = weight.clone() + errs = torch.zeros_like(weight[:, block_start:block_end]) + + for i in range(n_cols_blk): + w_ci = wblk[:, block_start + i] + d = h_inv_cho_blk[i, i] + qdq = quantize_fn(wblk) + weight[:, block_start + i] = qdq[:, block_start + i] + err = (w_ci - qdq[:, block_start + i]) / d + wblk[:, block_start + i : block_end].addr_(err, h_inv_cho_blk[i, i:], alpha=-1) + errs[:, i] = err + + weight[:, block_end:].addmm_(errs, h_inv[block_start:block_end, block_end:], alpha=-1) + + +def gptq_blockwise_update_fused_scalar(weight, scales_2d, h_inv, block_size, quant_block_size): + """Fused GPTQ blockwise update for NVFP4 scalar quantization. + + 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. + scales_2d: Pre-computed per-block scales ``[out_features, n_scale_blocks]``. + 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. + quant_block_size: Number of elements sharing one quantization scale factor. + """ + from modelopt.torch.quantization.triton.gptq_fused_kernel import gptq_fused_block_scalar + + num_cols = weight.shape[1] + for bs in range(0, num_cols, block_size): + be = min(bs + block_size, num_cols) + qw, err = gptq_fused_block_scalar( + weight[:, bs:be].clone().contiguous(), + scales_2d, + h_inv[bs:be, bs:be].contiguous(), + quant_block_size, + bs, + ) + weight[:, bs:be] = qw + if be < num_cols: + weight[:, be:].addmm_(err, h_inv[bs:be, be:], alpha=-1) + + +class FusedScalarGPTQHelper(GPTQHelper): + """GPTQHelper using the fused Triton kernel for NVFP4 scalar quantization. + + Overrides :meth:`_blockwise_update` to extract pre-computed scales from the + ``NVFP4StaticQuantizer`` and delegate to :func:`gptq_blockwise_update_fused_scalar`. + """ + + def _blockwise_update(self, block_size): + """Fused GPTQ using Triton kernel for NVFP4 scalar quantization.""" + assert self.weight is not None and self.h_inv is not None, ( + "_blockwise_update called before _prepare_hessian_inverse()" + ) + from modelopt.torch.quantization.triton.fp4_kernel import compute_fp4_scales + + quantizer = self.module.weight_quantizer + block_sizes = getattr(quantizer, "block_sizes", None) + quant_block_size = None + if block_sizes is not None: + 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 + + # Pre-compute scales from the calibrated amax (frozen during GPTQ). + amax = quantizer.amax.reshape(out_features, n_blocks) + scales_2d = compute_fp4_scales(amax, quantizer.global_amax, quantize_block_scales=True) + + gptq_blockwise_update_fused_scalar( + self.weight, scales_2d, self.h_inv, block_size, quant_block_size + ) + + _GPTQ_HELPER_REGISTRY: dict[str, type[GPTQHelper]] = {} @@ -242,3 +342,7 @@ def register_gptq_helper(backend: str, factory: type[GPTQHelper]) -> None: construct ``factory`` instead of the default ``GPTQHelper``. """ _GPTQ_HELPER_REGISTRY[backend] = factory + + +# Built-in registrations +register_gptq_helper("fused_gptq_nvfp4", FusedScalarGPTQHelper) diff --git a/tests/gpu/torch/quantization/test_gptq.py b/tests/gpu/torch/quantization/test_gptq.py index 7f011bb96ed..7f2ad253b26 100644 --- a/tests/gpu/torch/quantization/test_gptq.py +++ b/tests/gpu/torch/quantization/test_gptq.py @@ -25,7 +25,12 @@ from modelopt.torch.export.unified_export_hf import _export_quantized_weight from modelopt.torch.quantization.model_calib import gptq from modelopt.torch.quantization.qtensor.nvfp4_tensor import NVFP4QTensor -from modelopt.torch.quantization.utils.calib_utils import update_hessian +from modelopt.torch.quantization.utils.calib_utils import ( + compute_hessian_inverse, + gptq_blockwise_update, + gptq_blockwise_update_fused_scalar, + update_hessian, +) from modelopt.torch.utils.dataset_utils import create_forward_loop, get_dataset_dataloader RAND_SEED = 42 @@ -240,21 +245,6 @@ def test_gptq_e2e_flow(quant_cfg): # --------------------------------------------------------------------------- -# TODO(shiychen): This should be extracted out from production code path -def _compute_h_inv(hessian, weight, percdamp=0.01): - """Compute damped upper-Cholesky inverse Hessian.""" - h = hessian.clone() - zero_cols = torch.nonzero(weight.eq(0).all(dim=0)).unsqueeze(-1) - h[zero_cols, :] = 0 - h[:, zero_cols] = 0 - h[zero_cols, zero_cols] = 1 - damp = percdamp * torch.mean(torch.diag(h)) - diag_idx = torch.arange(h.shape[0], device=h.device) - h[diag_idx, diag_idx] += damp - h = torch.cholesky_inverse(torch.linalg.cholesky(h)) - return torch.linalg.cholesky(h, upper=True) - - def _make_nvfp4_test_data(quant_block_size, out_features, dim): """Create weight, h_inv, and scales_2d for NVFP4 GPTQ tests.""" from modelopt.torch.quantization.triton.fp4_kernel import compute_fp4_scales @@ -268,14 +258,13 @@ def _make_nvfp4_test_data(quant_block_size, out_features, dim): hessian = torch.zeros(dim, dim, dtype=torch.float32) hessian, _ = update_hessian(inp, hessian, 0) hessian = hessian.to("cuda") - h_inv = _compute_h_inv(hessian, weight) + h_inv = compute_hessian_inverse(hessian, weight, perc_damp=0.01) return weight, scales_2d, h_inv -# TODO(shiychen): This should be extracted out from production code path def _run_unfused_gptq_nvfp4(weight, scales_2d, h_inv, gptq_block_size, quant_block_size): - """Unfused NVFP4 GPTQ using the production Triton FP4 kernel per column. + """Unfused NVFP4 GPTQ using the production blockwise update with Triton FP4 kernel. Both fused and unfused use the same frozen pre-computed scales so the test verifies the fused kernel's correctness (not scale computation). @@ -285,52 +274,23 @@ def _run_unfused_gptq_nvfp4(weight, scales_2d, h_inv, gptq_block_size, quant_blo out_features, num_cols = weight.shape n_blocks = num_cols // quant_block_size w = weight.float().clone() - q = torch.zeros_like(w) # Recover amax from scales (scales = amax / 6.0, already FP8-quantized) amax_flat = (scales_2d * 6.0).reshape(out_features * n_blocks) - for i in range(0, num_cols, gptq_block_size): - j_end = min(i + gptq_block_size, num_cols) - e = torch.zeros(out_features, j_end - i, dtype=w.dtype, device=w.device) - - for j in range(i, j_end): - # Quantize full weight using production Triton FP4 kernel - w_blocked = w.reshape(out_features * n_blocks, quant_block_size) - qdq = static_blockwise_fp4_fake_quant( - w_blocked, - amax_flat, - quantize_block_scales=False, - ).reshape(out_features, num_cols) - q[:, j] = qdq[:, j] - - err = (w[:, j] - q[:, j]) / h_inv[j, j] - e[:, j - i] = err - w[:, j:j_end] -= err.unsqueeze(1) * h_inv[j, j:j_end].unsqueeze(0) + def quantize_fn(w_input): + w_blocked = w_input.reshape(out_features * n_blocks, quant_block_size) + return static_blockwise_fp4_fake_quant( + w_blocked, amax_flat, quantize_block_scales=False + ).reshape(out_features, num_cols) - if j_end < num_cols: - w[:, j_end:] -= e @ h_inv[i:j_end, j_end:] - - return q + gptq_blockwise_update(w, h_inv, gptq_block_size, quantize_fn) + return w def _run_fused_gptq_nvfp4(weight, scales_2d, h_inv, gptq_block_size, quant_block_size): - """Fused Triton GPTQ for NVFP4.""" - from modelopt.torch.quantization.triton.gptq_fused_kernel import gptq_fused_block_scalar - - dim = weight.shape[1] + """Fused Triton GPTQ for NVFP4 using the production fused update.""" w = weight.float().clone() - for bs in range(0, dim, gptq_block_size): - be = min(bs + gptq_block_size, dim) - qw, err = gptq_fused_block_scalar( - w[:, bs:be].clone().contiguous(), - scales_2d, - h_inv[bs:be, bs:be].contiguous(), - quant_block_size, - bs, - ) - w[:, bs:be] = qw - if be < dim: - w[:, be:].addmm_(err, h_inv[bs:be, be:], alpha=-1) + gptq_blockwise_update_fused_scalar(w, scales_2d, h_inv, gptq_block_size, quant_block_size) return w From b331ed2033fbf5daa3b37f6dc6dc3e67546db90a Mon Sep 17 00:00:00 2001 From: Shiyang Chen Date: Mon, 20 Apr 2026 21:45:41 -0700 Subject: [PATCH 5/5] add config for fused kernel; qdq inside kernel; even more extraction Signed-off-by: Shiyang Chen --- modelopt/torch/quantization/config.py | 6 ++ modelopt/torch/quantization/model_calib.py | 4 +- .../quantization/triton/fp4_kernel_hopper.py | 6 +- .../quantization/triton/gptq_fused_kernel.py | 45 ++++++---- .../torch/quantization/triton/nvfp4_quant.py | 49 ++++++++++ .../torch/quantization/utils/calib_utils.py | 89 +++++++------------ tests/gpu/torch/quantization/test_gptq.py | 85 ++++++++++-------- 7 files changed, 164 insertions(+), 120 deletions(-) diff --git a/modelopt/torch/quantization/config.py b/modelopt/torch/quantization/config.py index 3f24ac09a41..186ff1c7edd 100644 --- a/modelopt/torch/quantization/config.py +++ b/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] diff --git a/modelopt/torch/quantization/model_calib.py b/modelopt/torch/quantization/model_calib.py index 9b1cc5bc0c6..04aaa88a519 100644 --- a/modelopt/torch/quantization/model_calib.py +++ b/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(): diff --git a/modelopt/torch/quantization/triton/fp4_kernel_hopper.py b/modelopt/torch/quantization/triton/fp4_kernel_hopper.py index c46f9c69a91..624e723b957 100644 --- a/modelopt/torch/quantization/triton/fp4_kernel_hopper.py +++ b/modelopt/torch/quantization/triton/fp4_kernel_hopper.py @@ -24,7 +24,7 @@ import triton.language as tl from .fp4_kernel import _torch_dtype_to_tl -from .nvfp4_quant import fp4_round_magnitude +from .nvfp4_quant import fp4_round_magnitude, fp8_quantize_scale __all__ = ["fp4_fake_quant_block"] @@ -80,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( diff --git a/modelopt/torch/quantization/triton/gptq_fused_kernel.py b/modelopt/torch/quantization/triton/gptq_fused_kernel.py index 24aedb5a7fc..c070eac8800 100644 --- a/modelopt/torch/quantization/triton/gptq_fused_kernel.py +++ b/modelopt/torch/quantization/triton/gptq_fused_kernel.py @@ -15,30 +15,28 @@ """Fused Triton kernels for GPTQ blockwise weight-update. -A kernel for scalar (NVFP4) quantization. -Each kernel fuses quantization + per-column GPTQ error propagation into +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 (both kernels): +Architecture: - 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. + - 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_quant +from .nvfp4_quant import nvfp4_scalar_qdq __all__ = ["gptq_fused_block_scalar"] # --------------------------------------------------------------------------- -# Scalar kernel — NVFP4 quantization + error propagation +# Scalar kernel — NVFP4 QDQ + error propagation # --------------------------------------------------------------------------- @@ -47,10 +45,11 @@ def _gptq_scalar_kernel( w_ptr, qw_ptr, err_ptr, - scales_ptr, + amax_ptr, + global_scale, hinv_ptr, num_rows, - n_scale_blocks, + n_amax_blocks, quant_block_size, block_start, BLOCK_SIZE: tl.constexpr, @@ -62,19 +61,20 @@ def _gptq_scalar_kernel( 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 + 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): - scale = tl.load(scales_base + (block_start + col) // quant_block_size) + 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_quant( + nvfp4_scalar_qdq( tl.full([1], w_scalar, dtype=tl.float32), - scale, + block_amax, + global_scale, 1, ) ) @@ -91,16 +91,22 @@ def _gptq_scalar_kernel( def gptq_fused_block_scalar( w_block: torch.Tensor, - scales_2d: 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). - scales_2d: Pre-computed scales ``[num_rows, n_scale_blocks]`` (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. @@ -117,10 +123,11 @@ def gptq_fused_block_scalar( w_block.contiguous(), qw_block, err_block, - scales_2d.contiguous(), + block_amax.contiguous(), + global_scale, h_inv_cho_blk.contiguous(), num_rows, - scales_2d.shape[1], + block_amax.shape[1], quant_block_size, block_start, BLOCK_SIZE=block_size, diff --git a/modelopt/torch/quantization/triton/nvfp4_quant.py b/modelopt/torch/quantization/triton/nvfp4_quant.py index 63bf56b4b6c..32ab776b2b4 100644 --- a/modelopt/torch/quantization/triton/nvfp4_quant.py +++ b/modelopt/torch/quantization/triton/nvfp4_quant.py @@ -93,3 +93,52 @@ def nvfp4_scalar_quant( x_rescaled = q_val * scale_safe x_quant = tl.where(x >= 0, x_rescaled, -x_rescaled) return x_quant + + +@triton.jit +def fp8_quantize_scale(block_amax, global_scale): + """FP8 E4M3 fake-quantize the per-block NVFP4 scale. + + Computes ``scale = block_amax / 6.0``, then round-trips it through + FP8 E4M3 using ``global_scale`` for the second-level scaling. + + Works with any tensor shape (scalar, 1-D, or higher) since all ops + are element-wise. + + Args: + block_amax: Per-block amax value(s). + global_scale: Pre-computed ``global_amax / (6.0 * 448.0)``. + + Returns: + FP8-quantized per-block scale(s), same shape as ``block_amax``. + """ + FP8_E4M3_MAX: tl.constexpr = 448.0 + scale_in_fp8_range = block_amax / (6.0 * global_scale) + scale_clamped = tl.minimum(scale_in_fp8_range, FP8_E4M3_MAX) + return scale_clamped.to(tl.float8e4nv).to(tl.float32) * global_scale + + +@triton.jit +def nvfp4_scalar_qdq( + x, # [N] float32, already loaded + block_amax, # float32 scalar: per-block amax + global_scale, # float32 scalar: pre-computed global_amax / (6.0 * 448.0) + N: tl.constexpr, +): + """NVFP4 scalar fake quantization with inline two-level scale computation. + + Computes the per-block FP8-quantized scale from ``block_amax`` via + :func:`fp8_quantize_scale`, then quantizes each element to the nearest + FP4 (E2M1) value. + + Args: + x: [N] float32 tensor of values to quantize. + block_amax: Per-block amax (absolute maximum of the block). + global_scale: Pre-computed ``global_amax / (6.0 * 448.0)``. + N: Compile-time number of elements. + + Returns: + x_quant: [N] float32, fake-quantized values. + """ + scale = fp8_quantize_scale(block_amax, global_scale) + return nvfp4_scalar_quant(x, scale, N) diff --git a/modelopt/torch/quantization/utils/calib_utils.py b/modelopt/torch/quantization/utils/calib_utils.py index 5f511ea9f83..ac2ec7a2553 100644 --- a/modelopt/torch/quantization/utils/calib_utils.py +++ b/modelopt/torch/quantization/utils/calib_utils.py @@ -126,10 +126,11 @@ class GPTQHelper: CACHE_NAME = "_forward_no_gptq_hessian" - def __init__(self, module, name, offload_to_cpu=False): + def __init__(self, module, name, offload_to_cpu=False, fused=False): """Initialize GPTQHelper with module state and Hessian storage.""" self.module = module self.name = name + self.fused = fused in_features = module.weight.shape[-1] device = module.weight.device if device.type == "meta" or (offload_to_cpu and get_used_gpu_mem_fraction(device) > 0.65): @@ -195,23 +196,35 @@ def _prepare_hessian_inverse(self, hessian, perc_damp): self.h_inv = compute_hessian_inverse(hessian, self.weight, perc_damp) def _blockwise_update(self, block_size): - """Column-wise GPTQ update using full-matrix QDQ. + """Column-wise GPTQ update. - Delegates to :func:`gptq_blockwise_update` with the module's weight quantizer. + 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 + :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.module.weight_quantizer - block_sizes = getattr(quantizer, "block_sizes", None) - if block_sizes is not None: - group_size = block_sizes.get(-1) - if group_size is not None and block_size % group_size != 0: + + if self.fused and getattr(quantizer, "_is_nvfp4_static_quantizer", False): + block_sizes = quantizer.block_sizes + 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 ({group_size})" + f" group_size ({quant_block_size})" ) - gptq_blockwise_update(self.weight, self.h_inv, block_size, quantizer) + 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) + gptq_blockwise_update_fused_scalar( + self.weight, block_amax, global_scale, self.h_inv, block_size, quant_block_size + ) + else: + gptq_blockwise_update(self.weight, self.h_inv, block_size, quantizer) def _print_mse_error(self, hessian): """Log Hessian-weighted relative MSE between ``self.weight`` and original weights.""" @@ -260,17 +273,20 @@ 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, scales_2d, h_inv, block_size, quant_block_size): +def gptq_blockwise_update_fused_scalar( + weight, block_amax, global_scale, h_inv, block_size, quant_block_size +): """Fused GPTQ blockwise update for NVFP4 scalar quantization. - 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`. + 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`. Args: weight: Weight tensor ``[out_features, in_features]``, modified **in-place** with fake-quantized values. - scales_2d: Pre-computed per-block scales ``[out_features, n_scale_blocks]``. + block_amax: Per-block amax values ``[out_features, n_amax_blocks]``. + global_scale: Pre-computed ``global_amax / (6.0 * 448.0)`` (scalar). 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. @@ -283,7 +299,8 @@ def gptq_blockwise_update_fused_scalar(weight, scales_2d, h_inv, block_size, qua be = min(bs + block_size, num_cols) qw, err = gptq_fused_block_scalar( weight[:, bs:be].clone().contiguous(), - scales_2d, + block_amax, + global_scale, h_inv[bs:be, bs:be].contiguous(), quant_block_size, bs, @@ -293,44 +310,6 @@ def gptq_blockwise_update_fused_scalar(weight, scales_2d, h_inv, block_size, qua weight[:, be:].addmm_(err, h_inv[bs:be, be:], alpha=-1) -class FusedScalarGPTQHelper(GPTQHelper): - """GPTQHelper using the fused Triton kernel for NVFP4 scalar quantization. - - Overrides :meth:`_blockwise_update` to extract pre-computed scales from the - ``NVFP4StaticQuantizer`` and delegate to :func:`gptq_blockwise_update_fused_scalar`. - """ - - def _blockwise_update(self, block_size): - """Fused GPTQ using Triton kernel for NVFP4 scalar quantization.""" - assert self.weight is not None and self.h_inv is not None, ( - "_blockwise_update called before _prepare_hessian_inverse()" - ) - from modelopt.torch.quantization.triton.fp4_kernel import compute_fp4_scales - - quantizer = self.module.weight_quantizer - block_sizes = getattr(quantizer, "block_sizes", None) - quant_block_size = None - if block_sizes is not None: - 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 - - # Pre-compute scales from the calibrated amax (frozen during GPTQ). - amax = quantizer.amax.reshape(out_features, n_blocks) - scales_2d = compute_fp4_scales(amax, quantizer.global_amax, quantize_block_scales=True) - - gptq_blockwise_update_fused_scalar( - self.weight, scales_2d, self.h_inv, block_size, quant_block_size - ) - - _GPTQ_HELPER_REGISTRY: dict[str, type[GPTQHelper]] = {} @@ -342,7 +321,3 @@ def register_gptq_helper(backend: str, factory: type[GPTQHelper]) -> None: construct ``factory`` instead of the default ``GPTQHelper``. """ _GPTQ_HELPER_REGISTRY[backend] = factory - - -# Built-in registrations -register_gptq_helper("fused_gptq_nvfp4", FusedScalarGPTQHelper) diff --git a/tests/gpu/torch/quantization/test_gptq.py b/tests/gpu/torch/quantization/test_gptq.py index 7f2ad253b26..f04be0e4729 100644 --- a/tests/gpu/torch/quantization/test_gptq.py +++ b/tests/gpu/torch/quantization/test_gptq.py @@ -25,6 +25,7 @@ from modelopt.torch.export.unified_export_hf import _export_quantized_weight from modelopt.torch.quantization.model_calib import gptq from modelopt.torch.quantization.qtensor.nvfp4_tensor import NVFP4QTensor +from modelopt.torch.quantization.utils import promote_nvfp4_static_quantizers from modelopt.torch.quantization.utils.calib_utils import ( compute_hessian_inverse, gptq_blockwise_update, @@ -246,51 +247,57 @@ def test_gptq_e2e_flow(quant_cfg): def _make_nvfp4_test_data(quant_block_size, out_features, dim): - """Create weight, h_inv, and scales_2d for NVFP4 GPTQ tests.""" - from modelopt.torch.quantization.triton.fp4_kernel import compute_fp4_scales + """Create weight, weight_quantizer, block_amax, global_scale, and h_inv for NVFP4 GPTQ tests.""" + # Build a quantized Linear with NVFP4 static config at the desired block size + model = torch.nn.Linear(dim, out_features, bias=False, device="cuda") + weight = model.weight.data.clone() + + nvfp4_static_cfg = { + "num_bits": (2, 1), + "block_sizes": {-1: quant_block_size, "type": "static", "scale_bits": (4, 3)}, + } + quant_cfg = { + "quant_cfg": [ + {"quantizer_name": "*", "enable": False}, + {"quantizer_name": "*weight_quantizer", "cfg": nvfp4_static_cfg}, + ], + "algorithm": "max", + } + inp = torch.randn(4, 32, dim, device="cuda") + mtq.quantize(model, quant_cfg, forward_loop=lambda m: m(inp)) + promote_nvfp4_static_quantizers(model) - weight = torch.randn(out_features, dim, device="cuda", dtype=torch.float32) - n_blocks = dim // quant_block_size - amax = weight.reshape(out_features, n_blocks, quant_block_size).abs().amax(dim=-1) - scales_2d = compute_fp4_scales(amax) + # Restore original weight (GPTQ operates on original weights) + model.weight.data = weight.clone() - inp = torch.randn(4, 32, dim, device="cuda") + weight_quantizer = model.weight_quantizer + block_amax = weight_quantizer.amax.reshape(out_features, -1).float() + global_scale = weight_quantizer.global_amax.float().item() / (6.0 * 448.0) + + # Compute Hessian hessian = torch.zeros(dim, dim, dtype=torch.float32) hessian, _ = update_hessian(inp, hessian, 0) hessian = hessian.to("cuda") h_inv = compute_hessian_inverse(hessian, weight, perc_damp=0.01) - return weight, scales_2d, h_inv + return weight, weight_quantizer, block_amax, global_scale, h_inv -def _run_unfused_gptq_nvfp4(weight, scales_2d, h_inv, gptq_block_size, quant_block_size): - """Unfused NVFP4 GPTQ using the production blockwise update with Triton FP4 kernel. - - Both fused and unfused use the same frozen pre-computed scales so the - test verifies the fused kernel's correctness (not scale computation). - """ - from modelopt.torch.quantization.triton.fp4_kernel import static_blockwise_fp4_fake_quant - - out_features, num_cols = weight.shape - n_blocks = num_cols // quant_block_size +def _run_unfused_gptq_nvfp4(weight, weight_quantizer, h_inv, gptq_block_size): + """Unfused NVFP4 GPTQ using the production blockwise update with weight_quantizer.""" w = weight.float().clone() - # Recover amax from scales (scales = amax / 6.0, already FP8-quantized) - amax_flat = (scales_2d * 6.0).reshape(out_features * n_blocks) - - def quantize_fn(w_input): - w_blocked = w_input.reshape(out_features * n_blocks, quant_block_size) - return static_blockwise_fp4_fake_quant( - w_blocked, amax_flat, quantize_block_scales=False - ).reshape(out_features, num_cols) - - gptq_blockwise_update(w, h_inv, gptq_block_size, quantize_fn) + gptq_blockwise_update(w, h_inv, gptq_block_size, weight_quantizer) return w -def _run_fused_gptq_nvfp4(weight, scales_2d, h_inv, gptq_block_size, quant_block_size): +def _run_fused_gptq_nvfp4( + weight, block_amax, global_scale, h_inv, gptq_block_size, quant_block_size +): """Fused Triton GPTQ for NVFP4 using the production fused update.""" w = weight.float().clone() - gptq_blockwise_update_fused_scalar(w, scales_2d, h_inv, gptq_block_size, quant_block_size) + gptq_blockwise_update_fused_scalar( + w, block_amax, global_scale, h_inv, gptq_block_size, quant_block_size + ) return w @@ -307,7 +314,7 @@ def test_fused_vs_unfused_nvfp4(quant_block_size, gptq_block_size): dim = max(256, quant_block_size * 4) out_features = 64 - weight, scales_2d, h_inv = _make_nvfp4_test_data( + weight, weight_quantizer, block_amax, global_scale, h_inv = _make_nvfp4_test_data( quant_block_size, out_features, dim, @@ -315,17 +322,17 @@ def test_fused_vs_unfused_nvfp4(quant_block_size, gptq_block_size): weight_fused = _run_fused_gptq_nvfp4( weight, - scales_2d, + block_amax, + global_scale, h_inv, gptq_block_size, quant_block_size, ) weight_unfused = _run_unfused_gptq_nvfp4( weight, - scales_2d, + weight_quantizer, h_inv, gptq_block_size, - quant_block_size, ) assert not torch.equal(weight_fused, weight.float()), "Fused did not update weights" @@ -376,17 +383,17 @@ def _bench(fn, n_warmup=2, n_iters=5): for quant_block_size, gptq_block_size, out_features, dim in _NVFP4_BENCH_CONFIGS: torch.manual_seed(42) - weight, scales_2d, h_inv = _make_nvfp4_test_data(quant_block_size, out_features, dim) + weight, weight_quantizer, block_amax, global_scale, h_inv = _make_nvfp4_test_data( + quant_block_size, out_features, dim + ) def run_fused(): return _run_fused_gptq_nvfp4( - weight, scales_2d, h_inv, gptq_block_size, quant_block_size + weight, block_amax, global_scale, h_inv, gptq_block_size, quant_block_size ) def run_unfused(): - return _run_unfused_gptq_nvfp4( - weight, scales_2d, h_inv, gptq_block_size, quant_block_size - ) + return _run_unfused_gptq_nvfp4(weight, weight_quantizer, h_inv, gptq_block_size) t_fused = _bench(run_fused) t_unfused = _bench(run_unfused)