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

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

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()