From 26fd7fdaa2732abbad6d63b21cf0944aa88e977e Mon Sep 17 00:00:00 2001 From: Aurick Qiao Date: Thu, 27 Aug 2026 17:32:00 -0700 Subject: [PATCH] [Feature] Support PP in full prefill CUDA graphs (#35451) Co-authored-by: Yuwei An --- .../cuda_graph_buffer_registry.py | 24 +++++++++++++ .../cuda_graph_setup.py | 30 ++++++++++++++++ .../runner/prefill_cuda_graph_runner.py | 31 +++++++++++++--- .../model_executor/runner_utils/buffers.py | 30 ++++++++++++++++ .../test_full_cuda_graph_prefill.py | 29 +++++++++++++-- .../test_cuda_graph_setup.py | 26 ++++++++++++++ .../test_cuda_graph_buffer_registry.py | 20 +++++++++++ .../test_prefill_cuda_graph_runner_helpers.py | 35 +++++++++++++++++++ 8 files changed, 218 insertions(+), 7 deletions(-) diff --git a/python/sglang/srt/model_executor/cuda_graph_buffer_registry.py b/python/sglang/srt/model_executor/cuda_graph_buffer_registry.py index 5e60753b4..986e3545c 100644 --- a/python/sglang/srt/model_executor/cuda_graph_buffer_registry.py +++ b/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 diff --git a/python/sglang/srt/model_executor/model_runner_components/cuda_graph_setup.py b/python/sglang/srt/model_executor/model_runner_components/cuda_graph_setup.py index aecb95896..9914e6a63 100644 --- a/python/sglang/srt/model_executor/model_runner_components/cuda_graph_setup.py +++ b/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" diff --git a/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py b/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py index 0eb3ff600..1eb7d010e 100644 --- a/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py +++ b/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( - "PPProxyTensors is not supported in PrefillCudaGraphRunner yet." - ) + return output[: self.raw_num_tokens] def _validate_capture_hidden_mode(self, forward_batch: ForwardBatch) -> None: if self.capture_hidden_mode < forward_batch.capture_hidden_mode: diff --git a/python/sglang/srt/model_executor/runner_utils/buffers.py b/python/sglang/srt/model_executor/runner_utils/buffers.py index 23db4547b..5ac57f2f1 100644 --- a/python/sglang/srt/model_executor/runner_utils/buffers.py +++ b/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( diff --git a/test/registered/cuda_graph/full_prefill/test_full_cuda_graph_prefill.py b/test/registered/cuda_graph/full_prefill/test_full_cuda_graph_prefill.py index be53446c1..9001da960 100644 --- a/test/registered/cuda_graph/full_prefill/test_full_cuda_graph_prefill.py +++ b/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. -register_cuda_ci(est_time=170, stage="base-b", runner_config="4-gpu-b200") +# OSS FA4 coverage requires Blackwell. The PP test uses two GPUs; the other +# 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.""" diff --git a/test/registered/unit/model_executor/model_runner_components/test_cuda_graph_setup.py b/test/registered/unit/model_executor/model_runner_components/test_cuda_graph_setup.py index aec512936..6a054a331 100644 --- a/test/registered/unit/model_executor/model_runner_components/test_cuda_graph_setup.py +++ b/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"])) diff --git a/test/registered/unit/model_executor/test_cuda_graph_buffer_registry.py b/test/registered/unit/model_executor/test_cuda_graph_buffer_registry.py index 3ef7f5a10..214cb3f56 100644 --- a/test/registered/unit/model_executor/test_cuda_graph_buffer_registry.py +++ b/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, diff --git a/test/registered/unit/model_executor/test_prefill_cuda_graph_runner_helpers.py b/test/registered/unit/model_executor/test_prefill_cuda_graph_runner_helpers.py index b01de41d1..70c406d18 100644 --- a/test/registered/unit/model_executor/test_prefill_cuda_graph_runner_helpers.py +++ b/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()