minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[PP] Support prefill CUDA graph proxy tensors (#36248)

Browse Source 
Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
This commit is contained in:
YAMY
2026-08-30 23:11:45 -07:00
committed by GitHub
co-authored by github-actions[bot]
parent 5d92e60783
commit b77cac06a9
12 changed files with 353 additions and 64 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/unit/model_executor/model_runner_components/test_layer_setup.py
+14
View File
@@ -31,6 +31,20 @@ class TestComputeAttentionAndMoeLayers(unittest.TestCase):
self.assertEqual(mha_companion_layers, [attn_mha])
self.assertNotIn("_pcg_mha_companion", vars(attn_mqa))
def test_pipeline_placeholders_preserve_global_layer_ids(self):
local_attention = SimpleNamespace()
layer_model = SimpleNamespace(
layers=[SimpleNamespace(), SimpleNamespace()]
+ [SimpleNamespace(self_attn=SimpleNamespace(attn=local_attention))]
)
attention_layers, _, _, _, mha_companion_layers = (
compute_attention_and_moe_layers(layer_model)
)
self.assertEqual(attention_layers, [None, None, local_attention])
self.assertEqual(mha_companion_layers, [None, None, None])
if __name__ == "__main__":
unittest.main()
test/registered/unit/model_executor/test_prefill_cuda_graph_runner.py
+2
View File
@@ -100,6 +100,8 @@ class TestPrefillCudaGraphRunnerChunkedPrefix(CustomTestCase):
# out of the bags.
override = get_context().override_server_args(
enable_lora=False,
enable_prefill_cp=False,
pp_size=1,
cuda_graph_config=SimpleNamespace(
prefill=SimpleNamespace(bs=[1], backend=Backend.BREAKABLE)
),
test/registered/unit/model_executor/test_prefill_cuda_graph_runner_helpers.py
+127
View File
@@ -1,13 +1,19 @@
"""Unit tests for prefill CUDA graph wrapper helpers."""
import unittest
from contextlib import nullcontext
from functools import partial
from types import SimpleNamespace
import torch
from sglang.srt.model_executor.cuda_graph_buffer_registry import (
build_prefill_registry,
)
from sglang.srt.model_executor.forward_batch_info import PPProxyTensors
from sglang.srt.model_executor.runner.prefill_cuda_graph_runner import (
PrefillCudaGraphRunner,
_build_layer_model_forward_kwargs,
_resolve_transformer_layer_model,
)
from sglang.srt.model_executor.runner_utils.buffers import PrefillInputBuffers
@@ -26,7 +32,95 @@ class _LayerModel:
return input_embeds
def _make_pp_buffers_and_registry():
base = torch.zeros(3, dtype=torch.int64)
buffers = SimpleNamespace(
**{name: base.clone() for name in ("input_ids", "positions", "out_cache_loc")},
pp_proxy_tensors={
key: torch.zeros((3, 2)) for key in ("hidden_states", "residual")
},
)
registry = build_prefill_registry(
device=base.device,
max_bs=1,
max_num_token=len(base),
cache_loc_dtype=torch.int64,
share_pool=False,
source=buffers,
)
return buffers, registry
class TestPrefillCudaGraphRunnerHelpers(CustomTestCase):
def test_pp_proxy_stable_buffers_accept_full_and_hidden_only_contracts(self):
buffers, registry = _make_pp_buffers_and_registry()
full_proxy = PPProxyTensors(
{
"hidden_states": torch.full((3, 2), 2.0),
"residual": torch.full((3, 2), 3.0),
}
)
values = torch.arange(3)
fill = partial(
registry.fill_from,
SimpleNamespace(input_ids=values, positions=values, out_cache_loc=values),
raw_bs=1,
padded_bs=1,
raw_num_tokens=3,
padded_num_tokens=3,
)
fill(pp_proxy_tensors=full_proxy)
runner = PrefillCudaGraphRunner.__new__(PrefillCudaGraphRunner)
runner.buffers = buffers
runner.model_runner = SimpleNamespace(
pp_group=SimpleNamespace(is_first_rank=False)
)
capture_proxy = runner._capture_pp_proxy_tensors(3)
torch.testing.assert_close(capture_proxy.tensors, full_proxy.tensors)
self.assertEqual(
capture_proxy["hidden_states"].data_ptr(),
buffers.pp_proxy_tensors["hidden_states"].data_ptr(),
)
hidden_only_proxy = PPProxyTensors({"hidden_states": torch.full((3, 2), 4.0)})
fill(pp_proxy_tensors=hidden_only_proxy)
torch.testing.assert_close(
buffers.pp_proxy_tensors["hidden_states"][:3],
hidden_only_proxy["hidden_states"],
)
def test_layer_model_kwargs_bind_optional_inputs_by_signature(self):
def proxy_before_embeds(a, b, c, pp_proxy_tensors=None, inputs_embeds=None):
pass
cases = (
(_LayerModel(), {"input_embeds": "embeds"}),
(
SimpleNamespace(forward=proxy_before_embeds),
{"inputs_embeds": "embeds", "pp_proxy_tensors": "proxy"},
),
)
forward_batch = SimpleNamespace(input_embeds="embeds")
for layer_model, expected in cases:
with self.subTest(signature=layer_model.forward.__name__):
kwargs = _build_layer_model_forward_kwargs(
layer_model, forward_batch, "proxy"
)
self.assertEqual(kwargs, expected)
layer_model.forward(None, None, forward_batch, **kwargs)
def test_finalize_pp_proxy_trims_padded_token_rows(self):
runner = PrefillCudaGraphRunner.__new__(PrefillCudaGraphRunner)
runner.raw_num_tokens = 3
output = PPProxyTensors({"hidden_states": torch.arange(10).reshape(5, 2)})
trimmed = runner._finalize_execute_output(output)
self.assertIsInstance(trimmed, PPProxyTensors)
self.assertEqual(tuple(trimmed["hidden_states"].shape), (3, 2))
torch.testing.assert_close(
trimmed["hidden_states"][-1], output["hidden_states"][2]
)
def test_resolve_layer_model_from_language_model_wrapper(self):
layer_model = _LayerModel()
model = SimpleNamespace(language_model=SimpleNamespace(model=layer_model))
@@ -74,6 +168,7 @@ class TestPrefillCudaGraphRunnerHelpers(CustomTestCase):
dtype=torch.bfloat16,
enable_mamba_track=False,
pp_size=2,
is_first_pp_rank=False,
pp_proxy_residual_num_blocks=3,
)
@@ -93,6 +188,38 @@ class TestPrefillCudaGraphRunnerHelpers(CustomTestCase):
finalized = runner._finalize_execute_output(output)
self.assertEqual(finalized["hidden_states"].shape, (3, 8))
def test_bcg_eager_tail_uses_live_multimodal_embeddings(self):
live_embeds = object()
live_batch = SimpleNamespace(mm_input_embeds=live_embeds)
static_batch = SimpleNamespace(
input_ids=None,
positions=None,
mm_input_embeds=None,
)
runner = PrefillCudaGraphRunner.__new__(PrefillCudaGraphRunner)
runner._is_full_backend = False
runner._input_embeds_arg_idx = None
runner.buffer_registry = SimpleNamespace(has_slot=lambda _name: False)
runner.backend = SimpleNamespace(replay=lambda *_args, **_kwargs: None)
runner.layer_model = SimpleNamespace(forward=lambda *_args, **_kwargs: None)
runner.model_runner = SimpleNamespace(
model=SimpleNamespace(
forward=lambda _ids, _positions, batch, **_kwargs: batch.mm_input_embeds
)
)
runner._prefill_forward_context = lambda *_args, **_kwargs: nullcontext()
output = runner._execute_body_capture(
live_batch,
static_batch,
static_num_tokens=1,
raw_num_tokens=1,
shape_key=object(),
)
self.assertIs(output, live_embeds)
if __name__ == "__main__":
unittest.main()
test/registered/unit/server_args/test_server_args.py
+54
View File
@@ -15,6 +15,7 @@ from sglang.srt.arg_groups.attention_hook import (
handle_deterministic_inference,
)
from sglang.srt.arg_groups.cuda_graph_hook import (
apply_cuda_graph_compatibility,
disable_tc_piecewise_cudagraph_if_incompatible,
handle_cuda_graph_config,
)
@@ -31,6 +32,7 @@ from sglang.srt.arg_groups.kv_cache_hook import (
validate_prefill_only_disable_kv_cache_args,
)
from sglang.srt.arg_groups.mamba_hook import handle_mamba_backend
from sglang.srt.arg_groups.memory_hook import handle_gpu_memory_settings
from sglang.srt.arg_groups.model_path_hook import handle_load_format
from sglang.srt.arg_groups.moe_hook import (
handle_a2a_moe,
@@ -1852,6 +1854,58 @@ class TestCudaGraphConfigDataclassAccess(CustomTestCase):
self.assertEqual(config.compiler, "eager")
class TestPipelineParallelPrefillCudaGraphPolicy(CustomTestCase):
def test_pp_prefill_graph_is_opt_in(self):
cases = (
(set(), Backend.DISABLED),
({(Phase.PREFILL, "backend")}, Backend.BREAKABLE),
)
for locked, expected in cases:
with self.subTest(locked=locked):
args = ServerArgs(
model_path="dummy",
pp_size=4,
cuda_graph_config=CudaGraphConfig(
prefill=PhaseConfig(backend=Backend.BREAKABLE)
),
)
args._cuda_graph_config_locked = locked
apply_cuda_graph_compatibility(args)
self.assertEqual(
resolution_result(args, "cuda_graph_config").prefill.backend,
expected,
)
def test_pp_prefill_capture_limit_policy(self):
cases = (
(4096, None, 4096),
(32768, None, 8192),
(32768, 16384, 16384),
)
for chunked_prefill_size, max_bs, expected in cases:
with self.subTest(chunked_prefill_size=chunked_prefill_size, max_bs=max_bs):
args = ServerArgs(
model_path="dummy",
pp_size=4,
chunked_prefill_size=chunked_prefill_size,
mem_fraction_static=0.8,
cuda_graph_config=CudaGraphConfig(
decode=PhaseConfig(backend=Backend.DISABLED, max_bs=1, bs=[1]),
prefill=PhaseConfig(backend=Backend.BREAKABLE, max_bs=max_bs),
),
)
args._cuda_graph_config_locked = {(Phase.PREFILL, "backend")} | (
{(Phase.PREFILL, "max_bs")} if max_bs is not None else set()
)
with patch(
"sglang.srt.arg_groups.memory_hook.use_mla_backend",
return_value=False,
):
handle_gpu_memory_settings(args, gpu_mem=None)
prefill = resolution_result(args, "cuda_graph_config").prefill
self.assertEqual((prefill.max_bs, prefill.bs[-1]), (expected, expected))
class TestCudaGraphDisaggregationRoles(CustomTestCase):
def _handled_args(self, **overrides):
args = ServerArgs(model_path="dummy", **overrides)