diff --git a/cpp/tensorrt_llm/kernels/internal_cutlass_kernels/aarch64-linux-gnu/tensorrt_llm_internal_cutlass_kernels_static.tar.xz b/cpp/tensorrt_llm/kernels/internal_cutlass_kernels/aarch64-linux-gnu/tensorrt_llm_internal_cutlass_kernels_static.tar.xz index aa53b594f4aa..91687217b568 100644 --- a/cpp/tensorrt_llm/kernels/internal_cutlass_kernels/aarch64-linux-gnu/tensorrt_llm_internal_cutlass_kernels_static.tar.xz +++ b/cpp/tensorrt_llm/kernels/internal_cutlass_kernels/aarch64-linux-gnu/tensorrt_llm_internal_cutlass_kernels_static.tar.xz @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:090be002758e4fb864e16ec25c0db3f8eb562a0033e60a156bbbfd6bce67a5a1 -size 63577888 +oid sha256:e91036bf25c1fabca1d98dddb828db4bd7ddb7f14e166a9f5da1c71fca7a1cb6 +size 63543156 diff --git a/cpp/tensorrt_llm/kernels/internal_cutlass_kernels/aarch64-linux-gnu/version.txt b/cpp/tensorrt_llm/kernels/internal_cutlass_kernels/aarch64-linux-gnu/version.txt index ed1277dce44a..18c1f3898632 100644 --- a/cpp/tensorrt_llm/kernels/internal_cutlass_kernels/aarch64-linux-gnu/version.txt +++ b/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 diff --git a/cpp/tensorrt_llm/kernels/internal_cutlass_kernels/x86_64-linux-gnu/tensorrt_llm_internal_cutlass_kernels_static.tar.xz b/cpp/tensorrt_llm/kernels/internal_cutlass_kernels/x86_64-linux-gnu/tensorrt_llm_internal_cutlass_kernels_static.tar.xz index 6143671a5c82..5a8f56abf2ba 100644 --- a/cpp/tensorrt_llm/kernels/internal_cutlass_kernels/x86_64-linux-gnu/tensorrt_llm_internal_cutlass_kernels_static.tar.xz +++ b/cpp/tensorrt_llm/kernels/internal_cutlass_kernels/x86_64-linux-gnu/tensorrt_llm_internal_cutlass_kernels_static.tar.xz @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:16608c6e8ca9f1c3f26408c0ec97049ac05b1ee8b647480d969d488d9cef91f0 -size 63066816 +oid sha256:89df42f33d13660aa097f2ca1a23ceb17bc259d2030dc404262304f3d5ac3420 +size 63065300 diff --git a/cpp/tensorrt_llm/kernels/internal_cutlass_kernels/x86_64-linux-gnu/version.txt b/cpp/tensorrt_llm/kernels/internal_cutlass_kernels/x86_64-linux-gnu/version.txt index 3e93949bcee8..1e4fd5d947b9 100644 --- a/cpp/tensorrt_llm/kernels/internal_cutlass_kernels/x86_64-linux-gnu/version.txt +++ b/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 diff --git a/cpp/tensorrt_llm/kernels/preQuantScaleKernel.cu b/cpp/tensorrt_llm/kernels/preQuantScaleKernel.cu index e94fa9d362a8..fef285931a01 100644 --- a/cpp/tensorrt_llm/kernels/preQuantScaleKernel.cu +++ b/cpp/tensorrt_llm/kernels/preQuantScaleKernel.cu @@ -40,8 +40,8 @@ struct Vec2Type<__nv_bfloat16> }; // namespace template -__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]; @@ -49,6 +49,8 @@ __global__ void apply_per_channel_scale( 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(scale) = reinterpret_cast(per_channel_scale)[col_offset]; @@ -95,46 +97,46 @@ __global__ void apply_per_channel_scale( } template -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 - <<>>(smoothed_act, act, per_channel_scale, rows, cols); + <<>>(smoothed_act, act, per_channel_scale, rows, cols, num_valid_tokens_ptr); } template -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(rows) * static_cast(cols); if (elems < 2048 * 2048) { apply_per_channel_scale_kernel_launcher_( - 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_( - 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_( - 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_( - 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_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_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 -__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(act_vec) = reinterpret_cast(act + i * cols)[col_offset]; - if constexpr ((std::is_same_v -#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800) && defined(ENABLE_BF16)) - || std::is_same_v -#endif - ) &&(kElems % 2 == 0)) - { - using Vec2 = typename Vec2Type::type; -#pragma unroll - for (int j = 0; j < kElems; j += 2) - { - if constexpr (std::is_same_v) - { - *reinterpret_cast(act_vec + j) - = __hmul2(*reinterpret_cast(act_vec + j), __half2half2(scale)); - } - else - { - *reinterpret_cast(act_vec + j) - = __hmul2(*reinterpret_cast(act_vec + j), __bfloat162bfloat162(scale)); - } - } - } - else - { -#pragma unroll - for (int j = 0; j < kElems; ++j) - { - act_vec[j] = static_cast(static_cast(act_vec[j]) * static_cast(scale)); - } - } - if constexpr (std::is_same_v) - { - reinterpret_cast(smoothed_act + i * cols)[col_offset] - = *reinterpret_cast(act_vec); - } - else - { -#pragma unroll - for (int j = 0; j < kElems; ++j) - { - (smoothed_act + i * cols)[col_offset * kElems + j] = static_cast(act_vec[j]); - } - } - } -} - -template -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<<>>( - smoothed_act, act, per_expert_scale, permuted_token_selected_experts, num_valid_tokens_ptr, rows, cols); -} - -template -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_(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_(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_(smoothed_act, act, per_expert_scale, - permuted_token_selected_experts, num_valid_tokens_ptr, rows, cols, stream); - } - else - { - apply_per_expert_scale_kernel_launcher_(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_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 diff --git a/cpp/tensorrt_llm/kernels/preQuantScaleKernel.h b/cpp/tensorrt_llm/kernels/preQuantScaleKernel.h index 637983733133..11cf0f193a8e 100644 --- a/cpp/tensorrt_llm/kernels/preQuantScaleKernel.h +++ b/cpp/tensorrt_llm/kernels/preQuantScaleKernel.h @@ -36,13 +36,8 @@ namespace kernels { template -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 -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 diff --git a/cpp/tensorrt_llm/plugins/weightOnlyGroupwiseQuantMatmulPlugin/weightOnlyGroupwiseQuantMatmulPlugin.cpp b/cpp/tensorrt_llm/plugins/weightOnlyGroupwiseQuantMatmulPlugin/weightOnlyGroupwiseQuantMatmulPlugin.cpp index 5dc16460879e..c9b779f4f330 100644 --- a/cpp/tensorrt_llm/plugins/weightOnlyGroupwiseQuantMatmulPlugin/weightOnlyGroupwiseQuantMatmulPlugin.cpp +++ b/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( reinterpret_cast<__nv_fp8_e4m3*>(workspace), reinterpret_cast(inputs[0]), - reinterpret_cast(inputs[mPreQuantScaleInputIdx]), m, k, stream); + reinterpret_cast(inputs[mPreQuantScaleInputIdx]), m, k, nullptr, stream); } else { tensorrt_llm::kernels::apply_per_channel_scale_kernel_launcher( reinterpret_cast(workspace), reinterpret_cast(inputs[0]), - reinterpret_cast(inputs[mPreQuantScaleInputIdx]), m, k, stream); + reinterpret_cast(inputs[mPreQuantScaleInputIdx]), m, k, nullptr, stream); } } diff --git a/cpp/tests/unit_tests/kernels/weightOnly/weightOnlyKernelTest.cpp b/cpp/tests/unit_tests/kernels/weightOnly/weightOnlyKernelTest.cpp index 81e0fce059e4..3f22d594e2a9 100644 --- a/cpp/tests/unit_tests/kernels/weightOnly/weightOnlyKernelTest.cpp +++ b/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( reinterpret_cast(scaled_act), reinterpret_cast(params.act), - reinterpret_cast(params.act_scale), params.m, params.k, stream); + reinterpret_cast(params.act_scale), params.m, params.k, nullptr, stream); act = scaled_act; } if constexpr (QuantOp == cutlass::WeightOnlyQuantOp::PER_COLUMN_SCALE_ONLY) diff --git a/tensorrt_llm/_torch/modules/fused_moe.py b/tensorrt_llm/_torch/modules/fused_moe.py index 5b2b89d89d1b..a57e82ea65dc 100755 --- a/tensorrt_llm/_torch/modules/fused_moe.py +++ b/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"], diff --git a/tensorrt_llm/layers/moe.py b/tensorrt_llm/layers/moe.py index 6b55cb4643b2..46e48f9d96af 100755 --- a/tensorrt_llm/layers/moe.py +++ b/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: diff --git a/tests/unittest/trt/quantization/test_moe_weight_only_groupwise_quant_matmul.py b/tests/unittest/trt/quantization/test_moe_weight_only_groupwise_quant_matmul.py index 757f769d04ee..c48326e205d9 100644 --- a/tests/unittest/trt/quantization/test_moe_weight_only_groupwise_quant_matmul.py +++ b/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()