From 87992eeec4072995e8fa98fb2d0f3a7e5e581f2d Mon Sep 17 00:00:00 2001 From: Khoa Pham Date: Thu, 9 Jul 2026 16:59:39 -0700 Subject: [PATCH] [DeepSeek V2] Reorder dual-stream MoE to main-first to avoid CUDA graph stream explosion (#30460) Co-authored-by: Cursor Co-authored-by: Jiminator Co-authored-by: Claude Opus 4.8 (1M context) --- .../runner_utils/capture_mode.py | 5 ++++ python/sglang/srt/models/deepseek_v2.py | 25 +++++++++++-------- python/sglang/srt/runtime_context.py | 5 ++++ python/sglang/srt/utils/common.py | 5 ++++ 4 files changed, 30 insertions(+), 10 deletions(-) diff --git a/python/sglang/srt/model_executor/runner_utils/capture_mode.py b/python/sglang/srt/model_executor/runner_utils/capture_mode.py index 293ec1b25..54d802254 100644 --- a/python/sglang/srt/model_executor/runner_utils/capture_mode.py +++ b/python/sglang/srt/model_executor/runner_utils/capture_mode.py @@ -66,8 +66,13 @@ def _set_capture_lora_variant(variant: Optional[str]) -> None: @contextmanager def model_capture_mode(): global is_capture_mode + from sglang.srt.runtime_context import get_flags + + # Disable dispose_tensor() during capture: freeing mid-capture records data_ptr()==0 into the graph. is_capture_mode = True + get_flags().capture.disable_dispose_tensor = True try: yield finally: is_capture_mode = False + get_flags().capture.disable_dispose_tensor = False diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index 253c30afe..ef59bc0ea 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -929,23 +929,22 @@ class DeepseekV2MoE(nn.Module): *, use_flashinfer_trtllm_bypass: bool = False, ) -> torch.Tensor: - # Note(kpham-sgl): launch the shared expert BEFORE the routed call. - # The routed deep_gemm pre-permute calls `dispose_tensor` which - # `set_()`s `hidden_states` to empty (host-side); any later kernel - # launch consuming `hidden_states` then captures `data_ptr() == 0` - # into the decode CUDA graph and replays from null. + # Note(kpham-sgl): issue order satisfies 3 constraints: + # - no stream explosion: main (routed) issued before alt block -> capture reuses 1 alt stream; + # - PDL overlap: routed is the last main-stream kernel (fuses w/ residual add); + # - dispose_tensor: disabled during capture (CaptureFlags.disable_dispose_tensor) so the routed + # deep_gemm does not free hidden_states, which the shared expert reads on the alt stream. current_stream = torch.cuda.current_stream() self.alt_stream.wait_stream(current_stream) + has_shared_output = ( + hidden_states.shape[0] > 0 and self.num_fused_shared_experts == 0 + ) server_args = get_server_args() dispatch_info = ( ExpertLocationDispatchInfo.init_new(layer_id=self.layer_id) if server_args.enable_eplb and not self.is_nextn else None ) - with torch.cuda.stream(self.alt_stream): - shared_output = self._forward_shared_experts( - hidden_states, gemm_output_zero_allocator - ) # router_logits: (num_tokens, n_experts) router_logits = self.gate(hidden_states, gemm_output_zero_allocator) if use_flashinfer_trtllm_bypass: @@ -967,7 +966,7 @@ class DeepseekV2MoE(nn.Module): **topk_kwargs, ) deferred_finalize = ( - shared_output is not None + has_shared_output and not self._shared_expert_tp1 and topk_output.format == TopKOutputFormat.BYPASSED and self.experts.supports_deferred_finalize @@ -988,6 +987,12 @@ class DeepseekV2MoE(nn.Module): ): final_hidden_states *= self.routed_scaling_factor + # Shared expert on alt stream, issued AFTER the main (routed) branch. See note above. + with torch.cuda.stream(self.alt_stream): + shared_output = self._forward_shared_experts( + hidden_states, gemm_output_zero_allocator + ) + current_stream.wait_stream(self.alt_stream) if deferred_finalize: diff --git a/python/sglang/srt/runtime_context.py b/python/sglang/srt/runtime_context.py index 284814c95..bbe16696e 100644 --- a/python/sglang/srt/runtime_context.py +++ b/python/sglang/srt/runtime_context.py @@ -284,6 +284,11 @@ class CaptureFlags(_FlagGroupBase): # False clears it during warmup (the only post-publish writer). enable_torch_compile: bool = False + # Set for the duration of decode/spec graph capture (model_capture_mode). + # While set, dispose_tensor() is a no-op so deep_gemm's pre-permute does not + # free hidden_states that the dual-stream MoE shared expert reads afterward. + disable_dispose_tensor: bool = False + @dataclasses.dataclass class MoeFlags(_FlagGroupBase): diff --git a/python/sglang/srt/utils/common.py b/python/sglang/srt/utils/common.py index 09b9fbfde..005514f00 100644 --- a/python/sglang/srt/utils/common.py +++ b/python/sglang/srt/utils/common.py @@ -3417,6 +3417,11 @@ def dispose_tensor(x: torch.Tensor): if is_in_tc_piecewise_cuda_graph(): return + from sglang.srt.runtime_context import get_flags + + if get_flags().capture.disable_dispose_tensor: + return + x.set_(torch.empty((0,), device=x.device, dtype=x.dtype))