[https://nvbugs/5289907][fix] Restore per-channel pre-quant

cpp/tensorrt_llm/kernels/internal_cutlass_kernels/aarch64-linux-gnu/version.txt

@@ -1,2 +1,2 @@
-aff0f8e617f6ca2f95d121ab9cf0ab17c4e8077cf9e8896bf153d3942a4a50df  libtensorrt_llm_internal_cutlass_kernels_static.a
-commit d61e7684bc095c8ff5ec540363949bd1f491c960
+c05fd8d4a1e1fbcf5510ae3486f9094899688eb4d2f86a7befc71284da5433b0  libtensorrt_llm_internal_cutlass_kernels_static.a
+commit 1b03122f38f89025ec6d3af1f06531811e8c71b6

cpp/tensorrt_llm/kernels/internal_cutlass_kernels/x86_64-linux-gnu/version.txt

@@ -1,2 +1,2 @@
-9f0a29070b95a7db62f70cc45ef151e27c2a58697a2d50cbb002ff339035fb8e  libtensorrt_llm_internal_cutlass_kernels_static.a
-commit d61e7684bc095c8ff5ec540363949bd1f491c960
+981cba5aa914fef513dfd2e6daa90a5e0fec351a79ea0e73508f8afb1e9b982f  libtensorrt_llm_internal_cutlass_kernels_static.a
+commit 1b03122f38f89025ec6d3af1f06531811e8c71b6

cpp/tensorrt_llm/kernels/preQuantScaleKernel.cu

@@ -40,15 +40,17 @@ struct Vec2Type<__nv_bfloat16>
 }; // namespace
 
 template <typename T_in, typename T_out, int kProcessRows, typename AccessType>
-__global__ void apply_per_channel_scale(
-    T_out* smoothed_act, T_in const* act, T_in const* per_channel_scale, int rows, int cols)
+__global__ void apply_per_channel_scale(T_out* smoothed_act, T_in const* act, T_in const* per_channel_scale, int rows,
+    int cols, int64_t const* num_valid_tokens_ptr)
 {
     static constexpr int kElems = sizeof(AccessType) / sizeof(T_in);
     T_in scale[kElems], act_vec[kElems];
     int col_offset = blockIdx.y * blockDim.x + threadIdx.x;
     int row_offset = blockIdx.x;
     if (col_offset * kElems >= cols || row_offset * kProcessRows >= rows)
         return;
+    if (num_valid_tokens_ptr && (row_offset * kProcessRows >= *num_valid_tokens_ptr))
+        return;
     act += row_offset * kProcessRows * cols;
     smoothed_act += row_offset * kProcessRows * cols;
     *reinterpret_cast<AccessType*>(scale) = reinterpret_cast<AccessType const*>(per_channel_scale)[col_offset];
@@ -95,46 +97,46 @@ __global__ void apply_per_channel_scale(
 }
 
 template <typename T_in, typename T_out, int kProcessRows, typename AccessType = float4>
-void apply_per_channel_scale_kernel_launcher_(
-    T_out* smoothed_act, T_in const* act, T_in const* per_channel_scale, int rows, int cols, cudaStream_t stream = 0)
+void apply_per_channel_scale_kernel_launcher_(T_out* smoothed_act, T_in const* act, T_in const* per_channel_scale,
+    int rows, int cols, int64_t const* num_valid_tokens_ptr = nullptr, cudaStream_t stream = 0)
 {
     static constexpr int kElems = sizeof(AccessType) / sizeof(T_in);
     dim3 block(128);
     dim3 grid((rows + kProcessRows - 1) / kProcessRows, (cols / kElems + block.x - 1) / block.x);
     apply_per_channel_scale<T_in, T_out, kProcessRows, AccessType>
-        <<<grid, block, 0, stream>>>(smoothed_act, act, per_channel_scale, rows, cols);
+        <<<grid, block, 0, stream>>>(smoothed_act, act, per_channel_scale, rows, cols, num_valid_tokens_ptr);
 }
 
 template <typename T_in, typename T_out>
-void apply_per_channel_scale_kernel_launcher(
-    T_out* smoothed_act, T_in const* act, T_in const* per_channel_scale, int rows, int cols, cudaStream_t stream)
+void apply_per_channel_scale_kernel_launcher(T_out* smoothed_act, T_in const* act, T_in const* per_channel_scale,
+    int rows, int cols, int64_t const* num_valid_tokens_ptr, cudaStream_t stream)
 {
     uint64_t elems = static_cast<uint64_t>(rows) * static_cast<uint64_t>(cols);
     if (elems < 2048 * 2048)
     {
         apply_per_channel_scale_kernel_launcher_<T_in, T_out, 1, float4>(
-            smoothed_act, act, per_channel_scale, rows, cols, stream);
+            smoothed_act, act, per_channel_scale, rows, cols, num_valid_tokens_ptr, stream);
     }
     else if (elems < 4096 * 4096)
     {
         apply_per_channel_scale_kernel_launcher_<T_in, T_out, 4, float4>(
-            smoothed_act, act, per_channel_scale, rows, cols, stream);
+            smoothed_act, act, per_channel_scale, rows, cols, num_valid_tokens_ptr, stream);
     }
     else if (elems < 8192 * 8192)
     {
         apply_per_channel_scale_kernel_launcher_<T_in, T_out, 8, float4>(
-            smoothed_act, act, per_channel_scale, rows, cols, stream);
+            smoothed_act, act, per_channel_scale, rows, cols, num_valid_tokens_ptr, stream);
     }
     else
     {
         apply_per_channel_scale_kernel_launcher_<T_in, T_out, 16, float4>(
-            smoothed_act, act, per_channel_scale, rows, cols, stream);
+            smoothed_act, act, per_channel_scale, rows, cols, num_valid_tokens_ptr, stream);
     }
 }
 
 #define INSTANTIATE_PREQUANT_SCALE(T_in, T_out)                                                                        \
-    template void apply_per_channel_scale_kernel_launcher<T_in, T_out>(                                                \
-        T_out * smoothed_act, const T_in* act, const T_in* per_channel_scale, int rows, int cols, cudaStream_t stream)
+    template void apply_per_channel_scale_kernel_launcher<T_in, T_out>(T_out * smoothed_act, const T_in* act,          \
+        const T_in* per_channel_scale, int rows, int cols, int64_t const* num_valid_tokens_ptr, cudaStream_t stream)
 
 INSTANTIATE_PREQUANT_SCALE(half, half);
 #if defined(ENABLE_FP8)
@@ -148,128 +150,5 @@ INSTANTIATE_PREQUANT_SCALE(__nv_bfloat16, __nv_fp8_e4m3);
 #endif
 #endif
 
-template <typename T_in, typename T_out, int kProcessRows, typename AccessType>
-__global__ void apply_per_expert_scale(T_out* smoothed_act, T_in const* act, T_in const* per_expert_scale,
-    int const* permuted_token_selected_experts, int64_t const* num_valid_tokens_ptr, int rows, int cols)
-{
-    static constexpr int kElems = sizeof(AccessType) / sizeof(T_in);
-    T_in act_vec[kElems];
-    int col_offset = blockIdx.x * blockDim.x + threadIdx.x;
-    int row_offset = blockIdx.y;
-    int expert_idx = permuted_token_selected_experts[row_offset];
-    T_in scale = per_expert_scale[expert_idx];
-    if (col_offset * kElems >= cols || row_offset * kProcessRows >= rows)
-        return;
-    if (num_valid_tokens_ptr && (row_offset * kProcessRows >= *num_valid_tokens_ptr))
-        return;
-    act += row_offset * kProcessRows * cols;
-    smoothed_act += row_offset * kProcessRows * cols;
-#pragma unroll
-    for (int i = 0; i < kProcessRows; ++i)
-    {
-        *reinterpret_cast<AccessType*>(act_vec) = reinterpret_cast<AccessType const*>(act + i * cols)[col_offset];
-        if constexpr ((std::is_same_v<T_in, half>
-#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800) && defined(ENABLE_BF16))
-                          || std::is_same_v<T_in, __nv_bfloat16>
-#endif
-                          ) &&(kElems % 2 == 0))
-        {
-            using Vec2 = typename Vec2Type<T_in>::type;
-#pragma unroll
-            for (int j = 0; j < kElems; j += 2)
-            {
-                if constexpr (std::is_same_v<T_in, half>)
-                {
-                    *reinterpret_cast<Vec2*>(act_vec + j)
-                        = __hmul2(*reinterpret_cast<Vec2*>(act_vec + j), __half2half2(scale));
-                }
-                else
-                {
-                    *reinterpret_cast<Vec2*>(act_vec + j)
-                        = __hmul2(*reinterpret_cast<Vec2*>(act_vec + j), __bfloat162bfloat162(scale));
-                }
-            }
-        }
-        else
-        {
-#pragma unroll
-            for (int j = 0; j < kElems; ++j)
-            {
-                act_vec[j] = static_cast<T_in>(static_cast<float>(act_vec[j]) * static_cast<float>(scale));
-            }
-        }
-        if constexpr (std::is_same_v<T_in, T_out>)
-        {
-            reinterpret_cast<AccessType*>(smoothed_act + i * cols)[col_offset]
-                = *reinterpret_cast<AccessType*>(act_vec);
-        }
-        else
-        {
-#pragma unroll
-            for (int j = 0; j < kElems; ++j)
-            {
-                (smoothed_act + i * cols)[col_offset * kElems + j] = static_cast<T_out>(act_vec[j]);
-            }
-        }
-    }
-}
-
-template <typename T_in, typename T_out, int kProcessRows, typename AccessType = float4>
-void apply_per_expert_scale_kernel_launcher_(T_out* smoothed_act, T_in const* act, T_in const* per_expert_scale,
-    int const* permuted_token_selected_experts, int64_t const* num_valid_tokens_ptr, int rows, int cols,
-    cudaStream_t stream = 0)
-{
-    static constexpr int kElems = sizeof(AccessType) / sizeof(T_in);
-    dim3 block(128);
-    dim3 grid((cols / kElems + block.x - 1) / block.x, (rows + kProcessRows - 1) / kProcessRows);
-    apply_per_expert_scale<T_in, T_out, kProcessRows, AccessType><<<grid, block, 0, stream>>>(
-        smoothed_act, act, per_expert_scale, permuted_token_selected_experts, num_valid_tokens_ptr, rows, cols);
-}
-
-template <typename T_in, typename T_out>
-void apply_per_expert_scale_kernel_launcher(T_out* smoothed_act, T_in const* act, T_in const* per_expert_scale,
-    int const* permuted_token_selected_experts, int64_t const* num_valid_tokens_ptr, int rows, int cols,
-    cudaStream_t stream)
-{
-    int elems = rows * cols;
-    if (elems < 2048 * 2048)
-    {
-        apply_per_expert_scale_kernel_launcher_<T_in, T_out, 1, float4>(smoothed_act, act, per_expert_scale,
-            permuted_token_selected_experts, num_valid_tokens_ptr, rows, cols, stream);
-    }
-    else if (elems < 4096 * 4096)
-    {
-        apply_per_expert_scale_kernel_launcher_<T_in, T_out, 4, float4>(smoothed_act, act, per_expert_scale,
-            permuted_token_selected_experts, num_valid_tokens_ptr, rows, cols, stream);
-    }
-    else if (elems < 8192 * 8192)
-    {
-        apply_per_expert_scale_kernel_launcher_<T_in, T_out, 8, float4>(smoothed_act, act, per_expert_scale,
-            permuted_token_selected_experts, num_valid_tokens_ptr, rows, cols, stream);
-    }
-    else
-    {
-        apply_per_expert_scale_kernel_launcher_<T_in, T_out, 16, float4>(smoothed_act, act, per_expert_scale,
-            permuted_token_selected_experts, num_valid_tokens_ptr, rows, cols, stream);
-    }
-}
-
-#define INSTANTIATE_PEREXPERT_SCALE(T_in, T_out)                                                                       \
-    template void apply_per_expert_scale_kernel_launcher<T_in, T_out>(T_out * smoothed_act, T_in const* act,           \
-        T_in const* per_expert_scale, int const* permuted_token_selected_experts, int64_t const* num_valid_tokens_ptr, \
-        int rows, int cols, cudaStream_t stream)
-
-INSTANTIATE_PEREXPERT_SCALE(half, half);
-#if defined(ENABLE_FP8)
-INSTANTIATE_PEREXPERT_SCALE(half, __nv_fp8_e4m3);
-#endif
-
-#if defined(ENABLE_BF16)
-INSTANTIATE_PEREXPERT_SCALE(__nv_bfloat16, __nv_bfloat16);
-#if defined(ENABLE_FP8)
-INSTANTIATE_PEREXPERT_SCALE(__nv_bfloat16, __nv_fp8_e4m3);
-#endif
-#endif
-
 } // namespace kernels
 } // namespace tensorrt_llm

cpp/tensorrt_llm/kernels/preQuantScaleKernel.h

@@ -36,13 +36,8 @@ namespace kernels
 {
 
 template <typename T_in, typename T_out = T_in>
-void apply_per_channel_scale_kernel_launcher(
-    T_out* smoothed_act, T_in const* act, T_in const* per_channel_scale, int rows, int cols, cudaStream_t stream = 0);
-
-template <typename T_in, typename T_out = T_in>
-void apply_per_expert_scale_kernel_launcher(T_out* smoothed_act, T_in const* act, T_in const* per_expert_scale,
-    int const* permuted_token_selected_experts, int64_t const* num_valid_tokens_ptr, int rows, int cols,
-    cudaStream_t stream = 0);
+void apply_per_channel_scale_kernel_launcher(T_out* smoothed_act, T_in const* act, T_in const* per_channel_scale,
+    int rows, int cols, int64_t const* num_valid_tokens_ptr = nullptr, cudaStream_t stream = 0);
 
 } // namespace kernels
 } // namespace tensorrt_llm

cpp/tensorrt_llm/plugins/weightOnlyGroupwiseQuantMatmulPlugin/weightOnlyGroupwiseQuantMatmulPlugin.cpp

@@ -394,13 +394,13 @@ void pre_quant_scale_for_act(int const m, int const k, int const mQuantAlgo, int
     {
         tensorrt_llm::kernels::apply_per_channel_scale_kernel_launcher<ActType, __nv_fp8_e4m3>(
             reinterpret_cast<__nv_fp8_e4m3*>(workspace), reinterpret_cast<ActType const*>(inputs[0]),
-            reinterpret_cast<ActType const*>(inputs[mPreQuantScaleInputIdx]), m, k, stream);
+            reinterpret_cast<ActType const*>(inputs[mPreQuantScaleInputIdx]), m, k, nullptr, stream);
     }
     else
     {
         tensorrt_llm::kernels::apply_per_channel_scale_kernel_launcher<ActType, ActType>(
             reinterpret_cast<ActType*>(workspace), reinterpret_cast<ActType const*>(inputs[0]),
-            reinterpret_cast<ActType const*>(inputs[mPreQuantScaleInputIdx]), m, k, stream);
+            reinterpret_cast<ActType const*>(inputs[mPreQuantScaleInputIdx]), m, k, nullptr, stream);
     }
 }
 

cpp/tests/unit_tests/kernels/weightOnly/weightOnlyKernelTest.cpp

@@ -220,7 +220,7 @@ void exec_cutlass_kernel(
     {
         tensorrt_llm::kernels::apply_per_channel_scale_kernel_launcher<AType, AType>(
             reinterpret_cast<AType*>(scaled_act), reinterpret_cast<AType const*>(params.act),
-            reinterpret_cast<AType const*>(params.act_scale), params.m, params.k, stream);
+            reinterpret_cast<AType const*>(params.act_scale), params.m, params.k, nullptr, stream);
         act = scaled_act;
     }
     if constexpr (QuantOp == cutlass::WeightOnlyQuantOp::PER_COLUMN_SCALE_ONLY)

tensorrt_llm/_torch/modules/fused_moe.py

@@ -593,13 +593,17 @@ def create_weights(self):
                                    self.hidden_size,
                                    self.intermediate_size_per_partition // 2)
 
-                fc31_act_scale = nn.Parameter(torch.empty(
-                    self.expert_size_per_partition, 1, dtype=self.dtype),
+                fc31_act_scale = nn.Parameter(torch.empty(1,
+                                                          self.hidden_size,
+                                                          dtype=self.dtype),
                                               requires_grad=False)
                 self.register_parameter("fc31_act_scale", fc31_act_scale)
 
                 fc2_act_scale = nn.Parameter(torch.empty(
-                    self.expert_size_per_partition, 1, dtype=self.dtype),
+                    1,
+                    self.intermediate_size_per_partition,
+                    1,
+                    dtype=self.dtype),
                                              requires_grad=False)
                 self.register_parameter("fc2_act_scale", fc2_act_scale)
 
@@ -1670,12 +1674,14 @@ def _load_int4_groupwise_scales(self, weights: Dict):
             load_weight_shard(weights[f"{expert_id}.w1.input_scale"])
             for expert_id in range(self.expert_start, self.expert_end)
         ]
-        all_w3_w1_input_scales = torch.max(torch.stack(all_w3_input_scales),
-                                           torch.stack(all_w1_input_scales))
-        all_w3_w1_input_scales = torch.ones_like(
-            all_w3_w1_input_scales) * all_w3_w1_input_scales.max()
-        self.fc31_act_scale.data.copy_(1 / all_w3_w1_input_scales)
-        self.fc31_alpha.data.copy_(all_w3_w1_input_scales.float())
+        all_w3_w1_input_scales_max = torch.max(
+            torch.stack(all_w3_input_scales),
+            torch.stack(all_w1_input_scales)).max()
+        self.fc31_act_scale.data.copy_(
+            torch.ones_like(self.fc31_act_scale) *
+            (1 / all_w3_w1_input_scales_max))
+        self.fc31_alpha.data.copy_((torch.ones_like(self.fc31_alpha) *
+                                    all_w3_w1_input_scales_max).float())
 
         all_w3_scales = [
             load_weight_shard(weights[f"{expert_id}.w3.weight_scale_inv"],
@@ -1711,11 +1717,12 @@ def _load_int4_groupwise_scales(self, weights: Dict):
             load_weight_shard(weights[f"{expert_id}.w2.input_scale"])
             for expert_id in range(self.expert_start, self.expert_end)
         ]
-        all_w2_input_scales = torch.stack(all_w2_input_scales).to(self.dtype)
-        all_w2_input_scales = torch.ones_like(
-            all_w2_input_scales) * all_w2_input_scales.max()
-        self.fc2_act_scale.data.copy_(1 / all_w2_input_scales)
-        self.fc2_alpha.data.copy_(all_w2_input_scales.float())
+        all_w2_input_scales_max = torch.stack(all_w2_input_scales).to(
+            self.dtype).max()
+        self.fc2_act_scale.data.copy_(
+            torch.ones_like(self.fc2_act_scale) * (1 / all_w2_input_scales_max))
+        self.fc2_alpha.data.copy_(
+            (torch.ones_like(self.fc2_alpha) * all_w2_input_scales_max).float())
 
         all_w2_scales = [
             load_weight_shard(weights[f"{expert_id}.w2.weight_scale_inv"],

tensorrt_llm/layers/moe.py

@@ -500,8 +500,8 @@ def __init__(self, in_features: int, out_features: int,
             else:
                 self.register_parameter('zero', None)
             if groupwise_quant_algo & GroupwiseQuantAlgo.PRE_QUANT_SCALE:
-                self.prequant_scaling_factor = Parameter(
-                    shape=(experts_per_node, 1), dtype=dtype)
+                self.prequant_scaling_factor = Parameter(shape=(1, in_features),
+                                                         dtype=dtype)
             else:
                 self.register_parameter('prequant_scaling_factor', None)
             if groupwise_quant_algo & GroupwiseQuantAlgo.W4A8_ALPHA:

tests/unittest/trt/quantization/test_moe_weight_only_groupwise_quant_matmul.py

@@ -154,14 +154,14 @@ def _woq_moe_groupwise_matmul(self,
             2**31, (num_experts, n, k // num_weights_in_32_bits),
             dtype=torch.int32,
             device="cuda")
-        pre_quant_scale_1 = torch.ones(num_experts,
-                                       1,
-                                       dtype=activation_dtype,
-                                       device="cuda")
-        pre_quant_scale_2 = torch.ones(num_experts,
-                                       1,
-                                       dtype=activation_dtype,
-                                       device="cuda")
+        pre_quant_scale_1 = torch.randn(1,
+                                        k,
+                                        dtype=activation_dtype,
+                                        device="cuda")
+        pre_quant_scale_2 = torch.randn(1,
+                                        n,
+                                        dtype=activation_dtype,
+                                        device="cuda")
         scale_1 = torch.randn(num_experts,
                               k // group_size,
                               n * 2,
@@ -182,8 +182,10 @@ def _woq_moe_groupwise_matmul(self,
                              k,
                              dtype=activation_dtype,
                              device="cuda") * 0.01
-        alpha_1 = 1 / pre_quant_scale_1.float()
-        alpha_2 = 1 / pre_quant_scale_2.float()
+        alpha_1 = torch.randn(
+            num_experts, 1, dtype=torch.float32, device="cuda") * 0.1
+        alpha_2 = torch.randn(
+            num_experts, 1, dtype=torch.float32, device="cuda") * 0.1
 
         preprocessor = tensorrt_llm.quantization.functional.preprocess_weights_for_mixed_gemm
         unpacker = torch.ops.trtllm.unpack_int4_packed_tensor_to_int8
@@ -236,7 +238,7 @@ def _woq_moe_groupwise_matmul(self,
                 input = inputs_merged[i, :]
                 fc1_qd = ref_weight_1[expert].cuda().float()
                 if has_pre_quant:
-                    input = input * pre_quant_scale_1[expert]
+                    input = input * pre_quant_scale_1.squeeze()
                 if has_alpha:
                     input[input > 448.0] = 448.0
                     input = input.to(torch.float8_e4m3fn).float()
@@ -248,7 +250,7 @@ def _woq_moe_groupwise_matmul(self,
                 fc1 = fc1 * torch.nn.functional.silu(gate)
                 fc2_qd = ref_weight_2[expert].cuda().float()
                 if has_pre_quant:
-                    fc1 = fc1 * pre_quant_scale_2[expert]
+                    fc1 = fc1 * pre_quant_scale_2.squeeze()
                 if has_alpha:
                     fc1[fc1 > 448.0] = 448.0
                     fc1 = fc1.to(torch.float8_e4m3fn).float()