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 9, 00:00 UTC
[Bugfix] Disable TRTLLM FP8 MoE if router_logits_dtype==float32 and routing_method!=DeepSeekV3 (#33613)
Signed-off-by: mgoin <mgoin64@gmail.com>
Michael Goin committed
Feb 3, 2026 at 16:26 UTC
2a99c5a6c86daef8c766ba2dbf05c385b192c64b
5 files changed
+43
-33
vllm/model_executor/layers/fused_moe/flashinfer_trtllm_moe.py
+31
-4
@@ -98,7 +98,23 @@ def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bo
98
return not moe_parallel_config.enable_eplb
99
100
101
-def is_supported_config_trtllm(
101
+def _supports_router_logits_dtype(
102
+ router_logits_dtype: torch.dtype | None,
103
+ routing_method: RoutingMethodType,
104
+) -> bool:
105
+ """
106
+ The FlashInfer TRTLLM FP8 kernel expects bfloat16 router_logits by default.
107
+ Only DeepSeekV3 routing supports float32 router_logits (which is converted
108
+ internally in the kernel).
109
+ """
110
+ if router_logits_dtype == torch.float32:
111
+ # Only DeepSeekV3 routing handles float32 logits
112
+ # https://github.com/flashinfer-ai/flashinfer/issues/2469
113
+ return routing_method == RoutingMethodType.DeepSeekV3
114
+ return True
115
+
116
+
117
+def is_supported_config_trtllm_fp8(
118
moe_config: FusedMoEConfig,
119
weight_key: QuantKey | None,
120
activation_key: QuantKey | None,
@@ -127,6 +143,10 @@ def is_supported_config_trtllm(
143
return False, _make_reason("routing method")
144
elif activation_format != mk.FusedMoEActivationFormat.Standard:
145
return False, _make_reason("activation format")
146
+ elif not _supports_router_logits_dtype(
147
+ moe_config.router_logits_dtype, moe_config.routing_method
148
+ ):
149
+ return False, _make_reason("float32 router_logits with non-DeepSeekV3 routing")
150
151
return True, None
152
@@ -161,7 +181,7 @@ def is_supported_config_trtllm_bf16(
181
182
def flashinfer_fused_moe_blockscale_fp8(
183
routing_logits: torch.Tensor,
164
- routing_bias: torch.Tensor,
184
+ routing_bias: torch.Tensor | None,
185
x: torch.Tensor,
186
w13_weight: torch.Tensor,
187
w13_weight_scale_inv: torch.Tensor,
@@ -175,7 +195,7 @@ def flashinfer_fused_moe_blockscale_fp8(
195
expert_offset: int,
196
local_num_experts: int,
197
block_shape: list[int],
178
- routing_method_type: int = int(RoutingMethodType.DeepSeekV3),
198
+ routing_method_type: int,
199
routed_scaling: float | None = 1.0,
200
) -> torch.Tensor:
201
from vllm.utils.flashinfer import flashinfer_trtllm_fp8_block_scale_moe
@@ -188,6 +208,13 @@ def flashinfer_fused_moe_blockscale_fp8(
208
# Routing kernel expects #experts <= #threads 512
209
assert global_num_experts <= 512
210
211
+ # The DeepSeekV3 routing method requires float32 router logits.
212
+ if routing_method_type == RoutingMethodType.DeepSeekV3:
213
+ routing_logits = routing_logits.to(torch.float32)
214
+
215
+ if routing_bias is not None:
216
+ routing_bias = routing_bias.to(x.dtype)
217
+
218
a_q, a_sf = per_token_group_quant_fp8(x, block_shape[1])
219
# NOTE: scales of hidden states have to be transposed!
220
a_sf_t = a_sf.t().contiguous()
@@ -215,7 +242,7 @@ def flashinfer_fused_moe_blockscale_fp8(
242
243
def flashinfer_fused_moe_blockscale_fp8_fake(
244
routing_logits: torch.Tensor,
218
- routing_bias: torch.Tensor,
245
+ routing_bias: torch.Tensor | None,
246
x: torch.Tensor,
247
w13_weight: torch.Tensor,
248
w13_weight_scale_inv: torch.Tensor,
vllm/model_executor/layers/fused_moe/oracle/fp8.py
+5
-5
@@ -18,7 +18,7 @@ from vllm.model_executor.layers.fused_moe.config import (
18
fp8_w8a16_moe_quant_config,
19
)
20
from vllm.model_executor.layers.fused_moe.flashinfer_trtllm_moe import (
21
- is_supported_config_trtllm,
21
+ is_supported_config_trtllm_fp8,
22
)
23
from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
24
FlashinferMoeBackend,
@@ -213,7 +213,7 @@ def select_fp8_moe_backend(
213
214
if fi_backend == FlashinferMoeBackend.TENSORRT_LLM:
215
backend = Fp8MoeBackend.FLASHINFER_TRTLLM
216
- supported, reason = is_supported_config_trtllm(
216
+ supported, reason = is_supported_config_trtllm_fp8(
217
config, weight_key, activation_key, activation_format
218
)
219
if supported:
@@ -240,7 +240,7 @@ def select_fp8_moe_backend(
240
]:
241
if backend == Fp8MoeBackend.FLASHINFER_TRTLLM:
242
k_cls = None
243
- supported, reason = is_supported_config_trtllm(
243
+ supported, reason = is_supported_config_trtllm_fp8(
244
config,
245
weight_key,
246
activation_key,
@@ -309,7 +309,7 @@ def select_fp8_moe_backend(
309
for backend in AVAILABLE_BACKENDS:
310
if backend == Fp8MoeBackend.FLASHINFER_TRTLLM:
311
k_cls = None
312
- supported, reason = is_supported_config_trtllm(
312
+ supported, reason = is_supported_config_trtllm_fp8(
313
config,
314
weight_key,
315
activation_key,
@@ -482,7 +482,7 @@ def make_fp8_moe_kernel(
482
)
483
assert prepare_finalize is not None
484
485
- logger.info_once("Using %s", prepare_finalize.__class__.__name__)
485
+ logger.info_once("Using %s", prepare_finalize.__class__.__name__, scope="local")
486
487
# Create Experts.
488
if prepare_finalize.activation_format == mk.FusedMoEActivationFormat.BatchedExperts:
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py
+3
-12
@@ -27,7 +27,6 @@ from vllm.model_executor.layers.fused_moe import (
27
from vllm.model_executor.layers.fused_moe.config import (
28
FusedMoEConfig,
29
FusedMoEQuantConfig,
30
- RoutingMethodType,
30
int4_w4a16_moe_quant_config,
31
int4_w4afp8_moe_quant_config,
32
int8_w8a8_moe_quant_config,
@@ -1027,17 +1026,9 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):
1026
if self.block_quant:
1027
import vllm.model_executor.layers.fused_moe.flashinfer_trtllm_moe # noqa: E501, F401
1028
1030
- e_score_correction_bias = (
1031
- layer.e_score_correction_bias.to(x.dtype)
1032
- if layer.e_score_correction_bias is not None
1033
- else None
1034
- )
1035
- routing_method_type = layer.routing_method_type
1029
return torch.ops.vllm.flashinfer_fused_moe_blockscale_fp8(
1037
- routing_logits=router_logits.to(torch.float32)
1038
- if routing_method_type == RoutingMethodType.DeepSeekV3
1039
- else router_logits,
1040
- routing_bias=e_score_correction_bias,
1030
+ routing_logits=router_logits,
1031
+ routing_bias=layer.e_score_correction_bias,
1032
x=x,
1033
w13_weight=layer.w13_weight,
1034
w13_weight_scale_inv=layer.w13_weight_scale,
@@ -1051,7 +1042,7 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):
1042
expert_offset=layer.ep_rank * layer.local_num_experts,
1043
local_num_experts=layer.local_num_experts,
1044
block_shape=self.weight_block_size,
1054
- routing_method_type=routing_method_type,
1045
+ routing_method_type=layer.routing_method_type,
1046
routed_scaling=layer.routed_scaling_factor,
1047
)
1048
else:
vllm/model_executor/layers/quantization/fp8.py
+3
-12
@@ -26,7 +26,6 @@ from vllm.model_executor.layers.fused_moe import (
26
)
27
from vllm.model_executor.layers.fused_moe.config import (
28
FusedMoEQuantConfig,
29
- RoutingMethodType,
29
)
30
from vllm.model_executor.layers.fused_moe.layer import UnquantizedFusedMoEMethod
31
from vllm.model_executor.layers.fused_moe.oracle.fp8 import (
@@ -980,17 +979,9 @@ class Fp8MoEMethod(FusedMoEMethodBase):
979
if self.block_quant:
980
import vllm.model_executor.layers.fused_moe.flashinfer_trtllm_moe # noqa: E501, F401
981
983
- e_score_correction_bias = (
984
- layer.e_score_correction_bias.to(x.dtype)
985
- if layer.e_score_correction_bias is not None
986
- else None
987
- )
988
- routing_method_type = layer.routing_method_type
982
return torch.ops.vllm.flashinfer_fused_moe_blockscale_fp8(
990
- routing_logits=router_logits.to(torch.float32)
991
- if routing_method_type == RoutingMethodType.DeepSeekV3
992
- else router_logits,
993
- routing_bias=e_score_correction_bias,
983
+ routing_logits=router_logits,
984
+ routing_bias=layer.e_score_correction_bias,
985
x=x,
986
w13_weight=layer.w13_weight,
987
w13_weight_scale_inv=layer.w13_weight_scale_inv,
@@ -1004,7 +995,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):
995
expert_offset=layer.ep_rank * layer.local_num_experts,
996
local_num_experts=layer.local_num_experts,
997
block_shape=self.weight_block_size,
1007
- routing_method_type=routing_method_type,
998
+ routing_method_type=layer.routing_method_type,
999
routed_scaling=layer.routed_scaling_factor,
1000
)
1001
else:
vllm/model_executor/models/minimax_m2.py
+1
@@ -107,6 +107,7 @@ class MiniMaxM2MoE(nn.Module):
107
renormalize=True,
108
quant_config=quant_config,
109
prefix=f"{prefix}.experts",
110
+ router_logits_dtype=torch.float32,
111
)
112
113
self.gate = ReplicatedLinear(