Repository navigation
[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 |
|---|---|---|
|
|
@@ -333,15 +333,15 @@ class FusedMoeRunner : public EagerMoeWorkspaceOwner | |
| } | ||
|
|
||
| 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; | ||
|
|
@@ -360,6 +360,14 @@ class FusedMoeRunner : public EagerMoeWorkspaceOwner | |
| && 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.
Resolved
|
||
|
|
||
| // keep consistent with cpp/tensorrt_llm/plugins/mixtureOfExperts/mixtureOfExpertsPlugin.cpp | ||
| if (mActivationDtype == c10::ScalarType::Half && mWeightDtype == c10::ScalarType::Half) | ||
| { | ||
|
|
@@ -636,9 +644,9 @@ class FusedMoeRunner : public EagerMoeWorkspaceOwner | |
| 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.
Resolved
|
||
|
|
@@ -714,20 +722,20 @@ class FusedMoeRunner : public EagerMoeWorkspaceOwner | |
|
|
||
| // Validate the fc1/fc2 inter-size relationship now that the activation type (gated vs | ||
| // non-gated) is finalized: the swiglu_alpha/beta/limit handling above can promote the | ||
| // activation to SwigluBias, which is gated. Woq uses a transposed weight layout, so | ||
| // its fc1/fc2 dim ordering differs from the non-woq path. Gated activations (Swiglu/Geglu) | ||
| // require fc1's intermediate dim to be 2x fc2's (one half each for gate and up); non-gated | ||
| // activations (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."); | ||
| } | ||
| } | ||
|
|
@@ -985,9 +993,9 @@ class FusedMoeRunner : public EagerMoeWorkspaceOwner | |
| 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]; | ||
|
|
@@ -1043,20 +1051,20 @@ class FusedMoeRunner : public EagerMoeWorkspaceOwner | |
| swiglu_clamp_after_silu, mActivationParamBuffers, capture_buffers); | ||
|
|
||
| // Validate the fc1/fc2 inter-size relationship now that the activation type (gated vs | ||
| // non-gated) is finalized. Woq uses a transposed weight layout, so its fc1/fc2 dim | ||
| // 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."); | ||
| } | ||
| } | ||
|
|
@@ -1173,9 +1181,9 @@ class FusedMoeRunner : public EagerMoeWorkspaceOwner | |
| 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]; | ||
|
|
@@ -1185,8 +1193,13 @@ class FusedMoeRunner : public EagerMoeWorkspaceOwner | |
| std::nullopt, enable_alltoall, | ||
| /*has_input_sf=*/false, /*swizzled_input_sf=*/true, /*use_dynamic_fc2_scale=*/false); | ||
| #endif | ||
| // 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_; | ||
|
|
@@ -1273,7 +1286,7 @@ class FusedMoeRunner : public EagerMoeWorkspaceOwner | |
|
|
||
| bool mUseDeepSeekFP8BlockScaling = false; | ||
| bool mUseW4GroupScaling = false; | ||
| bool mUseWoqPerChannel = false; | ||
| bool mUseMxfp8ActScaling = false; | ||
| bool mUseFusedFinalize = true; | ||
| bool mUseMxfp8WeightScaling = false; | ||
|
|
@@ -2579,10 +2592,10 @@ class FusedMoeRunner : public EagerMoeWorkspaceOwner | |
| 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
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.