[DeepSeek V2] Reorder dual-stream MoE to main-first to avoid CUDA graph stream explosion (#30460)

Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Jiminator <jimmysh341@gmail.com>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

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

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:

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):

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))