minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Tiny fix trtllm_fp8_per_tensor_scale_moe_wrapper router_logits dtype (#22006)

Browse Source
This commit is contained in:
Qiaolin Yu
2026-04-05 21:11:45 -07:00
committed by GitHub
parent c4240218cb
commit f407461ec8
1 changed files with 8 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py
+8 -1
View File
@@ -518,8 +518,15 @@ def fused_experts_none_to_flashinfer_trtllm_fp8(
# Move kernel call outside context manager to avoid graph breaks # Move kernel call outside context manager to avoid graph breaks
# during torch.compile for piecewise cuda graph. # during torch.compile for piecewise cuda graph.
# Use custom op wrapper for torch.compile compatibility. # Use custom op wrapper for torch.compile compatibility.
# The DeepSeekV3 routing method requires float32 router logits.
if routing_method_type == RoutingMethodType.DeepSeekV3:
router_logits = router_logits.to(torch.float32)
else:
router_logits = router_logits.to(torch.bfloat16)
output = trtllm_fp8_per_tensor_scale_moe_wrapper( output = trtllm_fp8_per_tensor_scale_moe_wrapper(
routing_logits=router_logits.to(torch.bfloat16), routing_logits=router_logits,
routing_bias=routing_bias_cast, routing_bias=routing_bias_cast,
hidden_states=a_q, hidden_states=a_q,
gemm1_weights=quant_info.w13_weight, gemm1_weights=quant_info.w13_weight,