Skip to content

Bugfix: Cutlass FP8 FusedMoE bad scaling factors #27255

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 6 commits into from
Nov 5, 2025
Merged

Bugfix: Cutlass FP8 FusedMoE bad scaling factors #27255

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
5bcd35d
Bugfix: Cutlass FP8 FusedMoE
amirkl94 Oct 21, 2025
c45fbe4
Merge branch 'main' into bugfix/cutlass-fused-moe
tlrmchlsmth Oct 30, 2025
4925add
Merge branch 'main' into bugfix/cutlass-fused-moe
amirkl94 Nov 2, 2025
7277cd4
Fix breakage after merge
amirkl94 Nov 2, 2025
935fb9e
Merge branch 'main' into bugfix/cutlass-fused-moe
mgoin Nov 4, 2025
d7932ce
CR
amirkl94 Nov 4, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
24 changes: 18 additions & 6 deletions tests/kernels/moe/test_flashinfer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,10 @@
import torch

from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config
from vllm.model_executor.layers.fused_moe.config import fp8_w8a8_moe_quant_config
from vllm.model_executor.layers.fused_moe.config import (
FusedMoEQuantConfig,
fp8_w8a8_moe_quant_config,
)
from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts
from vllm.model_executor.layers.fused_moe.layer import FusedMoE
from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
Expand All @@ -22,10 +25,10 @@
from vllm.utils.flashinfer import has_flashinfer_cutlass_fused_moe

if not has_flashinfer_cutlass_fused_moe() or not current_platform.has_device_capability(
100
90
):
pytest.skip(
"Requires flashinfer_cutlass_fused_moe and nvfp4 support",
"Supported for sm >= 90",
allow_module_level=True,
)

Expand Down Expand Up @@ -131,6 +134,8 @@ def test_flashinfer_per_tensor_moe_fp8_no_graph(
topk: int,
monkeypatch,
):
if not current_platform.has_device_capability(100):
pytest.skip("Test is only supported for sm >= 100")
current_platform.seed_everything(7)
monkeypatch.setenv("VLLM_FUSED_MOE_CHUNK_SIZE", "8192")
with set_current_vllm_config(vllm_config):
Expand Down Expand Up @@ -184,9 +189,6 @@ def test_flashinfer_per_tensor_moe_fp8_no_graph(
torch.testing.assert_close(output, flashinfer_output, atol=5.5e-2, rtol=1e-2)


@pytest.mark.skip(
"Requires flashinfer version that contains https://github.com/flashinfer-ai/flashinfer/pull/1472"
)
@pytest.mark.parametrize("m,n,k", MNK_FACTORS)
@pytest.mark.parametrize("e", NUM_EXPERTS)
@pytest.mark.parametrize("topk", TOP_KS)
Expand Down Expand Up @@ -216,9 +218,13 @@ def test_flashinfer_cutlass_moe_fp8_no_graph(

quant_config = fp8_w8a8_moe_quant_config(
w1_scale=td.w13_weight_scale,
g1_alphas=(td.w13_weight_scale * td.a1_scale).squeeze(),
w2_scale=td.w2_weight_scale,
g2_alphas=(td.w2_weight_scale * td.a2_scale).squeeze(),
a1_scale=td.a1_scale,
a1_gscale=td.a1_scale,
a2_scale=td.a2_scale,
a2_gscale=1.0 / td.a2_scale,
Copy link
Contributor

@gemini-code-assist gemini-code-assist bot Oct 21, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

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.

Suggested change
a2_gscale=1.0 / td.a2_scale,
a2_gscale=1.0 / (td.a2_scale + 1e-6),

per_act_token_quant=False,
)

Expand All @@ -238,6 +244,12 @@ def test_flashinfer_cutlass_moe_fp8_no_graph(

td.layer.dp_size = 1

def get_fused_moe_quant_config(n: torch.nn.Module) -> FusedMoEQuantConfig:
return quant_config

td.layer.get_fused_moe_quant_config = get_fused_moe_quant_config
td.layer.quant_method = td.layer

flashinfer_cutlass_output = flashinfer_cutlass_moe_fp8(
td.hidden_states,
td.layer,
Expand Down
8 changes: 8 additions & 0 deletions vllm/model_executor/layers/fused_moe/config.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -463,16 +463,24 @@ def fp8_w8a8_moe_quant_config(
per_act_token_quant: bool = False,
per_out_ch_quant: bool = False,
block_shape: list[int] | None = None,
a1_gscale: torch.Tensor | None = None,
Copy link
Contributor

@wenscarl wenscarl Nov 4, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The additional 4 scales were not present before 27223 and no issue without them. Is it possible to deduce them from others?

Copy link
Contributor Author

@amirkl94 amirkl94 Nov 4, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

As far as I could tell, no

a2_gscale: torch.Tensor | None = None,
g1_alphas: torch.Tensor | None = None,
g2_alphas: torch.Tensor | None = None,
) -> FusedMoEQuantConfig:
"""
Construct a quant config for fp8 activations and fp8 weights.
"""
return FusedMoEQuantConfig.make(
torch.float8_e4m3fn,
w1_scale=w1_scale,
g1_alphas=g1_alphas,
w2_scale=w2_scale,
g2_alphas=g2_alphas,
a1_scale=a1_scale,
a1_gscale=a1_gscale,
a2_scale=a2_scale,
a2_gscale=a2_gscale,
per_act_token_quant=per_act_token_quant,
per_out_ch_quant=per_out_ch_quant,
block_shape=block_shape,
Expand Down
14 changes: 8 additions & 6 deletions vllm/model_executor/layers/fused_moe/flashinfer_cutlass_prepare_finalize.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -170,7 +170,7 @@ def prepare(
self._apply_router_weight_on_input(
a1, topk_weights, topk_ids, apply_router_weight_on_input
)
if not self.use_dp:
if not self.use_dp and quant_config.quant_dtype == "nvfp4":
return a1, None, None, topk_ids, topk_weights

a1q, a1q_scale = moe_kernel_quantize_input(
Expand All @@ -181,11 +181,13 @@ def prepare(
quant_config.block_shape,
is_fp4_scale_swizzled=not self.use_dp,
)
topk_weights, topk_ids, a1q, a1q_scale = get_dp_group().all_gatherv(
[topk_weights, topk_ids, a1q, a1q_scale],
dim=0,
sizes=get_local_sizes(),
)

if self.use_dp:
topk_weights, topk_ids, a1q, a1q_scale = get_dp_group().all_gatherv(
[topk_weights, topk_ids, a1q, a1q_scale],
dim=0,
sizes=get_local_sizes(),
)
if quant_config.quant_dtype == "nvfp4":
a1q_scale = nvfp4_block_scale_interleave(a1q_scale)

Expand Down
8 changes: 6 additions & 2 deletions vllm/model_executor/layers/quantization/modelopt.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -570,9 +570,13 @@ def get_fused_moe_quant_config(

return fp8_w8a8_moe_quant_config(
w1_scale=layer.w13_weight_scale,
g1_alphas=(layer.w13_weight_scale * layer.w13_input_scale).squeeze(),
Copy link
Contributor

@wenscarl wenscarl Nov 4, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The g1_alphas can be computed by w1_scale and a1_scale here, right? Same in test_flashinfer.py.

Copy link
Contributor Author

@amirkl94 amirkl94 Nov 5, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'd rather not move this calculation inside the function as I'm not sure if in the future other paths will require different factors. If it's important then sure I can move it to be calculated in fp8_w8a8_moe_quant_config .
@wenscarl wdyt?

Copy link
Contributor

@wenscarl wenscarl Nov 5, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Sound reasonable. Let's leave it as it is.

w2_scale=layer.w2_weight_scale,
g2_alphas=(layer.w2_weight_scale * layer.w2_input_scale).squeeze(),
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,
Copy link
Contributor

@gemini-code-assist gemini-code-assist bot Oct 21, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

critical

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.

Suggested change
a2_gscale=1.0 / layer.w2_input_scale,
a2_gscale=1.0 / (layer.w2_input_scale + 1e-6),

per_act_token_quant=False,
)

Expand Down Expand Up @@ -1159,8 +1163,8 @@ def __init__(
moe: FusedMoEConfig,
layer: torch.nn.Module,
) -> None:
from vllm.model_executor.layers.quantization.utils.nvfp4_moe_support import ( # noqa: E501
detect_nvfp4_moe_support,
from vllm.model_executor.layers.quantization.utils.nvfp4_moe_support import (
detect_nvfp4_moe_support, # noqa: E501
)

super().__init__(moe)
Expand Down