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 2, 02:15 UTC
[Bugfix] Update TrtLLM MoE routing methods (#44347)
Signed-off-by: wzhao18 <wzhao18.sz@gmail.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Wei Zhao committed
Jun 3, 2026 at 05:56 UTC
ace95c9cf8307301fbcdffacc4ab44e4d5e6829a
6 files changed
+24
-24
vllm/model_executor/layers/fused_moe/config.py
+17
-12
@@ -102,23 +102,26 @@ def _quant_flags_to_group_shape(
102
class RoutingMethodType(IntEnum):
103
# Default: Softmax -> TopK
104
Default = (0,)
105
- # Renormalize: TopK -> Softmax/Sigmoid
105
+ # Renormalize: TopK -> Softmax
106
Renormalize = (1,)
107
# DeepSeekV3: Sigmoid -> RoutingBiasAdd -> Top2 in group -> Top4 groups
108
# -> Top8 experts from the Top4 groups
109
DeepSeekV3 = (2,)
110
# Llama4: Top1 -> Sigmoid
111
Llama4 = (3,)
112
- # RenormalizeNaive: Softmax/Sigmoid -> TopK -> Renormalize
112
+ # RenormalizeNaive: Softmax -> TopK -> Renormalize
113
RenormalizeNaive = (4,)
114
# TopK: TopK (no softmax)
115
TopK = (5,)
116
# SigmoidRenorm: Sigmoid -> TopK -> Renormalize (divide by sum of top-K)
117
SigmoidRenorm = (6,)
118
# MiniMax2: Sigmoid + Bias -> TopK -> ScaledSumNormalize
119
+ # (routeScale=1.0, epsilon=1e-20)
120
MiniMax2 = (7,)
121
+ # Sigmoid: Sigmoid -> TopK (no renormalization)
122
+ Sigmoid = (8,)
123
# Unspecified
121
- Unspecified = (8,)
124
+ Unspecified = (9,)
125
# other routing types (not passed to FlashInfer kernels)
126
# Deepseek V4 -> sqrtsoftplus + Bias + Normalize
127
DeepseekV4 = (100,)
@@ -132,6 +135,7 @@ def get_routing_method_type(
135
renormalize: bool,
136
num_expert_group: int | None,
137
has_e_score_bias: bool,
138
+ routed_scaling_factor: float | None = 1.0,
139
) -> RoutingMethodType:
140
if scoring_func == "sqrtsoftplus":
141
# DeepSeek V4 uses sqrtsoftplus routing with optional routing bias
@@ -142,20 +146,21 @@ def get_routing_method_type(
146
return RoutingMethodType.Unspecified
147
148
if has_e_score_bias:
145
- if (num_expert_group or 0) > 0 and scoring_func == "sigmoid":
146
- return RoutingMethodType.DeepSeekV3
147
- elif scoring_func == "sigmoid":
148
- return RoutingMethodType.MiniMax2
149
+ if scoring_func == "sigmoid":
150
+ if not renormalize:
151
+ return RoutingMethodType.Unspecified
152
+ if (num_expert_group or 0) > 0:
153
+ return RoutingMethodType.DeepSeekV3
154
+ if routed_scaling_factor in (None, 1.0):
155
+ return RoutingMethodType.MiniMax2
156
+ return RoutingMethodType.Unspecified
157
else:
158
return RoutingMethodType.Unspecified
159
160
if scoring_func == "sigmoid":
153
- if top_k == 1:
154
- return RoutingMethodType.Llama4
155
- elif renormalize:
161
+ if renormalize:
162
return RoutingMethodType.SigmoidRenorm
157
- else:
158
- return RoutingMethodType.Unspecified
163
+ return RoutingMethodType.Sigmoid
164
165
if scoring_func == "softmax":
166
if renormalize:
vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py
+3
-6
@@ -257,7 +257,7 @@ class TrtLlmFp8ExpertsMonolithic(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsMonolit
257
router_logits_dtype: torch.dtype | None,
258
routing_method: RoutingMethodType,
259
) -> bool:
260
- return router_logits_dtype != torch.float32
260
+ return router_logits_dtype in [torch.bfloat16, torch.float32]
261
262
@staticmethod
263
def _supports_routing_method(
@@ -279,6 +279,7 @@ class TrtLlmFp8ExpertsMonolithic(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsMonolit
279
RoutingMethodType.Renormalize,
280
RoutingMethodType.RenormalizeNaive,
281
RoutingMethodType.SigmoidRenorm,
282
+ RoutingMethodType.Sigmoid,
283
RoutingMethodType.MiniMax2,
284
RoutingMethodType.Simulated,
285
]
@@ -290,6 +291,7 @@ class TrtLlmFp8ExpertsMonolithic(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsMonolit
291
RoutingMethodType.Renormalize,
292
RoutingMethodType.RenormalizeNaive,
293
RoutingMethodType.SigmoidRenorm,
294
+ RoutingMethodType.Sigmoid,
295
RoutingMethodType.MiniMax2,
296
RoutingMethodType.Simulated,
297
]
@@ -402,11 +404,6 @@ class TrtLlmFp8ExpertsMonolithic(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsMonolit
404
else:
405
assert not apply_router_weight_on_input
406
405
- # Currently FI requires bfloat16 routing bias.
406
- # https://github.com/flashinfer-ai/flashinfer/issues/2909
407
- if e_score_correction_bias is not None:
408
- e_score_correction_bias = e_score_correction_bias.to(torch.bfloat16)
409
-
407
out = flashinfer.fused_moe.trtllm_fp8_per_tensor_scale_moe(
408
routing_logits=router_logits,
409
routing_bias=e_score_correction_bias,
vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py
+1
-6
@@ -360,7 +360,7 @@ class TrtLlmNvFp4ExpertsMonolithic(
360
router_logits_dtype: torch.dtype | None,
361
routing_method: RoutingMethodType,
362
) -> bool:
363
- return router_logits_dtype != torch.float32
363
+ return router_logits_dtype in [torch.bfloat16, torch.float32]
364
365
def apply(
366
self,
@@ -393,11 +393,6 @@ class TrtLlmNvFp4ExpertsMonolithic(
393
and self.routing_method_type != RoutingMethodType.Llama4
394
)
395
396
- # Currently FI requires bfloat16 routing bias.
397
- # https://github.com/flashinfer-ai/flashinfer/issues/2909
398
- if e_score_correction_bias is not None:
399
- e_score_correction_bias = e_score_correction_bias.to(torch.bfloat16)
400
-
396
output1_scale_gate_scalar = self.quant_config.g1_alphas
397
398
# Invoke kernel.
vllm/model_executor/layers/fused_moe/router/fused_topk_bias_router.py
+1
@@ -326,6 +326,7 @@ class FusedTopKBiasRouter(BaseRouter):
326
renormalize=self.renormalize,
327
num_expert_group=None,
328
has_e_score_bias=True,
329
+ routed_scaling_factor=self.routed_scaling_factor,
330
)
331
332
def _compute_routing(
vllm/model_executor/layers/fused_moe/router/grouped_topk_router.py
+1
@@ -283,6 +283,7 @@ 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,
287
)
288
289
def _compute_routing(
vllm/model_executor/layers/fused_moe/router/zero_expert_router.py
+1
@@ -63,6 +63,7 @@ class ZeroExpertRouter(BaseRouter):
63
renormalize=self.renormalize,
64
num_expert_group=None,
65
has_e_score_bias=True,
66
+ routed_scaling_factor=self.routed_scaling_factor,
67
)
68
69
def _compute_routing(