[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_decode_cuda_graph",
-            "can_run_dp_breakable_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_decode_cuda_graph: bool = False
-    can_run_dp_breakable_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_decode_cuda_graph=self.can_run_decode_cuda_graph,
-            can_run_dp_breakable_cuda_graph=self.can_run_dp_breakable_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_decode_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_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_decode_cuda_graph: bool = False
-    can_run_dp_breakable_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_decode_cuda_graph=batch.can_run_decode_cuda_graph,
-            can_run_dp_breakable_cuda_graph=batch.can_run_dp_breakable_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
+        # Prefill CUDA graphs require every DP rank to run the same captured
-        # captured shape. Under SUM_LEN each rank pads to its own local token
+        # 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
+        # Larger prefills fall back to eager and keep the memory-efficient
-        # graph; larger prefills fall back to eager and keep the
+        # SUM_LEN. global_num_tokens is identical across ranks (all-gathered),
-        # memory-efficient SUM_LEN. global_num_tokens is identical across ranks
+        # so the decision is consistent cluster-wide.
-        # (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(
+        max_prefix_len = max(int(length) for length in prefix_lens)
-            int(length) for length in forward_batch.extend_prefix_lens_cpu
-        )
         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
+        # FullCG's chunked-prefix topology covers a bounded prefix. Its capture
-        # gating it is FULL-backend-only, so this is inert for the breakable
+        # flag is FullCG-only, so this is inert for the BreakableCG vote path.
-        # vote path.
+        if self._has_uncapturable_chunked_prefix(prefix_lens):
-        if chunked_prefix_uncapturable:
             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
 

test/registered/dp_attn/test_dp_attention_bcg_kl.py

@@ -1,6 +1,11 @@
 from __future__ import annotations
 
+import os
 import random
+import re
+import shutil
+import tempfile
+import time
 import unittest
 
 import numpy as np
@@ -23,7 +28,12 @@ from sglang.test.test_utils import (
     popen_launch_server,
 )
 
-register_cuda_ci(est_time=160, stage="base-b", runner_config="2-gpu-large")
+register_cuda_ci(est_time=320, stage="base-b", runner_config="2-gpu-large")
+
+PREFILL_GRAPH_REPLAY_PATTERN = re.compile(r"Prefill batch.*cuda graph: True")
+CACHED_PREFIX_GRAPH_REPLAY_PATTERN = re.compile(
+    r"Prefill batch.*#cached-token: [1-9][0-9]*.*cuda graph: True"
+)
 
 
 # ---------------------------------------------------------------------------
@@ -185,10 +195,11 @@ def _select_attention_backend():
     )
 
 
-class TestDPAttentionBreakablePrefillCudaGraphKL(CustomTestCase):
+class _DPAttentionPrefillCudaGraphKLMixin:
     num_samples = 48
     max_prompt_tokens = 1024
     max_new_tokens = 256
+    prefill_backend: str
 
     @classmethod
     def setUpClass(cls):
@@ -196,6 +207,11 @@ class TestDPAttentionBreakablePrefillCudaGraphKL(CustomTestCase):
         cls.model = DEFAULT_TARGET_MODEL_EAGLE_DP_ATTN
         cls.base_url = DEFAULT_URL_FOR_TEST
         cls.attention_backend = _select_attention_backend()
+        cls.log_dir = tempfile.mkdtemp(prefix=f"dp_attn_{cls.prefill_backend}_")
+        cls.stdout_path = os.path.join(cls.log_dir, "server.out")
+        cls.stderr_path = os.path.join(cls.log_dir, "server.err")
+        cls.stdout = open(cls.stdout_path, "w")
+        cls.stderr = open(cls.stderr_path, "w")
         cls.process = popen_launch_server(
             cls.model,
             cls.base_url,
@@ -212,12 +228,15 @@ class TestDPAttentionBreakablePrefillCudaGraphKL(CustomTestCase):
                 cls.attention_backend,
                 "--moe-runner-backend",
                 "triton",
-                "--cuda-graph-backend-prefill=breakable",
+                f"--cuda-graph-backend-prefill={cls.prefill_backend}",
                 "--chunked-prefill-size",
                 "2048",
+                "--prefill-max-requests",
+                "2",
                 "--mem-fraction-static",
                 "0.70",
             ],
+            return_stdout_stderr=(cls.stdout, cls.stderr),
         )
 
         server_info = requests.get(f"{cls.base_url}/server_info", timeout=30).json()
@@ -235,6 +254,36 @@ class TestDPAttentionBreakablePrefillCudaGraphKL(CustomTestCase):
     def tearDownClass(cls):
         if hasattr(cls, "process") and cls.process:
             kill_process_tree(cls.process.pid)
+        for attr in ("stdout", "stderr"):
+            output = getattr(cls, attr, None)
+            if output is not None and not output.closed:
+                output.close()
+        if hasattr(cls, "log_dir"):
+            shutil.rmtree(cls.log_dir, ignore_errors=True)
+
+    def _wait_for_prefill_graph_replay(
+        self,
+        offsets,
+        pattern=PREFILL_GRAPH_REPLAY_PATTERN,
+        case="a lone request",
+    ):
+        deadline = time.monotonic() + 30
+        while time.monotonic() < deadline:
+            chunks = []
+            for path, offset in zip(
+                (self.stdout_path, self.stderr_path), offsets, strict=True
+            ):
+                with open(path, "rb") as log:
+                    log.seek(offset)
+                    chunks.append(log.read().decode(errors="replace"))
+            logs = "\n".join(chunks)
+            if pattern.search(logs):
+                return
+            time.sleep(0.5)
+        self.fail(
+            f"No {self.prefill_backend} prefill CUDA graph replay was logged "
+            f"for {case}"
+        )
 
     def test_prefill_and_decode_cache_hit_kl_is_zero(self):
         server_info = requests.get(self.base_url + "/server_info", timeout=30).json()
@@ -243,15 +292,24 @@ class TestDPAttentionBreakablePrefillCudaGraphKL(CustomTestCase):
         self.assertTrue(server_info["enable_deterministic_inference"])
         self.assertEqual(server_info["attention_backend"], self.attention_backend)
         self.assertEqual(
-            server_info["cuda_graph_config"]["prefill"]["backend"], "breakable"
+            server_info["cuda_graph_config"]["prefill"]["backend"],
+            self.prefill_backend,
         )
 
         print("=== Radix Cache KL Divergence Eval ===")
         print(f"Server: {self.base_url}  Samples: {self.num_samples}\n")
 
+        offsets = [
+            os.path.getsize(path) for path in (self.stdout_path, self.stderr_path)
+        ]
         prefill_kl = test_prefill_cache_hit(
             self.base_url, self.input_ids, self.max_new_tokens
         )
+        self._wait_for_prefill_graph_replay(
+            offsets,
+            CACHED_PREFIX_GRAPH_REPLAY_PATTERN,
+            "a cached-prefix request",
+        )
         decode_kl = test_decode_cache_hit(
             self.base_url, self.input_ids, self.max_new_tokens
         )
@@ -259,6 +317,32 @@ class TestDPAttentionBreakablePrefillCudaGraphKL(CustomTestCase):
         self.assertEqual(prefill_kl, 0.0)
         self.assertEqual(decode_kl, 0.0)
 
+    def test_lone_request_replays_prefill_cuda_graph(self):
+        _flush_cache(self.base_url)
+        offsets = [
+            os.path.getsize(path) for path in (self.stdout_path, self.stderr_path)
+        ]
+        result = _generate(
+            self.base_url,
+            self.input_ids[0],
+            max_new_tokens=1,
+            return_logprob=True,
+        )
+        self.assertNotIn("error", result)
+        self._wait_for_prefill_graph_replay(offsets)
+
+
+class TestDPAttentionBreakablePrefillCudaGraphKL(
+    _DPAttentionPrefillCudaGraphKLMixin, CustomTestCase
+):
+    prefill_backend = "breakable"
+
+
+class TestDPAttentionFullPrefillCudaGraphKL(
+    _DPAttentionPrefillCudaGraphKLMixin, CustomTestCase
+):
+    prefill_backend = "full"
+
 
 if __name__ == "__main__":
     unittest.main()

test/registered/spec/dspark/test_dspark_dp_tier.py

@@ -67,7 +67,7 @@ class TestDraftDpSyncMetadata(CustomTestCase):
         batch = SimpleNamespace(
             global_num_tokens=[1, 3, 0, 2],
             global_num_tokens_for_logprob=[1, 3, 0, 2],
-            can_run_dp_cuda_graph=True,
+            can_run_decode_cuda_graph=True,
         )
 
         with patch(
@@ -84,7 +84,7 @@ class TestDraftDpSyncMetadata(CustomTestCase):
         self.assertEqual(forward_batch.num_token_non_padded.item(), 6)
         self.assertEqual(forward_batch.num_token_non_padded.dtype, torch.int32)
         self.assertEqual(forward_batch.num_token_non_padded_cpu, 6)
-        self.assertTrue(forward_batch.can_run_dp_cuda_graph)
+        self.assertTrue(forward_batch.can_run_decode_cuda_graph)
 
 
 class TestBusyIdleGraphKeyIdentity(CustomTestCase):

test/registered/unit/model_executor/test_mlp_sync_pad_unpad.py

@@ -59,7 +59,7 @@ class TestMlpSyncPadUnpad(CustomTestCase):
         batch = SimpleNamespace(
             global_num_tokens=[2, 0, 3],
             global_num_tokens_for_logprob=[2, 0, 3],
-            can_run_dp_cuda_graph=True,
+            can_run_decode_cuda_graph=True,
         )
 
         fb.init_mlp_sync_metadata(batch, torch.device("cpu"))
@@ -71,7 +71,7 @@ class TestMlpSyncPadUnpad(CustomTestCase):
         torch.testing.assert_close(
             fb.global_num_tokens_for_logprob_gpu, torch.tensor([4, 0, 6])
         )
-        self.assertTrue(fb.can_run_dp_cuda_graph)
+        self.assertTrue(fb.can_run_decode_cuda_graph)
 
     def test_draft_input_without_hidden_states_can_be_padded(self):
         spec_info = SimpleNamespace(