minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 13 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
python/sglang/srt/arg_groups/cuda_graph_hook.py
+18
View File
@@ -115,6 +115,24 @@ def apply_cuda_graph_compatibility(server_args: Any):
if (Phase.PREFILL, "backend") in server_args._cuda_graph_config_locked:
return
# PP prefill graph replay is opt-in. It is most useful for small
# aggregate forwards, while enabling it implicitly would also capture
# large buckets that can be slower than eager. An explicit backend
# selection bypasses this default policy.
if cfg.pp_size > 1 and cfg.cuda_graph_config.prefill.backend == Backend.BREAKABLE:
logger.info(
"Disabling breakable prefill CUDA graph by default for pipeline "
"parallelism. Set --cuda-graph-backend-prefill=breakable to opt in."
)
declare_resolution(
server_args,
"_apply_cuda_graph_compatibility",
cuda_graph_config=with_phase(
cfg.cuda_graph_config, Phase.PREFILL, backend=Backend.DISABLED
),
)
return
# Breakable is the CUDA default but not multimodal-compatible;
# piecewise-allowlisted archs run their validated decoder prefill
# there instead. Archs also on the breakable allowlist keep it --
python/sglang/srt/arg_groups/memory_hook.py
+15 -1
View File
@@ -16,10 +16,12 @@ from sglang.srt.arg_groups.overrides import (
use_mla_backend,
)
from sglang.srt.environ import envs
from sglang.srt.model_executor.cuda_graph_config import Backend
from sglang.srt.model_executor.cuda_graph_config import Backend, Phase
logger = logging.getLogger(__name__)
_DEFAULT_PP_PREFILL_CUDA_GRAPH_MAX_TOKENS = 8192
def handle_gpu_memory_settings(server_args: Any, gpu_mem):
"""
@@ -189,6 +191,18 @@ def handle_gpu_memory_settings(server_args: Any, gpu_mem):
else:
prefill_cuda_graph_config.max_bs = 2048
# For opt-in PP breakable graphs, capture small aggregate-token
# buckets by default and leave larger forwards on the eager path.
# Explicit max_bs or bs settings retain their existing semantics.
if (
cfg.pp_size > 1
and prefill_cuda_graph_config.backend == Backend.BREAKABLE
and (Phase.PREFILL, "bs") not in server_args._cuda_graph_config_locked
and prefill_cuda_graph_config.max_bs
> _DEFAULT_PP_PREFILL_CUDA_GRAPH_MAX_TOKENS
):
prefill_cuda_graph_config.max_bs = _DEFAULT_PP_PREFILL_CUDA_GRAPH_MAX_TOKENS
# If max_total_tokens is set, cap prefill max_bs to not exceed max_total_tokens.
if cfg.max_total_tokens is not None:
prefill_cuda_graph_config.max_bs = min(
python/sglang/srt/layers/cp/bcg.py
+1
View File
@@ -53,6 +53,7 @@ def supports_prefill_cp_bcg(server_args: ServerArgs) -> bool:
prefill_attention_backend, _ = attention_backends_of(resolved_view(server_args))
return (
cfg.enable_prefill_cp
and cfg.pp_size == 1
and resolved.attn_cp_size == cfg.tp_size
and cfg.cp_strategy == "zigzag"
and prefill_attention_backend == "trtllm_mha"
python/sglang/srt/model_executor/cuda_graph_buffer_registry.py
+4 -1
View File
@@ -953,7 +953,10 @@ def build_prefill_registry(
def _pp_source(key):
def _fn(_fb, ctx):
ppx = ctx.pp_proxy_tensors
return None if ppx is None else ppx.tensors[key]
# Proxy contracts vary by model. The capture buffers are a
# stable-address superset; only copy fields present in the
# live proxy for this model.
return None if ppx is None else ppx.tensors.get(key)
return _fn
python/sglang/srt/model_executor/model_runner_components/cuda_graph_setup.py
+14
View File
@@ -45,6 +45,7 @@ from sglang.srt.runtime_context import (
get_disagg,
get_exec,
get_flags,
get_parallel,
get_schedule,
get_spec,
)
@@ -72,6 +73,8 @@ def _align_pipeline_layers(layers: list, layer_model) -> list:
f"invalid pipeline layer range [{start_layer}, {end_layer}) for "
f"{len(layer_model.layers)} layers"
)
if len(layers) == len(layer_model.layers):
return layers
assert (
len(layers) <= end_layer - start_layer
), f"found {len(layers)} layers in PP range [{start_layer}, {end_layer})"
@@ -365,6 +368,17 @@ def capture_prefill_graph(
prefill_config = get_exec().graph.cuda_graph_config.prefill
prefill_backend = prefill_config.backend
parallel = get_parallel()
if (
prefill_backend == Backend.BREAKABLE
and parallel.enable_prefill_cp
and parallel.pp_size > 1
):
logger.warning(
"Disable prefill CUDA graph because pipeline parallelism combined "
"with prefill context parallelism is not validated."
)
return result(eager_runner)
context_length = model_runner.model_config.context_len
if prefill_backend == Backend.FULL:
max_capture_requests = prefill_config.full_prefill_max_req
python/sglang/srt/model_executor/model_runner_components/layer_setup.py
+5 -6
View File
@@ -62,12 +62,11 @@ def compute_attention_and_moe_layers(layer_model: Any) -> AttentionAndMoeLayers:
# Mamba layer with split op support - store the layer itself
attn_layer = layer
if attn_layer is not None:
attention_layers.append(attn_layer)
mha_companion_layers.append(mha_companion_layer)
elif hasattr(layer, "mixer"):
attention_layers.append(None)
mha_companion_layers.append(None)
# Keep these lists aligned with global layer ids. Pipeline-parallel
# models retain placeholders outside the local stage, while real
# attention modules use their global layer_id during graph replay.
attention_layers.append(attn_layer)
mha_companion_layers.append(mha_companion_layer)
moe_block = None
moe_fusion = None
python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py
+40 -10
View File
@@ -185,6 +185,23 @@ def _resolve_transformer_layer_model(model: torch.nn.Module) -> torch.nn.Module:
return layer_model
def _build_layer_model_forward_kwargs(
layer_model: torch.nn.Module,
forward_batch: ForwardBatch,
pp_proxy_tensors: Optional[PPProxyTensors],
) -> Dict[str, Any]:
"""Bind optional transformer inputs by name across model signatures."""
parameters = inspect.signature(layer_model.forward).parameters
kwargs = {}
for embeds_name in ("input_embeds", "inputs_embeds"):
if embeds_name in parameters:
kwargs[embeds_name] = forward_batch.input_embeds
break
if pp_proxy_tensors is not None and "pp_proxy_tensors" in parameters:
kwargs["pp_proxy_tensors"] = pp_proxy_tensors
return kwargs
def _slice_output_rows(output: Any, num_tokens: int) -> Any:
"""Slice every tensor leaf in a transformer-body output by token rows.
@@ -333,6 +350,7 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
dtype=self.model_runner.dtype,
enable_mamba_track=self.mamba_track_enabled,
pp_size=self.pp_size,
is_first_pp_rank=self.model_runner.pp_group.is_first_rank,
hc_hidden_size=getattr(
self.model_runner.model_config, "hc_hidden_size", None
),
@@ -653,12 +671,12 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
return forward_batch.positions
def _static_pp_proxy_tensors(self, num_tokens: int) -> Optional[PPProxyTensors]:
def _capture_pp_proxy_tensors(self, num_tokens: int) -> Optional[PPProxyTensors]:
buffers = self.buffers.pp_proxy_tensors
if buffers is None:
if buffers is None or self.model_runner.pp_group.is_first_rank:
return None
return PPProxyTensors(
{key: value[:num_tokens] for key, value in buffers.items()}
{name: buffer[:num_tokens] for name, buffer in buffers.items()}
)
@contextmanager
@@ -717,20 +735,28 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
set_is_extend_in_batch(False)
with self._prefill_forward_context(forward_batch):
pp_kwargs = self.model_runner._pp_kwargs(
self._static_pp_proxy_tensors(num_tokens)
)
pp_proxy_tensors = self._capture_pp_proxy_tensors(num_tokens)
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
kwargs = _build_layer_model_forward_kwargs(
self.layer_model, forward_batch, pp_proxy_tensors
)
if pp_proxy_tensors is not None:
input_ids = None
for embeds_name in ("input_embeds", "inputs_embeds"):
if embeds_name in kwargs:
kwargs[embeds_name] = None
break
return self.layer_model.forward(
forward_batch.input_ids,
input_ids,
positions,
forward_batch,
forward_batch.input_embeds,
**pp_kwargs,
**kwargs,
)
# tc_piecewise: compile/capture the outer model.forward path.
pp_kwargs = self.model_runner._pp_kwargs(pp_proxy_tensors)
return self.model_runner.model.forward(
forward_batch.input_ids,
forward_batch.positions,
@@ -1765,6 +1791,10 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
# uses real request metadata instead of padded slots. BCG has no
# request-slot padding, so static_forward_batch is already the serving batch.
tail_batch = forward_batch if full_path else static_forward_batch
if not full_path:
# MTP consumes the target model's live multimodal embeddings in its
# eager wrapper before the captured transformer body is replayed.
tail_batch.mm_input_embeds = forward_batch.mm_input_embeds
try:
with self._prefill_forward_context(
static_forward_batch,
@@ -1836,7 +1866,7 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
if isinstance(output, EmbeddingPoolerOutput):
return output
assert isinstance(output, PPProxyTensors)
return output[: self.raw_num_tokens]
return _slice_output_rows(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
+59 -46
View File
@@ -60,6 +60,38 @@ def _grouped_foreach_copy_(dsts: List[torch.Tensor], srcs: List[torch.Tensor]) -
foreach_copy(group_dsts, group_srcs)
def _allocate_pp_proxy_tensors(
*,
max_num_tokens: int,
max_hidden_tokens: int,
hidden_size: int,
dtype: torch.dtype,
hc_hidden_size: Optional[int] = None,
pp_proxy_topk_size: Optional[int] = None,
pp_proxy_residual_num_blocks: Optional[int] = None,
) -> Dict[str, torch.Tensor]:
"""Allocate the stable buffers consumed by an incoming PP proxy."""
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_hidden_tokens, pp_hidden_size), dtype=dtype),
}
if not is_mhc:
# Only Kimi K3 supplies num_blocks: its PP bank is token-major
# [T, blocks, H]. Other models use the phase-specific hidden-token bound.
residual_shape = (
(max_num_tokens, pp_proxy_residual_num_blocks, hidden_size)
if pp_proxy_residual_num_blocks is not None
else (max_hidden_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
)
return pp_proxy_tensors
@dataclass
class DecodeInputBuffers(ForwardInputBuffers):
input_ids: torch.Tensor
@@ -129,29 +161,19 @@ class DecodeInputBuffers(ForwardInputBuffers):
torch.zeros((max_bs,), dtype=torch.bool) if enable_mamba_track else None
)
if pp_size > 1:
is_mhc = hc_hidden_size is not None
hs = hc_hidden_size if is_mhc else hidden_size
pp_proxy_tensors = {
"hidden_states": torch.zeros((max_num_token, hs), dtype=dtype),
}
if not is_mhc:
# Only Kimi K3 supplies num_blocks: its PP bank is token-major
# [T, blocks, H]. Other models keep the legacy [max_bs, H].
residual_shape = (
(max_num_token, pp_proxy_residual_num_blocks, hidden_size)
if pp_proxy_residual_num_blocks is not None
else (max_num_token, 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_token, pp_proxy_topk_size), dtype=torch.int32
)
else:
pp_proxy_tensors = None
pp_proxy_tensors = (
_allocate_pp_proxy_tensors(
max_num_tokens=max_num_token,
max_hidden_tokens=max_num_token,
hidden_size=hidden_size,
dtype=dtype,
hc_hidden_size=hc_hidden_size,
pp_proxy_topk_size=pp_proxy_topk_size,
pp_proxy_residual_num_blocks=pp_proxy_residual_num_blocks,
)
if pp_size > 1
else None
)
if is_encoder_decoder:
encoder_lens = torch.full(
@@ -357,6 +379,7 @@ class PrefillInputBuffers(ForwardInputBuffers):
dtype: torch.dtype,
enable_mamba_track: bool,
pp_size: int = 1,
is_first_pp_rank: bool = False,
hc_hidden_size: Optional[int] = None,
pp_proxy_topk_size: Optional[int] = None,
pp_proxy_residual_num_blocks: Optional[int] = None,
@@ -387,29 +410,19 @@ 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
pp_proxy_tensors = (
_allocate_pp_proxy_tensors(
max_num_tokens=max_num_tokens,
max_hidden_tokens=max_num_tokens,
hidden_size=hidden_size,
dtype=dtype,
hc_hidden_size=hc_hidden_size,
pp_proxy_topk_size=pp_proxy_topk_size,
pp_proxy_residual_num_blocks=pp_proxy_residual_num_blocks,
)
if pp_size > 1 and not is_first_pp_rank
else None
)
return cls(
input_ids=input_ids,