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 4, 15:00 UTC
[Perf] Update TRTLLM supported MoE routing methods (#39141)
Signed-off-by: wzhao18 <wzhao18.sz@gmail.com> Signed-off-by: Wei Zhao <51183510+wzhao18@users.noreply.github.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Co-authored-by: root <root@bia0030.bia.clusters.nvidia.com> Co-authored-by: root <root@bia0036.bia.clusters.nvidia.com>
Wei Zhao committed
Apr 27, 2026 at 14:16 UTC
c8bbe05189babd69312876c1dcdc80912207e154
5 files changed
+26
-92
vllm/model_executor/layers/fused_moe/config.py
+14
-7
@@ -113,14 +113,17 @@ class RoutingMethodType(IntEnum):
113
RenormalizeNaive = (4,)
114
# TopK: TopK (no softmax)
115
TopK = (5,)
116
- # Custom
117
- Custom = (6,)
118
- # Simulated
119
- Simulated = (7,)
120
- # Deepseek V4 -> sqrtsoftplus + Bias + Normalize
121
- DeepseekV4 = (8,)
116
+ # SigmoidRenorm: Sigmoid -> TopK -> Renormalize (divide by sum of top-K)
117
+ SigmoidRenorm = (6,)
118
+ # MiniMax2: Sigmoid + Bias -> TopK -> ScaledSumNormalize
119
+ MiniMax2 = (7,)
120
# Unspecified
123
- Unspecified = 9.0
121
+ Unspecified = (8,)
122
+ # other routing types (not passed to FlashInfer kernels)
123
+ # Deepseek V4 -> sqrtsoftplus + Bias + Normalize
124
+ DeepseekV4 = (100,)
125
+ Custom = (101,)
126
+ Simulated = (102,)
127
128
129
def get_routing_method_type(
@@ -141,12 +144,16 @@ def get_routing_method_type(
144
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
else:
150
return RoutingMethodType.Unspecified
151
152
if scoring_func == "sigmoid":
153
if top_k == 1:
154
return RoutingMethodType.Llama4
155
+ elif renormalize:
156
+ return RoutingMethodType.SigmoidRenorm
157
else:
158
return RoutingMethodType.Unspecified
159
vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py
+9
-43
@@ -175,13 +175,6 @@ class TrtLlmFp8ExpertsModular(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsModular):
175
# Pack topk ids and weights into format expected by the kernel.
176
packed_topk_ids = trtllm_moe_pack_topk_ids_weights(topk_ids, topk_weights)
177
178
- # trtllm_fp8_block_scale_routed_moe does not support autotuning
179
- # so skip this kernel during dummy run for autotuning.
180
- import vllm.utils.flashinfer as fi_utils
181
-
182
- if fi_utils._is_fi_autotuning:
183
- return
184
-
178
assert a1q_scale is not None
179
180
is_mxfp8 = self.quant_config.block_shape == [1, 32]
@@ -196,11 +189,7 @@ class TrtLlmFp8ExpertsModular(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsModular):
189
weight_layout = WeightLayout.BlockMajorK
190
hidden_states_scale = a1q_scale.t().contiguous()
191
199
- # `trtllm_fp8_block_scale_routed_moe` has a bug and does not write to the
200
- # output tensor in-place so we need to manually copy the result to the
201
- # output tensor
202
- # https://github.com/flashinfer-ai/flashinfer/issues/2703
203
- result = flashinfer.fused_moe.trtllm_fp8_block_scale_routed_moe(
192
+ flashinfer.fused_moe.trtllm_fp8_block_scale_routed_moe(
193
topk_ids=packed_topk_ids,
194
routing_bias=None,
195
hidden_states=hidden_states,
@@ -217,13 +206,12 @@ class TrtLlmFp8ExpertsModular(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsModular):
206
local_expert_offset=self.ep_rank * self.local_num_experts,
207
local_num_experts=self.local_num_experts,
208
routed_scaling_factor=None,
220
- routing_method_type=1,
209
+ routing_method_type=1, # not used
210
use_shuffled_weight=use_shuffled_weight,
211
weight_layout=weight_layout,
212
fp8_quantization_type=fp8_quant_type,
224
- # output=output,
213
+ output=output,
214
)
226
- output.copy_(result)
215
216
217
class TrtLlmFp8ExpertsMonolithic(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsMonolithic):
@@ -275,20 +263,6 @@ class TrtLlmFp8ExpertsMonolithic(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsMonolit
263
router_logits_dtype: torch.dtype | None,
264
routing_method: RoutingMethodType,
265
) -> bool:
278
- """
279
- The FlashInfer TRTLLM FP8 kernel expects bfloat16 router_logits by default.
280
- DeepSeekV3 routing supports float32 router_logits (converted internally).
281
- Simulated routing generates synthetic decisions and is agnostic to dtype.
282
- """
283
- if router_logits_dtype == torch.float32:
284
- # DeepSeekV3 routing handles float32 logits internally.
285
- # Simulated routing generates synthetic decisions, so the
286
- # kernel doesn't care about the actual logits dtype.
287
- # https://github.com/flashinfer-ai/flashinfer/issues/2469
288
- return routing_method in (
289
- RoutingMethodType.DeepSeekV3,
290
- RoutingMethodType.Simulated,
291
- )
266
return True
267
268
@staticmethod
@@ -308,18 +282,22 @@ class TrtLlmFp8ExpertsMonolithic(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsMonolit
282
# NOTE(rob): potentially allow others here. This is a conservative list.
283
return routing_method in [
284
RoutingMethodType.DeepSeekV3,
311
- RoutingMethodType.Simulated,
285
RoutingMethodType.Renormalize,
286
RoutingMethodType.RenormalizeNaive,
287
+ RoutingMethodType.SigmoidRenorm,
288
+ RoutingMethodType.MiniMax2,
289
+ RoutingMethodType.Simulated,
290
]
291
elif (weight_key, activation_key) == (kFp8StaticTensorSym, kFp8StaticTensorSym):
292
# NOTE(dbari): as above, potentially allow others here.
293
return routing_method in [
294
RoutingMethodType.DeepSeekV3,
295
RoutingMethodType.Llama4,
320
- RoutingMethodType.Simulated,
296
RoutingMethodType.Renormalize,
297
RoutingMethodType.RenormalizeNaive,
298
+ RoutingMethodType.SigmoidRenorm,
299
+ RoutingMethodType.MiniMax2,
300
+ RoutingMethodType.Simulated,
301
]
302
else:
303
raise ValueError("Unsupported quantization scheme.")
@@ -355,14 +333,6 @@ class TrtLlmFp8ExpertsMonolithic(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsMonolit
333
# TODO: fuse into the quant kernel.
334
assert a1q_scale is not None
335
358
- if self.routing_method_type == RoutingMethodType.DeepSeekV3:
359
- router_logits = router_logits.to(torch.float32)
360
-
361
- # Currently FI requires bfloat16 routing bias.
362
- # https://github.com/flashinfer-ai/flashinfer/issues/2909
363
- if e_score_correction_bias is not None:
364
- e_score_correction_bias = e_score_correction_bias.to(torch.bfloat16)
365
-
336
is_mxfp8 = self.quant_config.block_shape == [1, 32]
337
if is_mxfp8:
338
fp8_quant_type = Fp8QuantizationType.MxFp8
@@ -429,10 +399,6 @@ class TrtLlmFp8ExpertsMonolithic(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsMonolit
399
else:
400
assert not apply_router_weight_on_input
401
432
- # The DeepSeekV3 routing method requires float32 router logits.
433
- if self.routing_method_type == RoutingMethodType.DeepSeekV3:
434
- router_logits = router_logits.to(torch.float32)
435
-
402
# Currently FI requires bfloat16 routing bias.
403
# https://github.com/flashinfer-ai/flashinfer/issues/2909
404
if e_score_correction_bias is not None:
vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py
+3
-29
@@ -198,13 +198,6 @@ class TrtLlmNvFp4ExpertsModular(TrtLlmNvFp4ExpertsBase, mk.FusedMoEExpertsModula
198
# Pack topk ids and weights into format expected by the kernel.
199
packed_tensor = trtllm_moe_pack_topk_ids_weights(topk_ids, topk_weights)
200
201
- # trtllm_fp4_block_scale_routed_moe does not support autotuning
202
- # so skip this kernel during dummy run for autotuning.
203
- import vllm.utils.flashinfer as fi_utils
204
-
205
- if fi_utils._is_fi_autotuning:
206
- return
207
-
201
# Invoke kernel.
202
flashinfer.fused_moe.trtllm_fp4_block_scale_routed_moe(
203
topk_ids=packed_tensor,
@@ -233,7 +226,7 @@ class TrtLlmNvFp4ExpertsModular(TrtLlmNvFp4ExpertsBase, mk.FusedMoEExpertsModula
226
local_expert_offset=self.ep_rank * self.local_num_experts,
227
local_num_experts=self.local_num_experts,
228
routed_scaling_factor=None,
236
- routing_method_type=1,
229
+ routing_method_type=1, # not used
230
do_finalize=True,
231
activation_type=activation_to_flashinfer_int(activation),
232
output=output,
@@ -267,6 +260,8 @@ class TrtLlmNvFp4ExpertsMonolithic(
260
RoutingMethodType.Renormalize,
261
RoutingMethodType.RenormalizeNaive,
262
RoutingMethodType.Llama4,
263
+ RoutingMethodType.SigmoidRenorm,
264
+ RoutingMethodType.MiniMax2,
265
RoutingMethodType.Simulated,
266
]
267
@@ -275,20 +270,6 @@ class TrtLlmNvFp4ExpertsMonolithic(
270
router_logits_dtype: torch.dtype | None,
271
routing_method: RoutingMethodType,
272
) -> bool:
278
- """
279
- The FlashInfer TRTLLM NvFp4 kernel expects bfloat16 router_logits by default.
280
- DeepSeekV3 routing supports float32 router_logits (converted internally).
281
- Simulated routing generates synthetic decisions and is agnostic to dtype.
282
- """
283
- if router_logits_dtype == torch.float32:
284
- # DeepSeekV3 routing handles float32 logits internally.
285
- # Simulated routing generates synthetic decisions, so the
286
- # kernel doesn't care about the actual logits dtype.
287
- # https://github.com/flashinfer-ai/flashinfer/issues/2469
288
- return routing_method in (
289
- RoutingMethodType.DeepSeekV3,
290
- RoutingMethodType.Simulated,
291
- )
273
return True
274
275
def apply(
@@ -322,13 +303,6 @@ class TrtLlmNvFp4ExpertsMonolithic(
303
and self.routing_method_type != RoutingMethodType.Llama4
304
)
305
325
- # Prepare router logits for kernel format.
326
- router_logits = (
327
- router_logits.to(torch.float32)
328
- if self.routing_method_type == RoutingMethodType.DeepSeekV3
329
- else router_logits
330
- )
331
-
306
# Currently FI requires bfloat16 routing bias.
307
# https://github.com/flashinfer-ai/flashinfer/issues/2909
308
if e_score_correction_bias is not None:
vllm/model_executor/layers/fused_moe/oracle/nvfp4.py
-1
@@ -275,7 +275,6 @@ def select_nvfp4_moe_backend(
275
activation_key,
276
activation_format,
277
)
278
-
278
if supported:
279
logger.info_once(_make_log_backend(backend))
280
return backend, k_cls
vllm/model_executor/models/deepseek_v2.py
-12
@@ -50,7 +50,6 @@ from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
50
from vllm.model_executor.layers.fused_moe import (
51
FusedMoE,
52
GateLinear,
53
- RoutingMethodType,
53
fused_moe_make_expert_params_mapping,
54
)
55
from vllm.model_executor.layers.layernorm import LayerNorm, RMSNorm
@@ -338,17 +337,6 @@ class DeepseekV2MoE(nn.Module):
337
else None,
338
)
339
341
- # NOTE(rob): this is a hack until we finish off the PR for
342
- # merging TRTLLM kernels into the MK framework. Then we can
343
- # query the MonolithicMK for the expected router logits.
344
- # NOTE(dbari): Use BF16 if routing is not Deepseek, e.g. Mistral Large 3
345
- self.gate.set_out_dtype(
346
- torch.float32
347
- if self.experts.quant_method.is_monolithic
348
- and self.experts.routing_method_type == RoutingMethodType.DeepSeekV3
349
- else torch.bfloat16
350
- )
351
-
340
# Pre-cast the bias to match the gate output dtype so the
341
# conversion is not repeated on every forward pass. All
342
# downstream references (FusedMoE, router) share the same