[TRTLLM-13581][test] Add moe fp4 gpu dequant and unpack optimization (NVIDIA#15714)

Signed-off-by: rosong11 <rosong@nvidia.com>

Author: rosong11

cpp/tensorrt_llm/kernels/quantization.cu

@@ -429,6 +429,71 @@ template void invokeFP4Quantization<__nv_fp8_e4m3, 32>(int b, int m, int n, __nv
     int multiProcessorCount, cudaStream_t stream);
 #endif
 
+__global__ void unpackInt4PackedTensorToInt8Kernel(int8_t* output, int8_t const* input, int64_t numPacked)
+{
+    for (int64_t idx = static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x; idx < numPacked;
+         idx += static_cast<int64_t>(gridDim.x) * blockDim.x)
+    {
+        int8_t const packed = input[idx];
+        // The double shift sign-extends the low nibble; the high nibble uses an arithmetic shift.
+        int8_t const elt0 = static_cast<int8_t>(packed << 4) >> 4;
+        int8_t const elt1 = static_cast<int8_t>(packed >> 4);
+        output[2 * idx] = elt0;
+        output[2 * idx + 1] = elt1;
+    }
+}
+
+__global__ void mxfp4DequantizeUnswizzledKernel(float* output, uint8_t const* weight, uint8_t const* scale,
+    int64_t numPacked, int64_t k, int64_t scaleStride, int64_t groupSize)
+{
+    // Keep the LUT identical to the CPU reference implementation in weightOnlyQuantOp.cpp.
+    float const fp4Lut[16]
+        = {0.0f, 0.5f, 1.0f, 1.5f, 2.0f, 3.0f, 4.0f, 6.0f, 0.0f, -0.5f, -1.0f, -1.5f, -2.0f, -3.0f, -4.0f, -6.0f};
+    auto const* scaleE8m0 = reinterpret_cast<__nv_fp8_e8m0 const*>(scale);
+
+    for (int64_t packedIdx = static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x; packedIdx < numPacked;
+         packedIdx += static_cast<int64_t>(gridDim.x) * blockDim.x)
+    {
+        uint8_t const packed = weight[packedIdx];
+        uint8_t const low = packed & 0xF;
+        uint8_t const high = (packed & 0xF0) >> 4;
+
+        int64_t const scaleNIdx = packedIdx / (k / 2);
+        int64_t const scaleKIdx = ((packedIdx * 2) % k) / groupSize;
+        // static_cast<float>(__nv_fp8_e8m0) decodes the e8m0 scale, matching the CPU reference exactly.
+        float const scaleValue = static_cast<float>(scaleE8m0[scaleNIdx * scaleStride + scaleKIdx]);
+
+        output[2 * packedIdx] = fp4Lut[low] * scaleValue;
+        output[2 * packedIdx + 1] = fp4Lut[high] * scaleValue;
+    }
+}
+
+void invokeUnpackInt4PackedTensorToInt8(int8_t* output, int8_t const* input, int64_t numPacked, cudaStream_t stream)
+{
+    if (numPacked <= 0)
+    {
+        return;
+    }
+    constexpr int block = 256;
+    int64_t const numBlocks = (numPacked + block - 1) / block;
+    dim3 const grid(static_cast<unsigned int>(std::min<int64_t>(numBlocks, 65535)));
+    unpackInt4PackedTensorToInt8Kernel<<<grid, block, 0, stream>>>(output, input, numPacked);
+}
+
+void invokeMxfp4DequantizeUnswizzled(float* output, uint8_t const* weight, uint8_t const* scale, int64_t numPacked,
+    int64_t k, int64_t scaleStride, int64_t groupSize, cudaStream_t stream)
+{
+    if (numPacked <= 0)
+    {
+        return;
+    }
+    constexpr int block = 256;
+    int64_t const numBlocks = (numPacked + block - 1) / block;
+    dim3 const grid(static_cast<unsigned int>(std::min<int64_t>(numBlocks, 65535)));
+    mxfp4DequantizeUnswizzledKernel<<<grid, block, 0, stream>>>(
+        output, weight, scale, numPacked, k, scaleStride, groupSize);
+}
+
 } // namespace kernels
 
 TRTLLM_NAMESPACE_END

cpp/tensorrt_llm/kernels/quantization.h

@@ -91,6 +91,17 @@ template <typename T>
 void computePerTokenGlobalScaleForFP4Quantization(int b, int m, int n, T const* input, int const* tokensPerBatch,
     float* globalScale, int multiProcessorCount, cudaStream_t stream = 0);
 
+// Unpack a packed int4 tensor (two signed 4-bit values per byte) into int8.
+// output must have 2 * numPacked elements.
+void invokeUnpackInt4PackedTensorToInt8(
+    int8_t* output, int8_t const* input, int64_t numPacked, cudaStream_t stream = 0);
+
+// Dequantize an unswizzled MXFP4 weight (two e2m1 values per byte) using a linear-layout
+// e8m0 block scale. output must have 2 * numPacked floats.
+// k is the unpacked column count (weight.size(1) * 2), scaleStride is scale.size(1).
+void invokeMxfp4DequantizeUnswizzled(float* output, uint8_t const* weight, uint8_t const* scale, int64_t numPacked,
+    int64_t k, int64_t scaleStride, int64_t groupSize, cudaStream_t stream = 0);
+
 } // namespace kernels
 
 TRTLLM_NAMESPACE_END

cpp/tensorrt_llm/thop/weightOnlyQuantOp.cpp

@@ -17,8 +17,11 @@
 #include "tensorrt_llm/common/config.h"
 #include "tensorrt_llm/common/cudaBf16Wrapper.h"
 #include "tensorrt_llm/kernels/cutlass_kernels/cutlass_preprocessors.h"
+#include "tensorrt_llm/kernels/quantization.h"
 #include "tensorrt_llm/thop/thUtils.h"
 
+#include <c10/cuda/CUDAGuard.h>
+
 #if defined(TORCH_VERSION_MAJOR)                                                                                       \
     && ((TORCH_VERSION_MAJOR > 1) || ((TORCH_VERSION_MAJOR == 1) && (TORCH_VERSION_MINOR >= 9)))
 #define TORCH_IS_AT_LEAST_v190
@@ -283,7 +286,6 @@ Tensor add_bias_and_interleave_int8s(Tensor weight)
 
 Tensor unpack_int4_packed_tensor_to_int8(Tensor weight)
 {
-    CHECK_CPU(weight);
     CHECK_CONTIGUOUS(weight);
     TORCH_CHECK(weight.numel() != 0, "weight should not be empty tensor");
     TORCH_CHECK(weight.dtype() == torch::kInt8, "Weight must be a packed int8 tensor");
@@ -295,6 +297,17 @@ Tensor unpack_int4_packed_tensor_to_int8(Tensor weight)
     }
     int8_tensor_size[weight.dim() - 1] *= 2;
 
+    if (weight.is_cuda())
+    {
+        at::cuda::CUDAGuard const device_guard(weight.device());
+        Tensor unpacked_weight
+            = torch::empty(int8_tensor_size, torch::dtype(torch::kInt8).device(weight.device()).requires_grad(false));
+        auto stream = at::cuda::getCurrentCUDAStream(weight.device().index());
+        tensorrt_llm::kernels::invokeUnpackInt4PackedTensorToInt8(
+            get_ptr<int8_t>(unpacked_weight), get_ptr<int8_t>(weight), weight.numel(), stream);
+        return unpacked_weight;
+    }
+
     Tensor unpacked_weight
         = torch::zeros(int8_tensor_size, torch::dtype(torch::kInt8).device(torch::kCPU).requires_grad(false));
 
@@ -357,8 +370,6 @@ Tensor mxfp4_dequantize_unswizzled(Tensor weight, Tensor scale, int64_t group_si
     // weight (n, k / 2)
     // scale (n, k / group_size)
 
-    CHECK_CPU(weight);
-    CHECK_CPU(scale);
     CHECK_CONTIGUOUS(weight);
     CHECK_CONTIGUOUS(scale);
     TORCH_CHECK(weight.numel() != 0, "weight should not be empty tensor");
@@ -368,12 +379,27 @@ Tensor mxfp4_dequantize_unswizzled(Tensor weight, Tensor scale, int64_t group_si
     TORCH_CHECK(weight.size(0) == scale.size(0))
     TORCH_CHECK(weight.size(1) * 2 == scale.size(1) * group_size)
 
-    uint8_t* weight_packed_ptr = get_ptr<uint8_t>(weight);
-    __nv_fp8_e8m0* scale_ptr = reinterpret_cast<__nv_fp8_e8m0*>(get_ptr<uint8_t>(scale));
-
     int const n = weight.size(0);
     int const k = weight.size(1) * 2;
 
+    if (weight.is_cuda())
+    {
+        TORCH_CHECK(scale.is_cuda(), "scale must be a CUDA tensor when weight is on CUDA");
+        TORCH_CHECK(weight.device() == scale.device(), "weight and scale must be on the same device");
+        at::cuda::CUDAGuard const device_guard(weight.device());
+        Tensor dequant_weight
+            = torch::empty({n, k}, torch::dtype(torch::kFloat).device(weight.device()).requires_grad(false));
+        auto stream = at::cuda::getCurrentCUDAStream(weight.device().index());
+        tensorrt_llm::kernels::invokeMxfp4DequantizeUnswizzled(get_ptr<float>(dequant_weight), get_ptr<uint8_t>(weight),
+            get_ptr<uint8_t>(scale), weight.numel(), k, scale.size(1), group_size, stream);
+        return dequant_weight;
+    }
+
+    TORCH_CHECK(scale.is_cpu(), "scale must be a CPU tensor when weight is on CPU");
+
+    uint8_t* weight_packed_ptr = get_ptr<uint8_t>(weight);
+    __nv_fp8_e8m0* scale_ptr = reinterpret_cast<__nv_fp8_e8m0*>(get_ptr<uint8_t>(scale));
+
     Tensor dequant_weight = torch::empty({n, k}, torch::dtype(torch::kFloat).device(torch::kCPU).requires_grad(false));
     float* dequant_weight_ptr = get_ptr<float>(dequant_weight);
 

tests/unittest/_torch/modules/moe/quantize_utils.py

@@ -2182,17 +2182,17 @@ def load_weights(self, weights_list: List[Dict]):
             # mxfp4_dequantize_unswizzled returns (out_features, in_features)
             # which matches F.linear weight layout (out, in). Do NOT transpose.
             w1_dequant = (
-                unpacker(w1.cpu(), s1.cpu(), scaling_group_size)
+                unpacker(w1, s1, scaling_group_size)
                 .to(dtype=self.dtype, device="cuda")
                 .contiguous()
             )
             w3_dequant = (
-                unpacker(w3.cpu(), s3.cpu(), scaling_group_size)
+                unpacker(w3, s3, scaling_group_size)
                 .to(dtype=self.dtype, device="cuda")
                 .contiguous()
             )
             w2_dequant = (
-                unpacker(w2.cpu(), s2.cpu(), scaling_group_size)
+                unpacker(w2, s2, scaling_group_size)
                 .to(dtype=self.dtype, device="cuda")
                 .contiguous()
             )
@@ -2439,17 +2439,17 @@ def load_weights(self, weights_list: List[Dict]):
             # Note: mxfp4_dequantize_unswizzled returns shape (out_features, in_features)
             # which matches F.linear weight layout (out, in). Do NOT transpose.
             w1_dequant = (
-                unpacker(w1.cpu(), s1.cpu(), scaling_group_size)
+                unpacker(w1, s1, scaling_group_size)
                 .to(dtype=self.dtype, device="cuda")
                 .contiguous()
             )
             w3_dequant = (
-                unpacker(w3.cpu(), s3.cpu(), scaling_group_size)
+                unpacker(w3, s3, scaling_group_size)
                 .to(dtype=self.dtype, device="cuda")
                 .contiguous()
             )
             w2_dequant = (
-                unpacker(w2.cpu(), s2.cpu(), scaling_group_size)
+                unpacker(w2, s2, scaling_group_size)
                 .to(dtype=self.dtype, device="cuda")
                 .contiguous()
             )
@@ -2770,10 +2770,10 @@ def _unpack_weights(self, weight: torch.Tensor) -> torch.Tensor:
         unpacker = torch.ops.trtllm.unpack_int4_packed_tensor_to_int8
         # ModelOpt W4A8 packs pairs of 4b weights in the output dimension into one 8b element.
         if self.weight_loading_mode == MoEWeightLoadingMode.VANILLA:
-            return unpacker(weight.cpu().T.contiguous()).cuda()
+            return unpacker(weight.T.contiguous()).cuda()
         # The custom W4A8 quantization script packs pairs of 4b weight in the input dimension.
         else:
-            return unpacker(weight.cpu()).T.contiguous().cuda()
+            return unpacker(weight.contiguous()).T.contiguous().cuda()
 
     def load_weights(self, weights: List[Dict]):
         """Store raw quantized weights for forward computation."""

tests/unittest/_torch/thop/parallel/test_weight_only_quant_ops.py

@@ -0,0 +1,110 @@
+# 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.
+"""GPU vs CPU golden tests for trtllm weight-only quant utility ops.
+
+These ops used to be CPU-only; the CUDA branch must produce bit-identical
+results to the original CPU reference implementation in weightOnlyQuantOp.cpp.
+"""
+
+import pytest
+import torch
+
+import tensorrt_llm  # noqa: F401  (registers the trtllm torch ops)
+
+pytestmark = pytest.mark.skipif(
+    not torch.cuda.is_available(), reason="CUDA is required for GPU vs CPU comparison"
+)
+
+
+@pytest.mark.parametrize(
+    "shape",
+    [
+        (8, 16),
+        (4, 32),
+        (1, 1),
+        (3, 5, 7),
+        (2, 3, 4, 9),
+    ],
+)
+def test_unpack_int4_packed_tensor_to_int8_gpu_matches_cpu(shape):
+    op = torch.ops.trtllm.unpack_int4_packed_tensor_to_int8
+
+    packed = torch.randint(-128, 128, shape, dtype=torch.int8)
+
+    cpu_out = op(packed)
+    gpu_out = op(packed.cuda())
+
+    assert gpu_out.is_cuda
+    assert gpu_out.dtype == torch.int8
+    # Output last dim is doubled, everything else identical.
+    assert list(gpu_out.shape[:-1]) == list(packed.shape[:-1])
+    assert gpu_out.shape[-1] == packed.shape[-1] * 2
+
+    # Integer unpack must be exactly equal.
+    assert torch.equal(gpu_out.cpu(), cpu_out)
+
+
+@pytest.mark.parametrize("group_size", [32, 128])
+@pytest.mark.parametrize(
+    "n, k",
+    [
+        (8, 256),
+        (4, 512),
+        (1, 128),
+        (16, 1024),
+    ],
+)
+def test_mxfp4_dequantize_unswizzled_gpu_matches_cpu(n, k, group_size):
+    op = torch.ops.trtllm.mxfp4_dequantize_unswizzled
+
+    # weight: (n, k/2), two packed e2m1 codes per byte -> any uint8 value is valid.
+    weight = torch.randint(0, 256, (n, k // 2), dtype=torch.uint8)
+
+    # scale: (n, k/group_size), interpreted as e8m0 (2^(e-127)).
+    # Restrict the exponent to a stable finite range so neither CPU nor GPU
+    # produces NaN (0xFF) or Inf (very large exponents); see plan doc.
+    scale = torch.randint(118, 135, (n, k // group_size), dtype=torch.uint8)
+
+    cpu_out = op(weight, scale, group_size)
+    gpu_out = op(weight.cuda(), scale.cuda(), group_size)
+
+    assert gpu_out.is_cuda
+    assert gpu_out.dtype == torch.float32
+    assert list(gpu_out.shape) == [n, k]
+
+    # LUT value * power-of-two scale, same decode on both sides -> bit exact.
+    torch.testing.assert_close(gpu_out.cpu(), cpu_out, rtol=0, atol=0)
+
+
+def test_unpack_int4_non_contiguous_raises():
+    op = torch.ops.trtllm.unpack_int4_packed_tensor_to_int8
+    packed = torch.randint(-128, 128, (8, 16), dtype=torch.int8).cuda()
+    non_contiguous = packed.t()
+    assert not non_contiguous.is_contiguous()
+    with pytest.raises(RuntimeError):
+        op(non_contiguous)
+
+
+def test_mxfp4_dequantize_device_mismatch_raises():
+    op = torch.ops.trtllm.mxfp4_dequantize_unswizzled
+    n, k, group_size = 4, 256, 32
+    weight = torch.randint(0, 256, (n, k // 2), dtype=torch.uint8).cuda()
+    scale = torch.randint(118, 135, (n, k // group_size), dtype=torch.uint8)  # CPU
+    with pytest.raises(RuntimeError):
+        op(weight, scale, group_size)
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])