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

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

Browse Source 
Co-authored-by: Yuwei An <ayw.sirius19@gmail.com>
This commit is contained in:
Aurick Qiao
2026-08-27 17:32:00 -07:00
committed by GitHub
co-authored by Yuwei An
parent 7cbe564829
commit 26fd7fdaa2
8 changed files with 218 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/model_executor/cuda_graph_buffer_registry.py
+24
View File
@@ -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
+30
View File
@@ -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
+26 -5
View File
@@ -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:
python/sglang/srt/model_executor/runner_utils/buffers.py
+30
View File
@@ -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(