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

[perf] DeepSeekV3: drop redundant FP32 upcasts in trtllm MoE paths (#25189)

Browse Source
This commit is contained in:
maocheng23
2026-05-22 14:23:57 -07:00
committed by GitHub
parent b73278e4e0
commit 2df9e8b4b3
2 changed files with 3 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py
+1 -13
View File
@@ -698,11 +698,7 @@ def fused_experts_none_to_flashinfer_trtllm_fp8(
assert TopKOutputChecker.format_is_bypassed(topk_output) assert TopKOutputChecker.format_is_bypassed(topk_output)
output = trtllm_fp8_block_scale_moe_wrapper( output = trtllm_fp8_block_scale_moe_wrapper(
routing_logits=( routing_logits=router_logits,
router_logits.to(torch.float32)
if routing_method_type == RoutingMethodType.DeepSeekV3
else router_logits
),
routing_bias=correction_bias, routing_bias=correction_bias,
hidden_states=a_q, hidden_states=a_q,
hidden_states_scale=a_sf_t, hidden_states_scale=a_sf_t,
@@ -758,10 +754,6 @@ def fused_experts_none_to_flashinfer_trtllm_fp8(
# 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) router_logits = router_logits.to(torch.bfloat16)
output = trtllm_fp8_per_tensor_scale_moe_wrapper( output = trtllm_fp8_per_tensor_scale_moe_wrapper(
@@ -977,10 +969,6 @@ def fused_experts_none_to_flashinfer_trtllm_fp4(
topk_config = topk_output.topk_config topk_config = topk_output.topk_config
routing_method_type = quant_info.routing_method_type routing_method_type = quant_info.routing_method_type
# DeepSeekV3 style routing requires float32 router logits
if routing_method_type == RoutingMethodType.DeepSeekV3:
router_logits = router_logits.to(torch.float32)
correction_bias = ( correction_bias = (
None None
if topk_config.correction_bias is None if topk_config.correction_bias is None
python/sglang/srt/models/deepseek_v2.py
+1
View File
@@ -403,6 +403,7 @@ class MoEGate(nn.Module):
and _device_sm >= 90 and _device_sm >= 90
): ):
if _device_sm in [100, 103] and self.weight.shape[0] == 256: if _device_sm in [100, 103] and self.weight.shape[0] == 256:
# TODO: will check the dtype to be bf16
# router gemm output float32 # router gemm output float32
logits = torch.empty( logits = torch.empty(
hidden_states.shape[0], hidden_states.shape[0],