[Perf][MoE] Use fused MiniMax2 routing with non-unit routed scaling (#58880)
Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>
Jiangyun Zhu committed
Sep 28, 2026 at 03:37 UTC
231fdb83cccf18b0c96ab874fe210820025b3857
8 files changed
+135
-21
tests/kernels/moe/test_flashinfer.py
+97
@@ -42,6 +42,9 @@ from vllm.model_executor.layers.fused_moe.oracle.fp8 import (
42
convert_to_fp8_moe_kernel_format,
43
make_fp8_moe_quant_config,
44
)
45
+from vllm.model_executor.layers.fused_moe.router.fused_topk_bias_router import (
46
+ fused_topk_bias,
47
+)
48
from vllm.model_executor.layers.fused_moe.router.fused_topk_router import fused_topk
49
from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
50
rotate_weights_for_fi_trtllm_fp8_per_tensor_moe,
@@ -887,3 +890,97 @@ def test_trtllm_fp8_block_moe_deferred_finalize(
890
),
891
router_weights=topk_weights,
892
)
893
+
894
+
895
+@pytest.mark.parametrize("m", [1, 16])
896
+@pytest.mark.skipif(
897
+ not current_platform.is_device_capability_family(100),
898
+ reason="Requires TRTLLM-Gen FP8 MoE (SM100)",
899
+)
900
+def test_trtllm_mxfp8_minimax2_routing_applies_routed_scale(m: int, workspace_init):
901
+ """FlashInfer's fused MiniMax2 routing must apply routed_scaling_factor,
902
+ matching vLLM's own sigmoid+bias top-k routing (MiniMax-M3 uses 2.0)."""
903
+ e, topk, n, k, scale = 32, 4, 1024, 1024, 2.0
904
+ block_shape = [1, 32]
905
+ set_random_seed(7)
906
+ with set_current_vllm_config(vllm_config):
907
+ w1, w2, w1_scale, w2_scale = convert_to_fp8_moe_kernel_format(
908
+ Fp8MoeBackend.FLASHINFER_TRTLLM,
909
+ SimpleNamespace(
910
+ weight_block_size=block_shape,
911
+ moe_config=SimpleNamespace(
912
+ is_act_and_mul=True, intermediate_size_per_partition=n
913
+ ),
914
+ activation=MoEActivation.SILU,
915
+ ),
916
+ *_make_mxfp8_moe_weights(e, n, k),
917
+ w13_input_scale=None,
918
+ w2_input_scale=None,
919
+ )
920
+ quant_config = make_fp8_moe_quant_config(
921
+ Fp8MoeBackend.FLASHINFER_TRTLLM,
922
+ w1_scale,
923
+ w2_scale,
924
+ a1_scale=None,
925
+ a2_scale=None,
926
+ block_shape=block_shape,
927
+ )
928
+ moe_config = FusedMoEConfig(
929
+ num_experts=e,
930
+ experts_per_token=topk,
931
+ hidden_dim=k,
932
+ intermediate_size=n,
933
+ num_local_experts=e,
934
+ num_logical_experts=e,
935
+ activation=MoEActivation.SILU,
936
+ device="cuda",
937
+ moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(),
938
+ in_dtype=torch.bfloat16,
939
+ routing_method=RoutingMethodType.MiniMax2,
940
+ max_num_tokens=next_power_of_2(m),
941
+ )
942
+
943
+ def make_kernel(experts_cls) -> mk.FusedMoEKernel:
944
+ return mk.FusedMoEKernel(
945
+ maybe_make_prepare_finalize(
946
+ moe=moe_config,
947
+ quant_config=quant_config,
948
+ allow_new_interface=True,
949
+ use_monolithic=experts_cls is TrtLlmFp8ExpertsMonolithic,
950
+ ),
951
+ experts_cls(moe_config=moe_config, quant_config=quant_config),
952
+ )
953
+
954
+ a = torch.randn((m, k), device="cuda", dtype=torch.bfloat16) / 10
955
+ logits = torch.randn((m, e), device="cuda", dtype=torch.float32)
956
+ bias = torch.randn(e, device="cuda", dtype=torch.float32) / 10
957
+ common = dict(
958
+ w1=w1,
959
+ w2=w2,
960
+ activation=MoEActivation.SILU,
961
+ global_num_experts=e,
962
+ expert_map=None,
963
+ apply_router_weight_on_input=False,
964
+ )
965
+
966
+ fused = make_kernel(TrtLlmFp8ExpertsMonolithic).apply_monolithic(
967
+ hidden_states=a,
968
+ router_logits=logits,
969
+ e_score_correction_bias=bias,
970
+ routed_scaling_factor=scale,
971
+ **common,
972
+ )
973
+ topk_weights, topk_ids = fused_topk_bias(
974
+ a,
975
+ logits,
976
+ scoring_func="sigmoid",
977
+ e_score_correction_bias=bias,
978
+ topk=topk,
979
+ renormalize=True,
980
+ routed_scaling_factor=scale,
981
+ )
982
+ reference = make_kernel(TrtLlmFp8ExpertsModular).apply(
983
+ hidden_states=a, topk_weights=topk_weights, topk_ids=topk_ids, **common
984
+ )
985
+
986
+ torch.testing.assert_close(fused, reference, atol=2e-2, rtol=2e-2)
tests/kernels/moe/test_routing.py
+32
@@ -167,6 +167,38 @@ def test_single_expert_group_with_non_unit_scale_uses_grouped_topk() -> None:
167
assert isinstance(router, GroupedTopKRouter)
168
169
170
+def test_sigmoid_bias_routing_with_routed_scale_is_minimax2() -> None:
171
+ """FlashInfer's MiniMax2 routing applies routed_scaling_factor, so a
172
+ non-unit scale (MiniMax-M3 uses 2.0) must not block fused routing."""
173
+ router = create_fused_moe_router(
174
+ top_k=4,
175
+ global_num_experts=128,
176
+ scoring_func="sigmoid",
177
+ renormalize=True,
178
+ routed_scaling_factor=2.0,
179
+ e_score_correction_bias=torch.empty(128),
180
+ )
181
+
182
+ assert isinstance(router, FusedTopKBiasRouter)
183
+ assert router.routing_method_type == RoutingMethodType.MiniMax2
184
+
185
+
186
+def test_zero_expert_routing_is_unspecified() -> None:
187
+ """Zero experts are resolved in the router, so kernels with built-in
188
+ routing must never be selected for them."""
189
+ router = create_fused_moe_router(
190
+ top_k=4,
191
+ global_num_experts=128,
192
+ scoring_func="sigmoid",
193
+ renormalize=True,
194
+ e_score_correction_bias=torch.empty(160),
195
+ zero_expert_type="identity",
196
+ num_logical_experts=160,
197
+ )
198
+
199
+ assert router.routing_method_type == RoutingMethodType.Unspecified
200
+
201
+
202
def setup_eplb_state(
203
enable_eplb: bool, global_num_experts: int
204
) -> EplbLayerState | None:
vllm/model_executor/layers/fused_moe/config.py
+2
-5
@@ -118,7 +118,7 @@ class RoutingMethodType(IntEnum):
118
# SigmoidRenorm: Sigmoid -> TopK -> Renormalize (divide by sum of top-K)
119
SigmoidRenorm = (6,)
120
# MiniMax2: Sigmoid + Bias -> TopK -> ScaledSumNormalize
121
- # (routeScale=1.0, epsilon=1e-20)
121
+ # (routeScale=routed_scaling_factor, epsilon=1e-20)
122
MiniMax2 = (7,)
123
# Sigmoid: Sigmoid -> TopK (no renormalization)
124
Sigmoid = (8,)
@@ -137,7 +137,6 @@ def get_routing_method_type(
137
renormalize: bool,
138
num_expert_group: int | None,
139
has_e_score_bias: bool,
140
- routed_scaling_factor: float | None = 1.0,
140
) -> RoutingMethodType:
141
if scoring_func == "sqrtsoftplus":
142
# DeepSeek V4 uses sqrtsoftplus routing with optional routing bias
@@ -153,9 +152,7 @@ def get_routing_method_type(
152
return RoutingMethodType.Unspecified
153
if (num_expert_group or 0) > 0:
154
return RoutingMethodType.DeepSeekV3
156
- if routed_scaling_factor in (None, 1.0):
157
- return RoutingMethodType.MiniMax2
158
- return RoutingMethodType.Unspecified
155
+ return RoutingMethodType.MiniMax2
156
else:
157
return RoutingMethodType.Unspecified
158
vllm/model_executor/layers/fused_moe/router/cpu_router.py
-1
@@ -198,7 +198,6 @@ class CPURouter(BaseRouter):
198
renormalize=self.renormalize,
199
num_expert_group=self.num_expert_group,
200
has_e_score_bias=self.e_score_correction_bias is not None,
201
- routed_scaling_factor=self.routed_scaling_factor,
201
)
202
203
def _compute_routing(
vllm/model_executor/layers/fused_moe/router/fused_topk_bias_router.py
-1
@@ -391,7 +391,6 @@ class FusedTopKBiasRouter(BaseRouter):
391
renormalize=self.renormalize,
392
num_expert_group=None,
393
has_e_score_bias=True,
394
- routed_scaling_factor=self.routed_scaling_factor,
394
)
395
396
def _compute_routing(
vllm/model_executor/layers/fused_moe/router/grouped_topk_router.py
-1
@@ -283,7 +283,6 @@ class GroupedTopKRouter(BaseRouter):
283
renormalize=self.renormalize,
284
num_expert_group=self.num_expert_group,
285
has_e_score_bias=self.e_score_correction_bias is not None,
286
- routed_scaling_factor=self.routed_scaling_factor,
286
)
287
288
def _compute_routing(
vllm/model_executor/layers/fused_moe/router/router_factory.py
-1
@@ -188,7 +188,6 @@ def create_fused_moe_router(
188
renormalize=renormalize,
189
num_expert_group=groups,
190
has_e_score_bias=e_score_correction_bias is not None,
191
- routed_scaling_factor=routed_scaling_factor,
191
)
192
193
routing_method_preserved = advertised_routing_method(
vllm/model_executor/layers/fused_moe/router/zero_expert_router.py
+4
-12
@@ -4,10 +4,7 @@
4
import torch
5
6
from vllm.distributed.eplb.eplb_state import EplbLayerState
7
-from vllm.model_executor.layers.fused_moe.config import (
8
- RoutingMethodType,
9
- get_routing_method_type,
10
-)
7
+from vllm.model_executor.layers.fused_moe.config import RoutingMethodType
8
from vllm.model_executor.layers.fused_moe.fused_moe import (
9
zero_experts_compute_triton,
10
)
@@ -53,14 +50,9 @@ class ZeroExpertRouter(BaseRouter):
50
51
@property
52
def routing_method_type(self) -> RoutingMethodType:
56
- return get_routing_method_type(
57
- scoring_func=self.scoring_func,
58
- top_k=self.top_k,
59
- renormalize=self.renormalize,
60
- num_expert_group=None,
61
- has_e_score_bias=True,
62
- routed_scaling_factor=self.routed_scaling_factor,
63
- )
53
+ # Zero experts are resolved by this router, which kernels with
54
+ # built-in routing (e.g. FlashInfer TRT-LLM) cannot reproduce.
55
+ return RoutingMethodType.Unspecified
56
57
def _compute_routing(
58
self,