From 0eea17e3b549fe9038c81df2146eb81a5ece2008 Mon Sep 17 00:00:00 2001 From: Trevor Morris Date: Sat, 22 Nov 2025 08:37:41 +0000 Subject: [PATCH] Use dual stream for DS MoE whenever cuda graph is used (instead of with token threshold) (#9405) --- python/sglang/srt/models/deepseek_v2.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index 584b15bf4..154a6de7f 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -787,12 +787,13 @@ class DeepseekV2MoE(nn.Module): gemm_output_zero_allocator: BumpAllocator = None, ) -> torch.Tensor: if not self._enable_a2a_moe: - DUAL_STREAM_TOKEN_THRESHOLD = 1024 + from sglang.srt.model_executor.cuda_graph_runner import get_is_capture_mode + if ( self.alt_stream is not None and self.num_fused_shared_experts == 0 and hidden_states.shape[0] > 0 - and hidden_states.shape[0] <= DUAL_STREAM_TOKEN_THRESHOLD + and get_is_capture_mode() ): return self.forward_normal_dual_stream( hidden_states,