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

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

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(