Use dual stream for DS MoE whenever cuda graph is used (instead of with token threshold) (#9405)

This commit is contained in:
Trevor Morris
2025-11-22 00:37:41 -08:00
committed by GitHub
parent 3805243245
commit 0eea17e3b5
+3 -2
View File
@@ -787,12 +787,13 @@ class DeepseekV2MoE(nn.Module):
gemm_output_zero_allocator: BumpAllocator = None, gemm_output_zero_allocator: BumpAllocator = None,
) -> torch.Tensor: ) -> torch.Tensor:
if not self._enable_a2a_moe: 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 ( if (
self.alt_stream is not None self.alt_stream is not None
and self.num_fused_shared_experts == 0 and self.num_fused_shared_experts == 0
and hidden_states.shape[0] > 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( return self.forward_normal_dual_stream(
hidden_states, hidden_states,