From 2df9e8b4b33112eed19d707b4e51e5a90ac39ce8 Mon Sep 17 00:00:00 2001
From: maocheng23 <35615230+maocheng23@users.noreply.github.com>
Date: Fri, 22 May 2026 14:23:57 -0700
Subject: [PATCH] [perf] DeepSeekV3: drop redundant FP32 upcasts in trtllm MoE
 paths (#25189)

---
 .../layers/moe/moe_runner/flashinfer_trtllm.py   | 16 ++--------------
 python/sglang/srt/models/deepseek_v2.py          |  1 +
 2 files changed, 3 insertions(+), 14 deletions(-)

diff --git a/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py b/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py
index 439840ad1..4e65c2a67 100644
--- a/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py
+++ b/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py
@@ -698,11 +698,7 @@ def fused_experts_none_to_flashinfer_trtllm_fp8(
             assert TopKOutputChecker.format_is_bypassed(topk_output)
 
             output = trtllm_fp8_block_scale_moe_wrapper(
-                routing_logits=(
-                    router_logits.to(torch.float32)
-                    if routing_method_type == RoutingMethodType.DeepSeekV3
-                    else router_logits
-                ),
+                routing_logits=router_logits,
                 routing_bias=correction_bias,
                 hidden_states=a_q,
                 hidden_states_scale=a_sf_t,
@@ -758,11 +754,7 @@ def fused_experts_none_to_flashinfer_trtllm_fp8(
         # during torch.compile for piecewise cuda graph.
         # 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(
             routing_logits=router_logits,
@@ -977,10 +969,6 @@ def fused_experts_none_to_flashinfer_trtllm_fp4(
         topk_config = topk_output.topk_config
         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 = (
             None
             if topk_config.correction_bias is None
diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py
index 44fd33cb6..1e5702b92 100644
--- a/python/sglang/srt/models/deepseek_v2.py
+++ b/python/sglang/srt/models/deepseek_v2.py
@@ -403,6 +403,7 @@ class MoEGate(nn.Module):
                 and _device_sm >= 90
             ):
                 if _device_sm in [100, 103] and self.weight.shape[0] == 256:
+                    # TODO: will check the dtype to be bf16
                     # router gemm output float32
                     logits = torch.empty(
                         hidden_states.shape[0],