biondizzle/vllm
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

[Bug][MoE] Fix TRTLLM NVFP4 Routing Kernel Precision (#36725)

Browse Source 
Signed-off-by: Robert Shaw <robshaw@redhat.com>
Co-authored-by: Robert Shaw <robshaw@redhat.com>
This commit is contained in:
Robert Shaw
2026-03-23 21:19:06 +01:00
committed by GitHub
parent 38364a7e32
commit 5bf3c42d4c
1 changed files with 2 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py
View File

@@ -298,10 +298,7 @@ class TrtLlmNvFp4ExpertsMonolithic(
and self.routing_method_type != RoutingMethodType.Llama4
)
# Prepare routing bias into kernel format.
routing_bias = e_score_correction_bias
if routing_bias is not None:
routing_bias = routing_bias.to(torch.bfloat16)
# Prepare router logits for kernel format.
router_logits = (
router_logits.to(torch.float32)
if self.routing_method_type == RoutingMethodType.DeepSeekV3
@@ -311,7 +308,7 @@ class TrtLlmNvFp4ExpertsMonolithic(
# Invoke kernel.
return flashinfer.fused_moe.trtllm_fp4_block_scale_moe(
routing_logits=router_logits,
routing_bias=routing_bias,
routing_bias=e_score_correction_bias,
hidden_states=hidden_states,
hidden_states_scale=a1q_scale.view(torch.float8_e4m3fn).reshape(
*hidden_states.shape[:-1], -1