Bugfix: Cutlass FP8 FusedMoE #27255
Bugfix: Cutlass FP8 FusedMoE #27255
Conversation
When running cutlass FusedMoE FP8 the scaling factors that are passed are None. This PR passes the correct scaling factors and enables the relevant test. Signed-off-by: Amir Klein <[email protected]>
amirkl94
requested review from
WoosukKwon,
mgoin,
robertgshaw2-redhat,
tlrmchlsmth and
yewentao256
as code owners
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request correctly addresses a bug in the Cutlass FP8 FusedMoE implementation by passing the necessary scaling factors. The changes are logical and enabling the previously skipped test test_flashinfer_cutlass_moe_fp8_no_graph validates the fix. However, I've identified a critical risk of a division-by-zero error in the calculation of the a2_gscale factor, which should be addressed to ensure numerical stability.
| a1_scale=layer.w13_input_scale, | ||
| a1_gscale=layer.w13_input_scale, | ||
| a2_scale=layer.w2_input_scale, | ||
| a2_gscale=1.0 / layer.w2_input_scale, |
There was a problem hiding this comment.
The calculation 1.0 / layer.w2_input_scale introduces a risk of a division-by-zero error if layer.w2_input_scale is zero. Although scales are typically positive, adding a small epsilon to the denominator is a crucial safeguard for numerical stability.
| a2_gscale=1.0 / layer.w2_input_scale, | |
| a2_gscale=1.0 / (layer.w2_input_scale + 1e-6), |
| a1_scale=td.a1_scale, | ||
| a1_gscale=td.a1_scale, | ||
| a2_scale=td.a2_scale, | ||
| a2_gscale=1.0 / td.a2_scale, |
There was a problem hiding this comment.
To prevent potential division-by-zero errors and for consistency with the recommended fix in the main logic, it's safer to add a small epsilon to the denominator here. While td.a2_scale is currently 1.0 in this test, this change improves the robustness of the test suite against future modifications.
| a2_gscale=1.0 / td.a2_scale, | |
| a2_gscale=1.0 / (td.a2_scale + 1e-6), |
1 participant
Purpose
When running cutlass FusedMoE FP8 the scaling factors that are passed are None. This PR passes the correct scaling factors and enables the relevant test.
Test Plan
Enabled previously disabled
test_flashinfer_cutlass_moe_fp8_no_graph.