[Feature] Support PP in full prefill CUDA graphs (#35451)

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

python/sglang/srt/model_executor/cuda_graph_buffer_registry.py

@@ -938,6 +938,30 @@ def build_prefill_registry(
                     "prefill registry; cannot adopt."
                 )
         reg.register_slot(slot, bind=bind)
+
+    if source is not None:
+        pp = getattr(source, "pp_proxy_tensors", None)
+        if pp is not None:
+
+            def _pp_source(key):
+                def _fn(_fb, ctx):
+                    ppx = ctx.pp_proxy_tensors
+                    return None if ppx is None else ppx.tensors[key]
+
+                return _fn
+
+            for _key, _backing in pp.items():
+                reg.register_slot(
+                    GraphSlot(
+                        name=f"pp_proxy_tensors.{_key}",
+                        shape_fn=lambda _bs, _mt, _s=tuple(_backing.shape): _s,
+                        dtype=_backing.dtype,
+                        axis="tokens",
+                        padding_policy=PaddingPolicy.ZERO,
+                        source_fn=_pp_source(_key),
+                    ),
+                    bind=_backing,
+                )
     return reg
 
 

python/sglang/srt/model_executor/model_runner_components/cuda_graph_setup.py

@@ -57,6 +57,29 @@ if TYPE_CHECKING:
 logger = logging.getLogger(__name__)
 
 
+def _align_pipeline_layers(layers: list, layer_model) -> list:
+    has_start_layer = hasattr(layer_model, "start_layer")
+    has_end_layer = hasattr(layer_model, "end_layer")
+    assert (
+        has_start_layer == has_end_layer
+    ), "pipeline layer ranges must define start_layer and end_layer together"
+    start_layer = layer_model.start_layer if has_start_layer else 0
+    end_layer = layer_model.end_layer if has_end_layer else len(layer_model.layers)
+    assert isinstance(start_layer, int) and isinstance(
+        end_layer, int
+    ), "pipeline layer ranges must define integer start_layer and end_layer"
+    assert 0 <= start_layer <= end_layer <= len(layer_model.layers), (
+        f"invalid pipeline layer range [{start_layer}, {end_layer}) for "
+        f"{len(layer_model.layers)} layers"
+    )
+    assert (
+        len(layers) <= end_layer - start_layer
+    ), f"found {len(layers)} layers in PP range [{start_layer}, {end_layer})"
+    return (
+        [None] * start_layer + layers + [None] * (len(layer_model.layers) - end_layer)
+    )
+
+
 class GraphCapture(msgspec.Struct, frozen=True, kw_only=True):
     runner: Optional[BaseRunner]
     memory_phase: str
@@ -370,6 +393,9 @@ def capture_prefill_graph(
         model_runner.mha_companion_layers,
     ) = compute_attention_and_moe_layers(layer_model)
 
+    model_runner.attention_layers = _align_pipeline_layers(
+        model_runner.attention_layers, layer_model
+    )
     if len(model_runner.attention_layers) < model_runner.model_config.num_hidden_layers:
         # TODO(yuwei): support Non-Standard GQA
         log_info_on_rank0(
@@ -378,6 +404,10 @@ def capture_prefill_graph(
         )
         return result(None)
 
+    model_runner.mha_companion_layers = _align_pipeline_layers(
+        model_runner.mha_companion_layers, layer_model
+    )
+
     tic = time.perf_counter()
     before_mem = get_available_gpu_memory(model_runner.device, model_runner.gpu_id)
     role = "draft" if model_runner.is_draft_worker else "target"

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

@@ -327,6 +327,12 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
             hidden_size=input_embeds_hidden_size,
             dtype=self.model_runner.dtype,
             enable_mamba_track=self.mamba_track_enabled,
+            pp_size=self.pp_size,
+            hc_hidden_size=model_runner.model_config.hc_hidden_size,
+            pp_proxy_topk_size=model_runner.get_pp_proxy_topk_size(),
+            pp_proxy_residual_num_blocks=(
+                model_runner.get_pp_proxy_residual_num_blocks()
+            ),
         )
         self.buffers.share_buffers()
         # Token-axis FB-shared slot registry adopting PrefillInputBuffers
@@ -379,6 +385,11 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
         self._capture_lora = False
         self.enable_cp_v2_bcg_capture = False
         self.prefill_cp_bcg_input: Optional[PrefillCPBCGInput] = None
+        self._static_pp_proxy_tensors = (
+            PPProxyTensors(self.buffers.pp_proxy_tensors)
+            if self.buffers.pp_proxy_tensors is not None
+            else None
+        )
         # TcPiecewise does its compile pass during backend construction.
         # Wrap only that path with the prefill CUDA graph failure hint.
         try:
@@ -694,11 +705,22 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
             if self._uses_eager_prefill_tail():
                 # BCG / Full: capture the transformer body only.
                 positions = self._get_layer_model_positions(forward_batch)
+                input_ids = forward_batch.input_ids
+                input_embeds = forward_batch.input_embeds
+                layer_kwargs = {}
+                if self._static_pp_proxy_tensors is not None:
+                    layer_kwargs["pp_proxy_tensors"] = self._static_pp_proxy_tensors[
+                        :num_tokens
+                    ]
+                    if not self.model_runner.pp_group.is_first_rank:
+                        input_ids = None
+                        input_embeds = None
                 return self.layer_model.forward(
-                    forward_batch.input_ids,
+                    input_ids,
                     positions,
                     forward_batch,
-                    forward_batch.input_embeds,
+                    input_embeds,
+                    **layer_kwargs,
                 )
             # tc_piecewise: compile/capture the outer model.forward path.
             return self.model_runner.model.forward(
@@ -1494,6 +1516,7 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
             padded_bs=bs,
             raw_num_tokens=num_tokens,
             padded_num_tokens=static_num_tokens,
+            pp_proxy_tensors=kwargs.get("pp_proxy_tensors"),
         )
 
         registry = self.buffer_registry
@@ -1788,9 +1811,7 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
         if isinstance(output, EmbeddingPoolerOutput):
             return output
         assert isinstance(output, PPProxyTensors)
-        raise NotImplementedError(
+        return output[: self.raw_num_tokens]
-            "PPProxyTensors is not supported in PrefillCudaGraphRunner yet."
-        )
 
     def _validate_capture_hidden_mode(self, forward_batch: ForwardBatch) -> None:
         if self.capture_hidden_mode < forward_batch.capture_hidden_mode:

python/sglang/srt/model_executor/runner_utils/buffers.py

@@ -342,6 +342,7 @@ class PrefillInputBuffers(ForwardInputBuffers):
     positions: torch.Tensor
     input_embeds: Optional[torch.Tensor]
     mrope_positions: Optional[torch.Tensor]
+    pp_proxy_tensors: Optional[Dict[str, torch.Tensor]]
 
     @classmethod
     def create(
@@ -355,6 +356,10 @@ class PrefillInputBuffers(ForwardInputBuffers):
         hidden_size: int,
         dtype: torch.dtype,
         enable_mamba_track: bool,
+        pp_size: int,
+        hc_hidden_size: Optional[int] = None,
+        pp_proxy_topk_size: Optional[int] = None,
+        pp_proxy_residual_num_blocks: Optional[int] = None,
     ) -> PrefillInputBuffers:
         with torch.device(device):
             input_ids = torch.zeros((max_num_tokens,), dtype=torch.int64)
@@ -382,6 +387,30 @@ class PrefillInputBuffers(ForwardInputBuffers):
                 input_embeds = None
                 mrope_positions = None
 
+            if pp_size > 1:
+                is_mhc = hc_hidden_size is not None
+                pp_hidden_size = hc_hidden_size if is_mhc else hidden_size
+                pp_proxy_tensors = {
+                    "hidden_states": torch.zeros(
+                        (max_num_tokens, pp_hidden_size), dtype=dtype
+                    )
+                }
+                if not is_mhc:
+                    residual_shape = (
+                        (max_num_tokens, pp_proxy_residual_num_blocks, hidden_size)
+                        if pp_proxy_residual_num_blocks is not None
+                        else (max_num_tokens, hidden_size)
+                    )
+                    pp_proxy_tensors["residual"] = torch.zeros(
+                        residual_shape, dtype=dtype
+                    )
+                if pp_proxy_topk_size is not None:
+                    pp_proxy_tensors["topk_indices"] = torch.zeros(
+                        (max_num_tokens, pp_proxy_topk_size), dtype=torch.int32
+                    )
+            else:
+                pp_proxy_tensors = None
+
         return cls(
             input_ids=input_ids,
             out_cache_loc=out_cache_loc,
@@ -392,6 +421,7 @@ class PrefillInputBuffers(ForwardInputBuffers):
             positions=positions,
             input_embeds=input_embeds,
             mrope_positions=mrope_positions,
+            pp_proxy_tensors=pp_proxy_tensors,
         )
 
     def populate_from_forward_batch(

test/registered/cuda_graph/full_prefill/test_full_cuda_graph_prefill.py

@@ -18,6 +18,7 @@ import requests
 
 from sglang.srt.utils import get_device_sm, kill_process_tree
 from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.mock_model.utils import run_mock_model_bench_serving
 from sglang.test.run_eval import run_eval
 from sglang.test.test_utils import (
     DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
@@ -27,8 +28,9 @@ from sglang.test.test_utils import (
     popen_launch_server,
 )
 
-# OSS FA4 coverage requires Blackwell. Each test still uses only one GPU.
+# OSS FA4 coverage requires Blackwell. The PP test uses two GPUs; the other
-register_cuda_ci(est_time=170, stage="base-b", runner_config="4-gpu-b200")
+# tests use one GPU.
+register_cuda_ci(est_time=240, stage="base-b", runner_config="4-gpu-b200")
 
 
 class TestFullCudaGraphPrefill(CustomTestCase):
@@ -68,6 +70,29 @@ class TestFullCudaGraphPrefill(CustomTestCase):
         self.assertGreaterEqual(score, 0.80)
 
 
+class TestFullCudaGraphPipelineParallel(CustomTestCase):
+    def test_pp_replays_full_prefill_cuda_graph(self) -> None:
+        result = run_mock_model_bench_serving(
+            extra_server_args=[
+                "--pp-size",
+                "2",
+                "--attention-backend",
+                "flashinfer",
+                "--cuda-graph-config",
+                '{"prefill":{"backend":"full","bs":[32,64],"max_bs":64,'
+                '"full_prefill_max_req":1}}',
+            ],
+            num_prompts=1,
+            random_input_len=47,
+            random_output_len=2,
+        )
+        self.assertRegex(
+            result.log_text,
+            r"Prefill batch.*cuda graph: True",
+            "The PP request did not replay a full prefill CUDA graph.",
+        )
+
+
 @unittest.skipIf(get_device_sm() < 100, "Test requires CUDA SM 100 or higher")
 class TestFullCudaGraphChunkedPrefix(unittest.TestCase):
     """A radix-cache hit replays the OSS FA4 FullCG prefix variant."""

test/registered/unit/model_executor/model_runner_components/test_cuda_graph_setup.py

@@ -5,6 +5,7 @@ import pytest
 
 from sglang.srt.model_executor.model_runner_components import cuda_graph_setup
 from sglang.srt.model_executor.model_runner_components.cuda_graph_setup import (
+    _align_pipeline_layers,
     capture_decode_graph,
 )
 from sglang.test.ci.ci_register import register_cpu_ci
@@ -56,5 +57,30 @@ def test_model_runner_can_override_decode_graph_runner(monkeypatch):
         override.restore()
 
 
+def test_align_pipeline_layers_uses_absolute_indices():
+    class PipelineStage:
+        start_layer = 3
+        end_layer = 5
+        layers = [object()] * 8
+
+    local_layers = ["layer-3", "layer-4"]
+    assert _align_pipeline_layers(local_layers, PipelineStage()) == [
+        None,
+        None,
+        None,
+        "layer-3",
+        "layer-4",
+        None,
+        None,
+        None,
+    ]
+    full_model = SimpleNamespace(layers=local_layers)
+    assert _align_pipeline_layers(local_layers, full_model) == local_layers
+    with pytest.raises(AssertionError, match="together"):
+        _align_pipeline_layers(
+            local_layers, SimpleNamespace(start_layer=0, layers=local_layers)
+        )
+
+
 if __name__ == "__main__":
     sys.exit(pytest.main([__file__, "-v"]))

test/registered/unit/model_executor/test_cuda_graph_buffer_registry.py

@@ -920,6 +920,7 @@ class TestBuildDecodeRegistry(unittest.TestCase):
     def test_source_with_pp_registers_proxy_slots(self):
         from sglang.srt.model_executor.cuda_graph_buffer_registry import (
             build_decode_registry,
+            build_prefill_registry,
         )
 
         hs = torch.zeros((8, 2), dtype=torch.int32)
@@ -964,6 +965,25 @@ class TestBuildDecodeRegistry(unittest.TestCase):
         self.assertTrue(torch.all(hs[:3] == 1))
         self.assertTrue(torch.all(hs[3:] == 0))  # tail untouched
 
+        hs.fill_(2)
+        reg = build_prefill_registry(
+            device=torch.device("cpu"),
+            max_bs=4,
+            max_num_token=8,
+            cache_loc_dtype=torch.int64,
+            source=src,
+        )
+        reg.fill_from(
+            fb,
+            raw_bs=3,
+            padded_bs=4,
+            raw_num_tokens=3,
+            padded_num_tokens=8,
+            pp_proxy_tensors=pp,
+        )
+        self.assertTrue(torch.all(hs[:3] == 1))
+        self.assertTrue(torch.all(hs[3:] == 0))
+
     def test_source_with_canary_registers_bs_slots(self):
         from sglang.srt.model_executor.cuda_graph_buffer_registry import (
             build_decode_registry,

test/registered/unit/model_executor/test_prefill_cuda_graph_runner_helpers.py

@@ -3,9 +3,14 @@
 import unittest
 from types import SimpleNamespace
 
+import torch
+
+from sglang.srt.model_executor.forward_batch_info import PPProxyTensors
 from sglang.srt.model_executor.runner.prefill_cuda_graph_runner import (
+    PrefillCudaGraphRunner,
     _resolve_transformer_layer_model,
 )
+from sglang.srt.model_executor.runner_utils.buffers import PrefillInputBuffers
 from sglang.srt.model_loader.utils import resolve_language_model
 from sglang.test.ci.ci_register import register_cpu_ci
 from sglang.test.test_utils import CustomTestCase
@@ -58,6 +63,36 @@ class TestPrefillCudaGraphRunnerHelpers(CustomTestCase):
         with self.assertRaises(AttributeError):
             resolve_language_model(SimpleNamespace())
 
+    def test_prefill_buffers_allocate_pipeline_proxy_token_rows(self):
+        buffers = PrefillInputBuffers.create(
+            device=torch.device("cpu"),
+            max_bs=4,
+            max_num_tokens=16,
+            cache_loc_dtype=torch.int64,
+            is_multimodal=False,
+            hidden_size=8,
+            dtype=torch.bfloat16,
+            enable_mamba_track=False,
+            pp_size=2,
+            pp_proxy_residual_num_blocks=3,
+        )
+
+        self.assertEqual(
+            {
+                key: tuple(value.shape)
+                for key, value in buffers.pp_proxy_tensors.items()
+            },
+            {"hidden_states": (16, 8), "residual": (16, 3, 8)},
+        )
+
+    def test_pipeline_proxy_output_is_supported(self):
+        runner = PrefillCudaGraphRunner.__new__(PrefillCudaGraphRunner)
+        runner.raw_num_tokens = 3
+        output = PPProxyTensors({"hidden_states": torch.zeros((8, 8))})
+
+        finalized = runner._finalize_execute_output(output)
+        self.assertEqual(finalized["hidden_states"].shape, (3, 8))
+
 
 if __name__ == "__main__":
     unittest.main()