[Bugfix][Model] Fix MiniMax-M3 NVFP4 inference correctness (#48929)
Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com> Co-authored-by: Pavani Majety <pmajety@nvidia.com>
Gabriel Wu committed
Aug 5, 2026 at 22:25 UTC
08b8613b7b15643e9675595997968b307724596b
2 files changed
+100
-31
tests/kernels/moe/test_flashinfer_moe.py
+69
@@ -20,6 +20,7 @@ from vllm.model_executor.layers.fused_moe.all2all_utils import (
20
from vllm.model_executor.layers.fused_moe.config import (
21
FusedMoEConfig,
22
FusedMoEParallelConfig,
23
+ FusedMoEQuantConfig,
24
RoutingMethodType,
25
)
26
from vllm.model_executor.layers.fused_moe.experts.flashinfer_cutlass_moe import (
@@ -52,6 +53,74 @@ MNK_FACTORS = [
53
]
54
55
56
+@pytest.mark.parametrize(
57
+ "activation",
58
+ [MoEActivation.SWIGLUOAI, MoEActivation.SWIGLUOAI_UNINTERLEAVE],
59
+)
60
+def test_flashinfer_swigluoai_params_are_forwarded(activation, monkeypatch):
61
+ from flashinfer.fused_moe.core import ActivationType
62
+
63
+ moe_config = FusedMoEConfig(
64
+ num_experts=2,
65
+ experts_per_token=1,
66
+ hidden_dim=128,
67
+ intermediate_size=128,
68
+ num_local_experts=2,
69
+ num_logical_experts=2,
70
+ activation=activation,
71
+ device="cuda",
72
+ moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(),
73
+ in_dtype=torch.bfloat16,
74
+ routing_method=RoutingMethodType.TopK,
75
+ )
76
+ quant_config = FusedMoEQuantConfig.make(
77
+ gemm1_alpha=1.702,
78
+ gemm1_beta=1.0,
79
+ gemm1_clamp_limit=7.0,
80
+ )
81
+ experts = FlashInferExperts(moe_config=moe_config, quant_config=quant_config)
82
+
83
+ call_args = {}
84
+
85
+ def fake_flashinfer_cutlass_fused_moe(**kwargs):
86
+ call_args.update(kwargs)
87
+
88
+ monkeypatch.setattr(
89
+ "vllm.model_executor.layers.fused_moe.experts."
90
+ "flashinfer_cutlass_moe.flashinfer_cutlass_fused_moe",
91
+ fake_flashinfer_cutlass_fused_moe,
92
+ )
93
+ experts.apply(
94
+ output=torch.empty((1, 128), device="cuda", dtype=torch.bfloat16),
95
+ hidden_states=torch.empty((1, 128), device="cuda", dtype=torch.bfloat16),
96
+ w1=torch.empty((2, 256, 128), device="cuda", dtype=torch.bfloat16),
97
+ w2=torch.empty((2, 128, 128), device="cuda", dtype=torch.bfloat16),
98
+ topk_weights=torch.ones((1, 1), device="cuda", dtype=torch.float32),
99
+ topk_ids=torch.zeros((1, 1), device="cuda", dtype=torch.int64),
100
+ activation=activation,
101
+ global_num_experts=2,
102
+ expert_map=None,
103
+ a1q_scale=None,
104
+ a2_scale=None,
105
+ workspace13=None,
106
+ workspace2=None,
107
+ expert_tokens_meta=None,
108
+ apply_router_weight_on_input=False,
109
+ )
110
+
111
+ assert experts._supports_activation(activation)
112
+ assert call_args["activation_type"] == ActivationType.Swiglu
113
+ for name, value in (
114
+ ("swiglu_alpha", 1.702),
115
+ ("swiglu_beta", 1.0),
116
+ ("swiglu_limit", 7.0),
117
+ ):
118
+ torch.testing.assert_close(
119
+ call_args[name],
120
+ torch.full((2,), value, device="cuda", dtype=torch.float32),
121
+ )
122
+
123
+
124
@pytest.mark.parametrize("m,n,k", MNK_FACTORS)
125
@pytest.mark.parametrize("e", [40, 64, 256])
126
@pytest.mark.parametrize("topk", [1, 6, 8])
vllm/model_executor/layers/fused_moe/experts/flashinfer_cutlass_moe.py
+31
-31
@@ -13,6 +13,9 @@ from vllm.model_executor.layers.fused_moe.config import (
13
from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
14
TopKWeightAndReduceNoOP,
15
)
16
+from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
17
+ activation_to_flashinfer_type,
18
+)
19
from vllm.model_executor.layers.quantization.utils.quant_utils import (
20
QuantKey,
21
kFp8Dynamic128Sym,
@@ -93,29 +96,30 @@ class FlashInferExperts(mk.FusedMoEExpertsModular):
96
# - pass per-block weight scales to the kernel
97
# - skip input activation quantization (kernel applies scaling)
98
self.use_deepseek_fp8_block_scale = quant_config.is_block_quantized
96
- self.gemm1_clamp_limit: torch.Tensor | None = None
97
- if quant_config.gemm1_clamp_limit is not None:
98
- self.gemm1_clamp_limit = torch.tensor(
99
- [quant_config.gemm1_clamp_limit] * self.num_experts,
99
+
100
+ def _per_expert(value: float | None) -> torch.Tensor | None:
101
+ if value is None:
102
+ return None
103
+ return torch.full(
104
+ (self.num_experts,),
105
+ float(value),
106
dtype=torch.float32,
107
device=self.device,
108
)
109
110
+ self.gemm1_clamp_limit = _per_expert(quant_config.gemm1_clamp_limit)
111
+ self.gemm1_alpha = _per_expert(quant_config.gemm1_alpha)
112
+ self.gemm1_beta = _per_expert(quant_config.gemm1_beta)
113
+
114
if quant_config.weight_quant_dtype == "mxfp4":
115
# This value is used specifically for gpt-oss,
116
# Need to revisit this for other models
107
- self.gemm1_alpha = torch.tensor(
108
- [1.702] * self.num_experts, dtype=torch.float32, device=self.device
109
- )
110
- self.gemm1_beta = torch.tensor(
111
- [1.0] * self.num_experts, dtype=torch.float32, device=self.device
112
- )
117
+ if self.gemm1_alpha is None:
118
+ self.gemm1_alpha = _per_expert(1.702)
119
+ if self.gemm1_beta is None:
120
+ self.gemm1_beta = _per_expert(1.0)
121
if self.gemm1_clamp_limit is None:
114
- self.gemm1_clamp_limit = torch.tensor(
115
- [7.0] * self.num_experts,
116
- dtype=torch.float32,
117
- device=self.device,
118
- )
122
+ self.gemm1_clamp_limit = _per_expert(7.0)
123
if quant_config.quant_dtype == "mxfp8":
124
self.fake_input_scale = torch.ones(
125
self.num_experts,
@@ -190,6 +194,7 @@ class FlashInferExperts(mk.FusedMoEExpertsModular):
194
MoEActivation.GELU_TANH,
195
MoEActivation.RELU2_NO_MUL,
196
MoEActivation.SWIGLUOAI,
197
+ MoEActivation.SWIGLUOAI_UNINTERLEAVE,
198
]
199
200
@staticmethod
@@ -263,18 +268,6 @@ class FlashInferExperts(mk.FusedMoEExpertsModular):
268
expert_tokens_meta: mk.ExpertTokensMetadata | None,
269
apply_router_weight_on_input: bool | None,
270
):
266
- from flashinfer.fused_moe.core import ActivationType
267
-
268
- activation_str_to_value_map = {
269
- MoEActivation.SILU: ActivationType.Swiglu, # This is the default
270
- MoEActivation.GELU_TANH: ActivationType.Geglu,
271
- MoEActivation.SWIGLUOAI: ActivationType.Swiglu, # gpt-oss alias
272
- MoEActivation.RELU2_NO_MUL: ActivationType.Relu2,
273
- }
274
- assert activation in activation_str_to_value_map, (
275
- f"{activation=} missing from {activation_str_to_value_map.keys()=}"
276
- )
277
-
271
quant_scales = None
272
fc1_expert_weights = None
273
fc2_expert_weights = None
@@ -282,9 +275,16 @@ class FlashInferExperts(mk.FusedMoEExpertsModular):
275
fc2_expert_biases = None
276
swiglu_alpha = None
277
swiglu_beta = None
285
- swiglu_limit = (
286
- self.gemm1_clamp_limit if activation == MoEActivation.SILU else None
287
- )
278
+ swiglu_limit = None
279
+ if activation == MoEActivation.SILU:
280
+ swiglu_limit = self.gemm1_clamp_limit
281
+ elif activation in (
282
+ MoEActivation.SWIGLUOAI,
283
+ MoEActivation.SWIGLUOAI_UNINTERLEAVE,
284
+ ):
285
+ swiglu_alpha = self.gemm1_alpha
286
+ swiglu_beta = self.gemm1_beta
287
+ swiglu_limit = self.gemm1_clamp_limit
288
use_mxfp8_act_scaling = False
289
use_w4_group_scaling = False
290
# Select quantization metadata based on FP8 format/path
@@ -392,7 +392,7 @@ class FlashInferExperts(mk.FusedMoEExpertsModular):
392
tp_rank=self.tp_rank,
393
ep_size=self.ep_size,
394
ep_rank=self.ep_rank,
395
- activation_type=activation_str_to_value_map[activation],
395
+ activation_type=activation_to_flashinfer_type(activation),
396
# Informs FlashInfer to use the block-scale decoding path when True
397
use_deepseek_fp8_block_scale=self.use_deepseek_fp8_block_scale,
398
use_mxfp8_act_scaling=use_mxfp8_act_scaling,