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, 04:30 UTC
[Bug][MoE] Fix TRTLLM NVFP4 Routing Kernel Precision (#36725)
Signed-off-by: Robert Shaw <robshaw@redhat.com> Co-authored-by: Robert Shaw <robshaw@redhat.com>
Robert Shaw committed
Mar 23, 2026 at 21:19 UTC
5bf3c42d4c7e99f4898b18a715be515f44a22f44
1 file changed
+2
-5
vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py
+2
-5
@@ -298,10 +298,7 @@ class TrtLlmNvFp4ExpertsMonolithic(
298
and self.routing_method_type != RoutingMethodType.Llama4
299
)
300
301
- # Prepare routing bias into kernel format.
302
- routing_bias = e_score_correction_bias
303
- if routing_bias is not None:
304
- routing_bias = routing_bias.to(torch.bfloat16)
301
+ # Prepare router logits for kernel format.
302
router_logits = (
303
router_logits.to(torch.float32)
304
if self.routing_method_type == RoutingMethodType.DeepSeekV3
@@ -311,7 +308,7 @@ class TrtLlmNvFp4ExpertsMonolithic(
308
# Invoke kernel.
309
return flashinfer.fused_moe.trtllm_fp4_block_scale_moe(
310
routing_logits=router_logits,
314
- routing_bias=routing_bias,
311
+ routing_bias=e_score_correction_bias,
312
hidden_states=hidden_states,
313
hidden_states_scale=a1q_scale.view(torch.float8_e4m3fn).reshape(
314
*hidden_states.shape[:-1], -1