[None][feat] Add W4A16 per-channel weight-only support for Cutlass fused MoE #18313
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -211,15 +211,15 @@ class FusedMoeRunner : public torch::CustomClassHolder | |
| } | ||
|
|
||
| FusedMoeRunner(c10::ScalarType activation_dtype, c10::ScalarType weight_dtype, c10::ScalarType output_dtype, | ||
| bool use_deepseek_fp8_block_scale, bool use_w4_group_scaling, bool use_woq_per_channel, | ||
| bool use_mxfp8_act_scaling, bool use_mxfp8_weight_scaling, bool use_fused_finalize) | ||
| { | ||
| mActivationDtype = activation_dtype; | ||
| mWeightDtype = weight_dtype; | ||
| mOutputDtype = output_dtype; | ||
| mUseDeepSeekFP8BlockScaling = use_deepseek_fp8_block_scale; | ||
| mUseW4GroupScaling = use_w4_group_scaling; | ||
| mUseWoqPerChannel = use_woq_per_channel; | ||
| mUseMxfp8ActScaling = use_mxfp8_act_scaling; | ||
| mUseMxfp8WeightScaling = use_mxfp8_weight_scaling; | ||
| mUseFusedFinalize = use_fused_finalize; | ||
|
|
@@ -234,6 +234,14 @@ class FusedMoeRunner : public torch::CustomClassHolder | |
| && mWeightDtype == c10::ScalarType::Float8_e4m3fn), | ||
| "use_mxfp8_weight_scaling requires both activation and weight dtypes to be Float8_e4m3fn."); | ||
|
|
||
| // The per-channel weight-only path reinterprets fc2's dimensions as | ||
| // [num_experts, inter_size, hidden_size] and, for INT4, treats the trailing | ||
| // dim as packed two-per-byte. That is only meaningful for integer | ||
| // weight-only quantization, so reject other weight dtypes here rather | ||
| // than silently transposing them downstream. | ||
| TORCH_CHECK( | ||
| !mUseWoqPerChannel || isIntWeightOnlyQuant(), "use_woq_per_channel requires an INT8 or INT4 weight dtype."); | ||
|
coderabbitai[bot] marked this conversation as resolved.
|
||
|
|
||
| // keep consistent with cpp/tensorrt_llm/plugins/mixtureOfExperts/mixtureOfExpertsPlugin.cpp | ||
| if (mActivationDtype == c10::ScalarType::Half && mWeightDtype == c10::ScalarType::Half) | ||
| { | ||
|
|
@@ -471,21 +479,28 @@ class FusedMoeRunner : public torch::CustomClassHolder | |
| ActivationType base_activation_type = activation_type.has_value() | ||
| ? static_cast<ActivationType>(activation_type.value()) | ||
| : ActivationType::Swiglu; | ||
| if (mUseWoqPerChannel) | ||
| { | ||
| // Note: The weight shape for per-channel weight-only quantization is dim-swapped, e.g., | ||
| // fc2_expert_weights: [num_experts, inter_size, hidden_size] | ||
| // Mirror the non-woq else-branch below: gated activations (Swiglu/Geglu) require fc1's | ||
| // intermediate dim to be 2x fc2's (one half each for gate and up), while non-gated | ||
| // activations (Relu2/Identity/ReLU/SiLU/Gelu, e.g. Nemotron-H) require them to be equal. | ||
| // | ||
| // Under this dim-swapped layout the sub-byte packing sits on fc1's inter dim | ||
| // (sizes()[2]), while fc2's inter dim (sizes()[1]) is unpacked, so mInnerDimMultiplier | ||
| // multiplies the fc1 side here -- unlike the non-swapped branch below. For INT8 the | ||
| // multiplier is 1 and this is identical to the previous form; for INT4 it is 2 (see the | ||
| // isInt4Quant() branch in the constructor) and the previous form rejected every valid | ||
| // shape. | ||
| if (isGatedActivation(base_activation_type)) | ||
| { | ||
| TORCH_CHECK(fc1_expert_weights.sizes()[2] * mInnerDimMultiplier == fc2_expert_weights.sizes()[1] * 2, | ||
| "fc1_expert_weights inter size must be 2 times fc2_expert_weights inter size."); | ||
| } | ||
| else | ||
| { | ||
| TORCH_CHECK(fc1_expert_weights.sizes()[2] * mInnerDimMultiplier == fc2_expert_weights.sizes()[1], | ||
| "fc1_expert_weights inter size must be equal to fc2_expert_weights inter size."); | ||
| } | ||
| } | ||
|
|
@@ -506,16 +521,19 @@ class FusedMoeRunner : public torch::CustomClassHolder | |
| int experts_per_token = token_selected_experts.sizes()[1]; | ||
| int64_t num_rows = input.sizes()[0]; | ||
| int64_t hidden_size = fc2_expert_weights.sizes()[1]; | ||
| int64_t inter_size = fc2_expert_weights.sizes()[2] * mInnerDimMultiplier; | ||
| if (mUseWoqPerChannel) | ||
| { | ||
| // Note: The weight shape for per-channel weight-only quantization is different, e.g., fc2_expert_weights: | ||
| // [num_experts, inter_size, hidden_size] | ||
| hidden_size = fc2_expert_weights.sizes()[2] * mInnerDimMultiplier; | ||
| inter_size = fc2_expert_weights.sizes()[1]; | ||
|
coderabbitai[bot] marked this conversation as resolved.
|
||
| } | ||
| // Default the output width only after the per-channel weight-only layout is resolved: fc2's dims | ||
| // are transposed on that path (and packed two-per-byte for INT4), so hidden_size is not the | ||
| // logical width until the swap above. | ||
| int64_t unpadded_hidden_size_val | ||
| = unpadded_hidden_size.has_value() ? unpadded_hidden_size.value() : hidden_size; | ||
|
|
||
| if (isWMxfp4AMxfp8Quant() || isWMxfp4AFp8Quant()) | ||
| { | ||
|
|
@@ -789,16 +807,19 @@ class FusedMoeRunner : public torch::CustomClassHolder | |
| int experts_per_token = token_selected_experts.sizes()[1]; | ||
| int64_t num_rows = input.sizes()[0]; | ||
| int64_t hidden_size = fc2_expert_weights.sizes()[1]; | ||
| int64_t inter_size = fc2_expert_weights.sizes()[2] * mInnerDimMultiplier; | ||
| if (mUseWoqPerChannel) | ||
| { | ||
| // Note: The weight shape for per-channel weight-only quantization is different, e.g., fc2_expert_weights: | ||
| // [num_experts, inter_size, hidden_size] | ||
| hidden_size = fc2_expert_weights.sizes()[2] * mInnerDimMultiplier; | ||
| inter_size = fc2_expert_weights.sizes()[1]; | ||
| } | ||
| // Default the output width only after the per-channel weight-only layout is resolved: fc2's dims | ||
| // are transposed on that path (and packed two-per-byte for INT4), so hidden_size is not the | ||
| // logical width until the swap above. | ||
| int64_t unpadded_hidden_size_val | ||
| = unpadded_hidden_size.has_value() ? unpadded_hidden_size.value() : hidden_size; | ||
| int const num_experts_on_rank = fc2_expert_weights.sizes()[0]; | ||
| auto const num_experts_total = static_cast<int>(num_experts_on_rank * ep_size); | ||
| auto parallelism_config | ||
|
|
@@ -851,16 +872,16 @@ class FusedMoeRunner : public torch::CustomClassHolder | |
| // ordering differs from the non-woq path; both mirror the gated/non-gated split used in | ||
| // runMoe(). Gated activations (Swiglu/Geglu) require fc1's intermediate dim to be 2x fc2's; | ||
| // non-gated (Relu2/Identity/ReLU/SiLU/Gelu, e.g. Nemotron-H) require them to be equal. | ||
| if (mUseWoqPerChannel) | ||
| { | ||
| if (isGatedActivation(base_activation_type)) | ||
| { | ||
| TORCH_CHECK(fc1_expert_weights.sizes()[2] * mInnerDimMultiplier == fc2_expert_weights.sizes()[1] * 2, | ||
| "fc1_expert_weights inter size must be 2 times fc2_expert_weights inter size."); | ||
| } | ||
| else | ||
| { | ||
| TORCH_CHECK(fc1_expert_weights.sizes()[2] * mInnerDimMultiplier == fc2_expert_weights.sizes()[1], | ||
| "fc1_expert_weights inter size must be equal to fc2_expert_weights inter size."); | ||
| } | ||
| } | ||
|
|
@@ -978,15 +999,20 @@ class FusedMoeRunner : public torch::CustomClassHolder | |
| int64_t const num_rows = input.sizes()[0]; | ||
| int64_t hidden_size = fc2_expert_weights.sizes()[1]; | ||
| int64_t inter_size = fc2_expert_weights.sizes()[2] * mInnerDimMultiplier; | ||
| if (mUseWoqPerChannel) | ||
|
Collaborator
There was a problem hiding this comment.
Contributor
Author
There was a problem hiding this comment. Confirmed and fixed in de766ac. Until this change |
||
| { | ||
| // Note: The weight shape for per-channel weight-only quantization is different, e.g., fc2_expert_weights: | ||
| // [num_experts, inter_size, hidden_size] | ||
| hidden_size = fc2_expert_weights.sizes()[2] * mInnerDimMultiplier; | ||
| inter_size = fc2_expert_weights.sizes()[1]; | ||
| } | ||
| // Only group-scaled INT4 (W4A8_AWQ) carries a group size. Plain per-channel | ||
| // W4A16 must profile with -1 so the profiler builds QuantParams::Int, matching | ||
| // getQuantParams() and the runner chosen in the constructor; deriving this from | ||
| // isInt4Quant() alone would profile a groupwise configuration runMoe never uses. | ||
| int64_t const group_size_ = (isInt4Quant() && mUseW4GroupScaling) | ||
| ? TmaWarpSpecializedGroupedGemmInput::INT4GroupwiseParams::int4_group_size | ||
| : -1; | ||
| int64_t const group_size = isWFP4A16Quant() | ||
| ? TmaWarpSpecializedGroupedGemmInput::INT4GroupwiseParams::wfp4a16_group_size | ||
| : group_size_; | ||
|
|
@@ -1072,7 +1098,7 @@ class FusedMoeRunner : public torch::CustomClassHolder | |
|
|
||
| bool mUseDeepSeekFP8BlockScaling = false; | ||
| bool mUseW4GroupScaling = false; | ||
| bool mUseWoqPerChannel = false; | ||
| bool mUseMxfp8ActScaling = false; | ||
| bool mUseFusedFinalize = true; | ||
| bool mUseMxfp8WeightScaling = false; | ||
|
|
@@ -2348,10 +2374,10 @@ class FusedMoeRunner : public torch::CustomClassHolder | |
| else if (isIntWeightOnlyQuant()) | ||
| { | ||
| TORCH_CHECK(quant_scales.has_value(), "Expecting quant scales for weight only quantization"); | ||
| if (mUseWoqPerChannel) | ||
| { | ||
| TORCH_CHECK(quant_scales.value().size() == 2, | ||
| "Expecting 2 quant scales for per-channel weight only quantization"); | ||
| auto& fc1_weight_scales = quant_scales.value()[0]; | ||
| auto& fc2_weight_scales = quant_scales.value()[1]; | ||
| return kernels::QuantParams::Int(static_cast<float const*>(fc1_weight_scales.data_ptr()), | ||
|
|
||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Now that the flag is no longer INT8-specific by name, nothing ties
use_woq_per_channelto an integer weight dtype, yet it unconditionally swaps the hidden/inter interpretation inrunMoe,runMoeMinLantencyandrunGemmProfile. A caller passing it with, say, an FP8 weight dtype would get silently transposed dimensions with no error. Worth a constructor check next to theuse_mxfp8_weight_scalingone:TORCH_CHECK(!mUseWoqPerChannel || isIntWeightOnlyQuant(), ...).There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Agreed, added in de766ac next to the
use_mxfp8_weight_scalingcheck:TORCH_CHECK(!mUseWoqPerChannel || isIntWeightOnlyQuant(), "use_woq_per_channel requires an INT8 or INT4 weight dtype.");. After the rename the flag name was the only thing implying an integer weight dtype, and it swaps the hidden/inter interpretation in all three call sites.