From 1244f173bc9ce0066319c22f22c928b1580a36ee Mon Sep 17 00:00:00 2001 From: Chenjie Luo Date: Mon, 19 Oct 2026 20:41:00 +0000 Subject: [PATCH] Launch the fused NVFP4 GPTQ kernel on its tensors' GPU Triton launches on the current CUDA device. A model that device_map="auto" spreads over several GPUs keeps later decoder layers off the current one, so GPTQ with fused: true launched the kernel on the wrong GPU with pointers into another GPU's memory: a CUDA illegal memory access on NVLink systems (GB300), a Triton pointer error on PCIe ones. Launch it under torch.cuda.device of the weight block, as the other quantization kernels do. Co-Authored-By: Claude Opus 5.5 (1M context) Signed-off-by: Chenjie Luo --- CHANGELOG.rst | 1 + .../quantization/gemm/gptq_fused_kernel.py | 29 ++++++++++--------- tests/gpu/torch/quantization/test_gptq.py | 14 +++++++++ 3 files changed, 31 insertions(+), 13 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index c3a9530c4c3..351b25ce71d 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -44,6 +44,7 @@ Changelog - The default PTQ recipes now leave multi-token-prediction (MTP) layers unquantized in Megatron vision-language models too (``language_model.mtp.*``), as they already did at the model root. - Fix ONNX quantization placing a per-tensor Q/DQ pair on a rank-0 (scalar) constant input, which TensorRT's ONNX parser rejects (``convertAxis`` asserts the quantization axis lies in ``[0, nbDims]`` but ``nbDims == 0``). Rank-0 constant operands are now kept unquantized in the int8, fp8, and autotune export paths; the existing single-element Mul/Div/Pad stripping stays gated under ``target_dla``. - Fix PTQ of native Transformers FP8 experts to apply source weight scales during calibration and export. Re-quantize affected checkpoints from the original FP8 source. +- Fix GPTQ with ``fused: true`` failing with a CUDA illegal memory access, or a Triton pointer error, on a model that ``device_map="auto"`` spreads over several GPUs once calibration reaches a layer outside the current GPU. - Fix ONNX preprocessing with custom TensorRT plugin libraries on TensorRT 11. - Fix ``INT8_DEFAULT_CFG``, ``FP8_DEFAULT_CFG``, ``NVFP4_DEFAULT_CFG`` and the other presets and recipes that import ``configs/ptq/units/default_disabled_quantizers`` and enable ``*input_quantizer``: since 0.44 they also quantized ``nn.LayerNorm`` inputs, which adds quantization error to the residual stream of models built on ``nn.LayerNorm`` such as BERT-style encoders. Re-quantize affected models; a recipe that wants ``nn.LayerNorm`` quantizers must enable them after importing that unit. - Fix ONNX quantization with the default ``concat_elimination`` pass crashing in calibration with ``KeyError: '_ReduceMin'`` when a ``Concat`` input is a constant, such as the class token of a ViT. diff --git a/modelopt/torch/kernels/quantization/gemm/gptq_fused_kernel.py b/modelopt/torch/kernels/quantization/gemm/gptq_fused_kernel.py index fb13386f851..3b410e97bae 100644 --- a/modelopt/torch/kernels/quantization/gemm/gptq_fused_kernel.py +++ b/modelopt/torch/kernels/quantization/gemm/gptq_fused_kernel.py @@ -119,18 +119,21 @@ def gptq_fused_block_scalar( 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, - ) + # Triton launches on the current device; a model sharded with device_map="auto" keeps later + # layers on other GPUs, where launching from the current one reads another GPU's memory. + with torch.cuda.device(w_block.device): + _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/tests/gpu/torch/quantization/test_gptq.py b/tests/gpu/torch/quantization/test_gptq.py index 12b387450d3..5388d8c6517 100644 --- a/tests/gpu/torch/quantization/test_gptq.py +++ b/tests/gpu/torch/quantization/test_gptq.py @@ -208,6 +208,20 @@ def _run_fused_gptq_nvfp4( return w +@requires_triton +def test_fused_nvfp4_gptq_runs_on_the_weights_gpu(need_2_gpus): + """A layer device_map placed on another GPU is updated there, not from the current GPU.""" + torch.manual_seed(42) + weight, _, block_amax, global_scale, h_inv = _make_nvfp4_test_data(16, 64, 256) + expected = _run_fused_gptq_nvfp4(weight, block_amax, global_scale, h_inv, 128, 16) + + weight, block_amax, h_inv = (t.to("cuda:1") for t in (weight, block_amax, h_inv)) + got = _run_fused_gptq_nvfp4(weight, block_amax, global_scale, h_inv, 128, 16) + + assert torch.cuda.current_device() == 0 + torch.testing.assert_close(got.cpu(), expected.cpu(), rtol=0, atol=0) + + _NVFP4_QUANT_BLOCK_SIZES = [16, 128] _NVFP4_GPTQ_BLOCK_SIZES = [16, 128]