[Feature] Coordinate FullCG prefill across DP-attention ranks (#35640)

Co-authored-by: Yuwei An <ayw.sirius19@gmail.com>

python/sglang/srt/batch_overlap/two_batch_overlap.py

@@ -736,8 +736,8 @@ class TboForwardBatchPreparer:
             "forward_mode",
             "is_extend_in_batch",
             "return_logprob",
-            "can_run_dp_cuda_graph",
-            "can_run_dp_breakable_cuda_graph",
+            "can_run_decode_cuda_graph",
+            "can_run_dp_prefill_cuda_graph",
             "dp_padding_mode",
             "global_forward_mode",
             "is_prefill_only",

python/sglang/srt/managers/schedule_batch.py

@@ -2161,8 +2161,8 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
 
     # For DP attention
     is_extend_in_batch: bool = False
-    can_run_dp_cuda_graph: bool = False
-    can_run_dp_breakable_cuda_graph: bool = False
+    can_run_decode_cuda_graph: bool = False
+    can_run_dp_prefill_cuda_graph: bool = False
     tbo_split_seq_index: Optional[int] = None
     # Rank-consistent forward mode for the recv skipper, derived from the MLP
     # sync all-gather (the TBO-only `global_forward_mode` is None without TBO).
@@ -3355,8 +3355,8 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
             spec_info=self.spec_info,
             global_num_tokens=self.global_num_tokens,
             global_num_tokens_for_logprob=self.global_num_tokens_for_logprob,
-            can_run_dp_cuda_graph=self.can_run_dp_cuda_graph,
-            can_run_dp_breakable_cuda_graph=self.can_run_dp_breakable_cuda_graph,
+            can_run_decode_cuda_graph=self.can_run_decode_cuda_graph,
+            can_run_dp_prefill_cuda_graph=self.can_run_dp_prefill_cuda_graph,
             is_extend_in_batch=self.is_extend_in_batch,
             is_prefill_only=self.is_prefill_only,
             seq_lens_cpu=self.seq_lens_cpu,

python/sglang/srt/managers/scheduler_components/dp_attn.py

@@ -221,8 +221,8 @@ def _update_gather_batch(
         batch.global_forward_mode = mlp_sync_info.global_forward_mode
 
     # Check forward mode for cuda graph
-    batch.can_run_dp_cuda_graph = mlp_sync_info.can_run_decode_cuda_graph
-    batch.can_run_dp_breakable_cuda_graph = mlp_sync_info.can_run_prefill_cuda_graph
+    batch.can_run_decode_cuda_graph = mlp_sync_info.can_run_decode_cuda_graph
+    batch.can_run_dp_prefill_cuda_graph = mlp_sync_info.can_run_prefill_cuda_graph
 
 
 def prepare_mlp_sync_batch_raw(
@@ -271,14 +271,15 @@ def prepare_mlp_sync_batch_raw(
         or local_batch.forward_mode.is_decode_or_idle()
         or local_batch.forward_mode.is_prebuilt()
     ) and not disable_cuda_graph
-    breakable_prefill = check_cuda_graph_backend(Phase.PREFILL, Backend.BREAKABLE)
+    coordinated_prefill = check_cuda_graph_backend(
+        Phase.PREFILL, Backend.BREAKABLE
+    ) or check_cuda_graph_backend(Phase.PREFILL, Backend.FULL)
     prefill_graph_runner = (
-        model_runner.prefill_cuda_graph_runner if breakable_prefill else None
+        model_runner.prefill_cuda_graph_runner if coordinated_prefill else None
     )
     can_run_prefill_cuda_graph = (
         local_batch is None
         or local_batch.forward_mode.is_idle()
-        # Breakable Cuda Graph Backend Check.
         or (
             local_batch.forward_mode in (ForwardMode.EXTEND, ForwardMode.MIXED)
             and (
@@ -287,7 +288,7 @@ def prepare_mlp_sync_batch_raw(
                     batch_size=local_batch.batch_size(),
                     num_tokens=local_batch.extend_num_tokens,
                     input_embeds=local_batch.input_embeds,
-                    replace_embeds=None,
+                    replace_embeds=local_batch.replace_embeds,
                     prefix_lens=local_batch.prefix_lens,
                     is_target_verify=local_batch.forward_mode.is_target_verify(),
                     capture_hidden_mode=None,
@@ -295,7 +296,7 @@ def prepare_mlp_sync_batch_raw(
                     lora_ineligible=prefill_graph_runner.enable_lora,
                 )
             )
-            and breakable_prefill
+            and coordinated_prefill
         )
     )
 

python/sglang/srt/model_executor/forward_batch_info.py

@@ -446,8 +446,8 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
 
     # For DP attention
     is_extend_in_batch: bool = False
-    can_run_dp_cuda_graph: bool = False
-    can_run_dp_breakable_cuda_graph: bool = False
+    can_run_decode_cuda_graph: bool = False
+    can_run_dp_prefill_cuda_graph: bool = False
     global_forward_mode: Optional[ForwardMode] = None
 
     # For two-batch overlap
@@ -703,7 +703,7 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
         self.global_num_tokens_for_logprob_gpu = torch.tensor(
             global_num_tokens_for_logprob, dtype=torch.int64
         ).to(device, non_blocking=True)
-        self.can_run_dp_cuda_graph = batch.can_run_dp_cuda_graph
+        self.can_run_decode_cuda_graph = batch.can_run_decode_cuda_graph
 
     @classmethod
     def init_new(
@@ -788,8 +788,8 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
             # Scalar config / flags
             return_logprob=batch.return_logprob,
             is_extend_in_batch=batch.is_extend_in_batch,
-            can_run_dp_cuda_graph=batch.can_run_dp_cuda_graph,
-            can_run_dp_breakable_cuda_graph=batch.can_run_dp_breakable_cuda_graph,
+            can_run_decode_cuda_graph=batch.can_run_decode_cuda_graph,
+            can_run_dp_prefill_cuda_graph=batch.can_run_dp_prefill_cuda_graph,
             global_forward_mode=batch.global_forward_mode,
             is_prefill_only=batch.is_prefill_only,
             spec_algorithm=batch.spec_algorithm,
@@ -1312,21 +1312,20 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
         ):
             # Joined ranks require real token counts instead of MAX_LEN padding.
             dp_padding_mode = DpPaddingMode.SUM_LEN
-        # Prefill breakable CUDA graph requires every DP rank to run the SAME
-        # captured shape. Under SUM_LEN each rank pads to its own local token
+        # Prefill CUDA graphs require every DP rank to run the same captured
+        # shape. Under SUM_LEN each rank pads to its own local token
         # count and can select a different capture bucket. This mismatches the
         # rank-coupled communication geometry: DP gather/combine uses
         # all_gather_into_tensor / reduce_scatter_tensor, while MoE backends may
         # use A2A dispatch/combine. Force MAX_LEN so every rank pads to the global
         # max and picks the same bucket.
         #
-        # Only force MAX_LEN when the batch fits a captured breakable prefill
-        # graph; larger prefills fall back to eager and keep the
-        # memory-efficient SUM_LEN. global_num_tokens is identical across ranks
-        # (all-gathered), so the decision is consistent cluster-wide.
+        # Larger prefills fall back to eager and keep the memory-efficient
+        # SUM_LEN. global_num_tokens is identical across ranks (all-gathered),
+        # so the decision is consistent cluster-wide.
         prefill_cg = get_exec().graph.cuda_graph_config.prefill
         if (
-            self.can_run_dp_breakable_cuda_graph
+            self.can_run_dp_prefill_cuda_graph
             and self.is_extend_in_batch
             and prefill_cg.bs
             and max(global_num_tokens) <= max(prefill_cg.bs)

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

@@ -694,7 +694,9 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
         )
 
         if self.require_mlp_sync:
-            is_bs_supported = is_bs_supported and forward_batch.can_run_dp_cuda_graph
+            is_bs_supported = (
+                is_bs_supported and forward_batch.can_run_decode_cuda_graph
+            )
 
         # NOTE: cuda graph cannot handle mixed batch (encoder_len = 0)
         # If mixed batch cannot be supported, then encoder_lens can be removed in cuda graph
@@ -735,7 +737,7 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
         ] and forward_batch.batch_size <= self._ragged_capture_slots(admission_tokens)
 
         is_dp_supported = (
-            forward_batch.can_run_dp_cuda_graph if self.require_mlp_sync else True
+            forward_batch.can_run_decode_cuda_graph if self.require_mlp_sync else True
         )
 
         is_encoder_lens_supported = (

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

@@ -40,6 +40,7 @@ from __future__ import annotations
 import copy
 import inspect
 import logging
+from collections.abc import Sequence
 from contextlib import contextmanager
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any, Dict, Optional, Union
@@ -826,19 +827,29 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
         return prefix_chunk_len, prefix_chunk_len * capture_req_slots
 
     def _select_prefix_capture_chunks(
-        self, forward_batch: ForwardBatch
+        self, prefix_lens: Sequence[int]
     ) -> Optional[int]:
         """Smallest captured variant covering the batch's max prefix, or None."""
-        max_prefix_len = max(
-            int(length) for length in forward_batch.extend_prefix_lens_cpu
-        )
+        max_prefix_len = max(int(length) for length in prefix_lens)
         real_n = _ceil_div(max_prefix_len, self._prefix_chunk_len)
         return next((n for n in self._prefix_capture_variants if n >= real_n), None)
 
+    def _has_uncapturable_chunked_prefix(
+        self, prefix_lens: Sequence[int] | None
+    ) -> bool:
+        return (
+            self._capture_chunked_prefix
+            and prefix_lens is not None
+            and any(int(length) > 0 for length in prefix_lens)
+            and self._select_prefix_capture_chunks(prefix_lens) is None
+        )
+
     def _shape_key(self, num_tokens: int, forward_batch: ForwardBatch) -> ShapeKey:
         variant = None
         if self._capture_chunked_prefix and self._has_prefix_hit(forward_batch):
-            captured_n = self._select_prefix_capture_chunks(forward_batch)
+            captured_n = self._select_prefix_capture_chunks(
+                forward_batch.extend_prefix_lens_cpu
+            )
             assert captured_n is not None, "prefix batch has no captured FullCG variant"
             variant = _chunked_prefix_variant(captured_n)
         return ShapeKey(size=num_tokens, variant_label=variant)
@@ -1055,7 +1066,6 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
         capture_hidden_mode,
         return_logprob: bool,
         lora_ineligible: bool = False,
-        chunked_prefix_uncapturable: bool = False,
     ) -> bool:
         """Rank-local replay eligibility: the single source of truth for
         ``can_run_graph`` (ForwardBatch, forward time) and the dp mlp-sync
@@ -1084,10 +1094,9 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
             and any(prefix_lens)
         ):
             return False
-        # FullCG's chunked-prefix topology covers a bounded prefix. The flag
-        # gating it is FULL-backend-only, so this is inert for the breakable
-        # vote path.
-        if chunked_prefix_uncapturable:
+        # FullCG's chunked-prefix topology covers a bounded prefix. Its capture
+        # flag is FullCG-only, so this is inert for the BreakableCG vote path.
+        if self._has_uncapturable_chunked_prefix(prefix_lens):
             return False
         # tc_piecewise captures with ForwardMode.EXTEND and spec_info=None.
         if is_target_verify:
@@ -1115,7 +1124,7 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
         # (min-reduced votes; also requires every rank to hold tokens).
         if (
             forward_batch.global_num_tokens_cpu is not None
-            and not forward_batch.can_run_dp_breakable_cuda_graph
+            and not forward_batch.can_run_dp_prefill_cuda_graph
         ):
             return False
 
@@ -1141,11 +1150,6 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
                     forward_batch
                 )
             ),
-            chunked_prefix_uncapturable=(
-                self._capture_chunked_prefix
-                and self._has_prefix_hit(forward_batch)
-                and self._select_prefix_capture_chunks(forward_batch) is None
-            ),
         ):
             return False
         if getattr(self, "enable_cp_v2_bcg_capture", False) and is_cp_v2_active(

python/sglang/srt/speculative/dspark_components/dspark_draft.py

@@ -444,7 +444,7 @@ class DraftBlockProposer:
     ) -> None:
         # The dense DSpark draft still reuses the target batch's graph tier.
         # Set graph eligibility before the DP-MoE-only metadata early return.
-        forward_batch.can_run_dp_cuda_graph = batch.can_run_dp_cuda_graph
+        forward_batch.can_run_decode_cuda_graph = batch.can_run_decode_cuda_graph
         if not self._dp_moe_sync or batch.global_num_tokens is None:
             return
         # Graph bucket selection uses the raw per-rank request counts.  Keep

python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py

@@ -318,7 +318,9 @@ class EAGLEDraftCudaGraphRunner(DecodeCudaGraphRunner):
         )
 
         if self.require_mlp_sync:
-            is_bs_supported = is_bs_supported and forward_batch.can_run_dp_cuda_graph
+            is_bs_supported = (
+                is_bs_supported and forward_batch.can_run_decode_cuda_graph
+            )
 
         return is_bs_supported
 

python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py

@@ -309,7 +309,9 @@ class EAGLEDraftExtendCudaGraphRunner(DecodeCudaGraphRunner):
         )
 
         if self.require_mlp_sync:
-            is_bs_supported = is_bs_supported and forward_batch.can_run_dp_cuda_graph
+            is_bs_supported = (
+                is_bs_supported and forward_batch.can_run_decode_cuda_graph
+            )
 
         return is_bs_supported
 

python/sglang/srt/speculative/frozen_kv_mtp_cuda_graph_runner.py

@@ -237,7 +237,9 @@ class FrozenKVMTPCudaGraphRunner(DecodeCudaGraphRunner):
             else cuda_graph_bs <= self.max_bs
         )
         if self.require_mlp_sync:
-            is_bs_supported = is_bs_supported and forward_batch.can_run_dp_cuda_graph
+            is_bs_supported = (
+                is_bs_supported and forward_batch.can_run_decode_cuda_graph
+            )
         return is_bs_supported
 
     def capture_one_shape(

python/sglang/srt/speculative/frozen_kv_mtp_worker_v2.py

@@ -485,7 +485,7 @@ class FrozenKVMTPDraftWorker(EagleDraftWorkerBase, TpModelWorker):
                 self.cuda_graph_runner.execute(forward_batch)
             )
         else:
-            forward_batch.can_run_dp_cuda_graph = False
+            forward_batch.can_run_decode_cuda_graph = False
             parent_list, top_scores_index, draft_tokens = self.draft_forward(
                 forward_batch
             )

python/sglang/srt/speculative/multi_layer_eagle_draft_extend_cuda_graph_runner.py

@@ -255,7 +255,9 @@ class MultiLayerEagleDraftExtendCudaGraphRunner(DecodeCudaGraphRunner):
         )
 
         if self.require_mlp_sync:
-            is_bs_supported = is_bs_supported and forward_batch.can_run_dp_cuda_graph
+            is_bs_supported = (
+                is_bs_supported and forward_batch.can_run_decode_cuda_graph
+            )
 
         return is_bs_supported