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.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..624e723b957 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, 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) 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..c070eac8800 --- /dev/null +++ b/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 diff --git a/modelopt/torch/quantization/triton/nvfp4_quant.py b/modelopt/torch/quantization/triton/nvfp4_quant.py new file mode 100644 index 00000000000..32ab776b2b4 --- /dev/null +++ b/modelopt/torch/quantization/triton/nvfp4_quant.py @@ -0,0 +1,144 @@ +# 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 + + +@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 252f0af6fc8..ac2ec7a2553 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. @@ -90,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): @@ -154,73 +191,40 @@ 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. + """Column-wise GPTQ update. - 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. + 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})" ) - 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 + 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.""" @@ -231,6 +235,81 @@ 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, 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 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. + 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. + 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(), + block_amax, + global_scale, + 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) + + _GPTQ_HELPER_REGISTRY: dict[str, type[GPTQHelper]] = {} 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..f04be0e4729 100644 --- a/tests/gpu/torch/quantization/test_gptq.py +++ b/tests/gpu/torch/quantization/test_gptq.py @@ -14,16 +14,24 @@ # 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.utils.calib_utils import update_hessian +from modelopt.torch.quantization.utils import promote_nvfp4_static_quantizers +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 @@ -231,3 +239,172 @@ 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 _make_nvfp4_test_data(quant_block_size, out_features, dim): + """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) + + # Restore original weight (GPTQ operates on original weights) + model.weight.data = weight.clone() + + 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, weight_quantizer, block_amax, global_scale, h_inv + + +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() + gptq_blockwise_update(w, h_inv, gptq_block_size, weight_quantizer) + return w + + +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, block_amax, global_scale, h_inv, gptq_block_size, quant_block_size + ) + 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, weight_quantizer, block_amax, global_scale, h_inv = _make_nvfp4_test_data( + quant_block_size, + out_features, + dim, + ) + + weight_fused = _run_fused_gptq_nvfp4( + weight, + block_amax, + global_scale, + h_inv, + gptq_block_size, + quant_block_size, + ) + weight_unfused = _run_unfused_gptq_nvfp4( + weight, + weight_quantizer, + h_inv, + gptq_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), +] + + +def bench_fused_nvfp4(): + """Benchmark fused Triton NVFP4 GPTQ vs unfused production loop (informational-only). + + 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): + 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 + + for quant_block_size, gptq_block_size, out_features, dim in _NVFP4_BENCH_CONFIGS: + torch.manual_seed(42) + 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, block_amax, global_scale, h_inv, gptq_block_size, quant_block_size + ) + + def run_unfused(): + return _run_unfused_gptq_nvfp4(weight, weight_quantizer, h_inv, gptq_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" + ) + + +if __name__ == "__main__": + bench_fused_nvfp4()