Launch the fused NVFP4 GPTQ kernel on its tensors' GPU

CHANGELOG.rst

@@ -45,6 +45,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 ONNX remote AutoQDQ safety benchmarks returning infinite latency by running the generated engine with ``trtexec_safe`` on the configured target.
 - 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.

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

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]