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 6, 10:45 UTC
[Bug] Fix fp8 trtllm MoE modular kernel supported routing methods (#37346)
Signed-off-by: wzhao18 <wzhao18.sz@gmail.com>
Wei Zhao committed
Mar 19, 2026 at 11:43 UTC
e27b8ba3d17df1330c81adf755988e8ee0fd6ab8
1 file changed
+16
-59
vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py
+16
-59
@@ -69,54 +69,11 @@ class TrtLlmFp8ExpertsBase:
69
"""Does not support non-gated MoE (i.e. Nanotron-3-Nano)."""
70
return True
71
72
- @staticmethod
73
- def _supports_quant_scheme(
74
- weight_key: QuantKey | None,
75
- activation_key: QuantKey | None,
76
- ) -> bool:
77
- """Supports Fp8 per-tensor, Fp8 block, and MXFP8."""
78
- SUPPORTED_W_A = [
79
- (kFp8Static128BlockSym, kFp8Dynamic128Sym),
80
- (kFp8StaticTensorSym, kFp8StaticTensorSym),
81
- (kMxfp8Static, kMxfp8Dynamic),
82
- ]
83
- return (weight_key, activation_key) in SUPPORTED_W_A
84
-
72
@staticmethod
73
def _supports_activation(activation: MoEActivation) -> bool:
74
"""Supports only SiLU and RELU^2 non-gated activation."""
75
return activation in [MoEActivation.SILU, MoEActivation.RELU2_NO_MUL]
76
90
- @staticmethod
91
- def _supports_routing_method(
92
- routing_method: RoutingMethodType,
93
- weight_key: QuantKey | None,
94
- activation_key: QuantKey | None,
95
- ) -> bool:
96
- """Monolithic kernels need to express router support."""
97
- # NOTE(dbari): TopK routing could also be enabled, but need to validate models
98
- # NOTE(dbari): Default is not implemented and should not be enabled until it is
99
- if (weight_key, activation_key) in [
100
- (kFp8Static128BlockSym, kFp8Dynamic128Sym),
101
- (kMxfp8Static, kMxfp8Dynamic),
102
- ]:
103
- # NOTE(rob): potentially allow others here. This is a conservative list.
104
- return routing_method in [
105
- RoutingMethodType.DeepSeekV3,
106
- RoutingMethodType.Renormalize,
107
- RoutingMethodType.RenormalizeNaive,
108
- ]
109
- elif (weight_key, activation_key) == (kFp8StaticTensorSym, kFp8StaticTensorSym):
110
- # NOTE(dbari): as above, potentially allow others here.
111
- return routing_method in [
112
- RoutingMethodType.DeepSeekV3,
113
- RoutingMethodType.Llama4,
114
- RoutingMethodType.Renormalize,
115
- RoutingMethodType.RenormalizeNaive,
116
- ]
117
- else:
118
- raise ValueError("Unsupported quantization scheme.")
119
-
77
@staticmethod
78
def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool:
79
"""Monolithic kernel so only use with naive DP/EP and TP."""
@@ -125,22 +82,6 @@ class TrtLlmFp8ExpertsBase:
82
or moe_parallel_config.use_naive_all2all_kernels
83
) and not moe_parallel_config.enable_eplb
84
128
- @staticmethod
129
- def _supports_router_logits_dtype(
130
- router_logits_dtype: torch.dtype | None,
131
- routing_method: RoutingMethodType,
132
- ) -> bool:
133
- """
134
- The FlashInfer TRTLLM FP8 kernel expects bfloat16 router_logits by default.
135
- Only DeepSeekV3 routing supports float32 router_logits (which is converted
136
- internally in the kernel).
137
- """
138
- if router_logits_dtype == torch.float32:
139
- # Only DeepSeekV3 routing handles float32 logits
140
- # https://github.com/flashinfer-ai/flashinfer/issues/2469
141
- return routing_method == RoutingMethodType.DeepSeekV3
142
- return True
143
-
85
def supports_chunking(self) -> bool:
86
return False
87
@@ -306,6 +247,22 @@ class TrtLlmFp8ExpertsMonolithic(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsMonolit
247
]
248
return (weight_key, activation_key) in SUPPORTED_W_A
249
250
+ @staticmethod
251
+ def _supports_router_logits_dtype(
252
+ router_logits_dtype: torch.dtype | None,
253
+ routing_method: RoutingMethodType,
254
+ ) -> bool:
255
+ """
256
+ The FlashInfer TRTLLM FP8 kernel expects bfloat16 router_logits by default.
257
+ Only DeepSeekV3 routing supports float32 router_logits (which is converted
258
+ internally in the kernel).
259
+ """
260
+ if router_logits_dtype == torch.float32:
261
+ # Only DeepSeekV3 routing handles float32 logits
262
+ # https://github.com/flashinfer-ai/flashinfer/issues/2469
263
+ return routing_method == RoutingMethodType.DeepSeekV3
264
+ return True
265
+
266
@staticmethod
267
def _supports_routing_method(
268
routing_method: RoutingMethodType,