A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 13, 19:30 UTC
Bugfix: Cutlass FP8 FusedMoE bad scaling factors (#27255)
Signed-off-by: Amir Klein <203507526+amirkl94@users.noreply.github.com> Co-authored-by: Tyler Michael Smith <tyler@neuralmagic.com> Co-authored-by: Michael Goin <mgoin64@gmail.com>
amirkl94 committed
Nov 5, 2025 at 13:06 UTC
6b7a81185d7e0f0fbad23d59db857eeafee688f7
4 files changed
+40
-14
tests/kernels/moe/test_flashinfer.py
+18
-6
@@ -6,7 +6,10 @@ import pytest
6
import torch
7
8
from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config
9
-from vllm.model_executor.layers.fused_moe.config import fp8_w8a8_moe_quant_config
9
+from vllm.model_executor.layers.fused_moe.config import (
10
+ FusedMoEQuantConfig,
11
+ fp8_w8a8_moe_quant_config,
12
+)
13
from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts
14
from vllm.model_executor.layers.fused_moe.layer import FusedMoE
15
from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
@@ -22,10 +25,10 @@ from vllm.platforms import current_platform
25
from vllm.utils.flashinfer import has_flashinfer_cutlass_fused_moe
26
27
if not has_flashinfer_cutlass_fused_moe() or not current_platform.has_device_capability(
25
- 100
28
+ 90
29
):
30
pytest.skip(
28
- "Requires flashinfer_cutlass_fused_moe and nvfp4 support",
31
+ "Supported for sm >= 90",
32
allow_module_level=True,
33
)
34
@@ -131,6 +134,8 @@ def test_flashinfer_per_tensor_moe_fp8_no_graph(
134
topk: int,
135
monkeypatch,
136
):
137
+ if not current_platform.has_device_capability(100):
138
+ pytest.skip("Test is only supported for sm >= 100")
139
current_platform.seed_everything(7)
140
monkeypatch.setenv("VLLM_FUSED_MOE_CHUNK_SIZE", "8192")
141
with set_current_vllm_config(vllm_config):
@@ -184,9 +189,6 @@ def test_flashinfer_per_tensor_moe_fp8_no_graph(
189
torch.testing.assert_close(output, flashinfer_output, atol=5.5e-2, rtol=1e-2)
190
191
187
-@pytest.mark.skip(
188
- "Requires flashinfer version that contains https://github.com/flashinfer-ai/flashinfer/pull/1472"
189
-)
192
@pytest.mark.parametrize("m,n,k", MNK_FACTORS)
193
@pytest.mark.parametrize("e", NUM_EXPERTS)
194
@pytest.mark.parametrize("topk", TOP_KS)
@@ -216,9 +218,13 @@ def test_flashinfer_cutlass_moe_fp8_no_graph(
218
219
quant_config = fp8_w8a8_moe_quant_config(
220
w1_scale=td.w13_weight_scale,
221
+ g1_alphas=(td.w13_weight_scale * td.a1_scale).squeeze(),
222
w2_scale=td.w2_weight_scale,
223
+ g2_alphas=(td.w2_weight_scale * td.a2_scale).squeeze(),
224
a1_scale=td.a1_scale,
225
+ a1_gscale=td.a1_scale,
226
a2_scale=td.a2_scale,
227
+ a2_gscale=1.0 / td.a2_scale,
228
per_act_token_quant=False,
229
)
230
@@ -238,6 +244,12 @@ def test_flashinfer_cutlass_moe_fp8_no_graph(
244
245
td.layer.dp_size = 1
246
247
+ def get_fused_moe_quant_config(n: torch.nn.Module) -> FusedMoEQuantConfig:
248
+ return quant_config
249
+
250
+ td.layer.get_fused_moe_quant_config = get_fused_moe_quant_config
251
+ td.layer.quant_method = td.layer
252
+
253
flashinfer_cutlass_output = flashinfer_cutlass_moe_fp8(
254
td.hidden_states,
255
td.layer,
vllm/model_executor/layers/fused_moe/config.py
+8
@@ -463,6 +463,10 @@ def fp8_w8a8_moe_quant_config(
463
per_act_token_quant: bool = False,
464
per_out_ch_quant: bool = False,
465
block_shape: list[int] | None = None,
466
+ a1_gscale: torch.Tensor | None = None,
467
+ a2_gscale: torch.Tensor | None = None,
468
+ g1_alphas: torch.Tensor | None = None,
469
+ g2_alphas: torch.Tensor | None = None,
470
) -> FusedMoEQuantConfig:
471
"""
472
Construct a quant config for fp8 activations and fp8 weights.
@@ -470,9 +474,13 @@ def fp8_w8a8_moe_quant_config(
474
return FusedMoEQuantConfig.make(
475
torch.float8_e4m3fn,
476
w1_scale=w1_scale,
477
+ g1_alphas=g1_alphas,
478
w2_scale=w2_scale,
479
+ g2_alphas=g2_alphas,
480
a1_scale=a1_scale,
481
+ a1_gscale=a1_gscale,
482
a2_scale=a2_scale,
483
+ a2_gscale=a2_gscale,
484
per_act_token_quant=per_act_token_quant,
485
per_out_ch_quant=per_out_ch_quant,
486
block_shape=block_shape,
vllm/model_executor/layers/fused_moe/flashinfer_cutlass_prepare_finalize.py
+8
-6
@@ -170,7 +170,7 @@ class FlashInferAllGatherMoEPrepareAndFinalize(FlashInferCutlassMoEPrepareAndFin
170
self._apply_router_weight_on_input(
171
a1, topk_weights, topk_ids, apply_router_weight_on_input
172
)
173
- if not self.use_dp:
173
+ if not self.use_dp and quant_config.quant_dtype == "nvfp4":
174
return a1, None, None, topk_ids, topk_weights
175
176
a1q, a1q_scale = moe_kernel_quantize_input(
@@ -181,11 +181,13 @@ class FlashInferAllGatherMoEPrepareAndFinalize(FlashInferCutlassMoEPrepareAndFin
181
quant_config.block_shape,
182
is_fp4_scale_swizzled=not self.use_dp,
183
)
184
- topk_weights, topk_ids, a1q, a1q_scale = get_dp_group().all_gatherv(
185
- [topk_weights, topk_ids, a1q, a1q_scale],
186
- dim=0,
187
- sizes=get_local_sizes(),
188
- )
184
+
185
+ if self.use_dp:
186
+ topk_weights, topk_ids, a1q, a1q_scale = get_dp_group().all_gatherv(
187
+ [topk_weights, topk_ids, a1q, a1q_scale],
188
+ dim=0,
189
+ sizes=get_local_sizes(),
190
+ )
191
if quant_config.quant_dtype == "nvfp4":
192
a1q_scale = nvfp4_block_scale_interleave(a1q_scale)
193
vllm/model_executor/layers/quantization/modelopt.py
+6
-2
@@ -567,9 +567,13 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):
567
568
return fp8_w8a8_moe_quant_config(
569
w1_scale=layer.w13_weight_scale,
570
+ g1_alphas=(layer.w13_weight_scale * layer.w13_input_scale).squeeze(),
571
w2_scale=layer.w2_weight_scale,
572
+ g2_alphas=(layer.w2_weight_scale * layer.w2_input_scale).squeeze(),
573
a1_scale=layer.w13_input_scale,
574
+ a1_gscale=layer.w13_input_scale,
575
a2_scale=layer.w2_input_scale,
576
+ a2_gscale=1.0 / layer.w2_input_scale,
577
per_act_token_quant=False,
578
)
579
@@ -1138,8 +1142,8 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
1142
moe: FusedMoEConfig,
1143
layer: torch.nn.Module,
1144
) -> None:
1141
- from vllm.model_executor.layers.quantization.utils.nvfp4_moe_support import ( # noqa: E501
1142
- detect_nvfp4_moe_support,
1145
+ from vllm.model_executor.layers.quantization.utils.nvfp4_moe_support import (
1146
+ detect_nvfp4_moe_support, # noqa: E501
1147
)
1148
1149
super().__init__(moe)