[Bugfix] Keep a shared MAX_LEN prefill CUDA graph bucket when the graph captures a DP gather (MegaMoE sparse-DP hang) (#37933)

Co-authored-by: shyeh25 <206795756+shyeh25@users.noreply.github.com>
Co-authored-by: Po-Han Huang (NVIDIA) <53919306+nvpohanh@users.noreply.github.com>

python/sglang/srt/layers/dp_attention.py

@@ -820,12 +820,19 @@ def _dp_gather_via_all_gatherv(
     get_tp_group().all_gatherv(local_real, sizes=sizes, output=global_tokens)
 
 
+def _note_dp_gather_in_prefill_graph() -> None:
+    dp = get_flags().dp
+    if dp.capturing_prefill_graph:
+        dp.prefill_graph_has_dp_gather = True
+
+
 def _dp_gather(
     global_tokens: torch.Tensor,
     local_tokens: torch.Tensor,
     forward_batch: ForwardBatch,
     is_partial: bool,
 ):
+    _note_dp_gather_in_prefill_graph()
     if (
         is_dp_gatherv_active()
         and forward_batch.dp_padding_mode is not None
@@ -883,6 +890,7 @@ def dp_scatter(
     global_tokens: torch.Tensor,  # input
     forward_batch: ForwardBatch,
 ):
+    _note_dp_gather_in_prefill_graph()
     # local_num_tokens is not necessarily the same as local_tokens.shape[0],
     # since local_tokens may be padded for cuda graph
     local_start_pos, local_num_tokens = get_dp_local_info(forward_batch)
@@ -899,6 +907,7 @@ def dp_scatter(
 
 
 def dp_reduce_scatter_tensor(output: torch.Tensor, input: torch.Tensor):
+    _note_dp_gather_in_prefill_graph()
     if is_dp_gatherv_active():
         # Variable-length combine matching all_gatherv dispatch: scatter the
         # global (sum_len) tensor back to per-rank token counts. Fall through to

python/sglang/srt/model_executor/forward_batch_info.py

@@ -53,6 +53,7 @@ from sglang.srt.model_executor.forward_batch_deepseek_mha_mixin import (
 )
 from sglang.srt.runtime_context import (
     get_exec,
+    get_flags,
     get_lora,
     get_parallel,
     mamba_cache_chunk_size,
@@ -303,13 +304,19 @@ def compute_local_num_token_non_padded_cpu(
 
 
 def prefill_graph_tolerates_sum_len() -> bool:
-    """Whether MegaMoE may replay prefill graphs with local shapes."""
+    """Whether MegaMoE may replay prefill graphs with per-rank SUM_LEN buckets.
+
+    The graph body is captured with MAX_LEN geometry, so a graph that recorded
+    a DP gather/scatter only replays correctly when every rank uses one bucket.
+    """
     from sglang.srt.layers.attention.dsa.utils import is_dsa_enable_prefill_cp
     from sglang.srt.layers.cp.utils import is_mla_cp_enabled
     from sglang.srt.layers.moe.utils import get_moe_a2a_backend
 
     if not get_moe_a2a_backend().is_megamoe():
         return False
+    if get_flags().dp.prefill_graph_has_dp_gather:
+        return False
     return not (is_dsa_enable_prefill_cp() or is_mla_cp_enabled())
 
 

python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py

@@ -126,6 +126,7 @@ from sglang.srt.model_executor.runner_utils.pool import (
 from sglang.srt.model_loader.utils import resolve_language_model
 from sglang.srt.runtime_context import (
     get_exec,
+    get_flags,
     get_memory,
     get_parallel,
     get_schedule,
@@ -844,9 +845,8 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
         # DSV4 DP attention / DeepEP collectives need every DP rank to enter
         # the same replay path. Sparse-DP batches (one or more ranks with
         # zero local tokens) fall back to eager to avoid hanging ranks.
-        # MegaMoE is exempt (prefill_graph_tolerates_sum_len): its idle ranks
-        # still execute MegaMoE with 0 tokens, so per-rank SUM_LEN buckets stay
-        # collective-safe and need no eager fallback.
+        # MegaMoE graphs without a captured DP gather tolerate per-rank buckets
+        # and an eager idle rank; graphs with one fall through to the check.
         global_num_tokens = forward_batch.global_num_tokens_cpu
         if global_num_tokens is None:
             return False
@@ -1412,13 +1412,23 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
         # Warm up + autotune kernels once before capture (run-once across the
         # decode + prefill runners; see BaseRunner.warmup).
         self.warmup()
-        with freeze_gc(get_exec().graph.enable_cudagraph_gc):
-            with graph_capture(
-                stream=get_or_create_global_graph_capture_stream()
-            ) as graph_capture_context:
-                self.stream = graph_capture_context.stream
-                with self.backend.capture_session(self.stream):
-                    self._capture_one_stream()
+        dp_flags = get_flags().dp
+        dp_flags.capturing_prefill_graph = True
+        try:
+            with freeze_gc(get_exec().graph.enable_cudagraph_gc):
+                with graph_capture(
+                    stream=get_or_create_global_graph_capture_stream()
+                ) as graph_capture_context:
+                    self.stream = graph_capture_context.stream
+                    with self.backend.capture_session(self.stream):
+                        self._capture_one_stream()
+        finally:
+            dp_flags.capturing_prefill_graph = False
+        if dp_flags.prefill_graph_has_dp_gather:
+            logger.info(
+                "Prefill CUDA graph captured a DP gather/scatter; "
+                "DP ranks will replay a shared MAX_LEN bucket."
+            )
 
     def _capture_one_stream(self) -> None:
         avail_mem = get_available_gpu_memory(

python/sglang/srt/runtime_context.py

@@ -565,6 +565,10 @@ class DpFlags(_FlagGroupBase):
     # Hybrid-SSM models materialize idle ranks via the MAX_LEN fabricated-row
     # conversion (set when hf_config has hybrid_override_pattern).
     max_len_with_idle: bool = False
+    # Set while the prefill CUDA graph runner captures; latched by the DP
+    # gather/scatter helpers, whose captured geometry needs one shared bucket.
+    capturing_prefill_graph: bool = False
+    prefill_graph_has_dp_gather: bool = False
     # DP gathered-buffer allocation metadata (model hidden size / dtype /
     # device), set by initialize_dp_attention alongside the flags above.
     buffer_hidden_size: Any = None