From 0bf0640b9d945cb9e3a96a66bb7d0ba47a72f3b5 Mon Sep 17 00:00:00 2001 From: Baizhou Zhang Date: Sun, 2 Aug 2026 23:11:57 -0700 Subject: [PATCH] [CP] Support breakable CUDA graphs for zigzag strategy (#33136) --- python/sglang/srt/layers/cp/bcg.py | 240 ++++++++++++++++++ .../sglang/srt/model_executor/model_runner.py | 15 +- .../runner/prefill_cuda_graph_runner.py | 54 +++- python/sglang/srt/server_args.py | 4 +- .../cp/test_gpt_oss_4gpu_mxfp4_cp.py | 2 + 5 files changed, 311 insertions(+), 4 deletions(-) create mode 100644 python/sglang/srt/layers/cp/bcg.py diff --git a/python/sglang/srt/layers/cp/bcg.py b/python/sglang/srt/layers/cp/bcg.py new file mode 100644 index 000000000..ed6ea4b79 --- /dev/null +++ b/python/sglang/srt/layers/cp/bcg.py @@ -0,0 +1,240 @@ +# Copyright 2023-2026 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== + +"""Breakable CUDA graph helpers for context-parallel prefill.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Any, Dict + +import torch + +from sglang.srt.layers.cp.utils import ( + cp_gather_after_forward, + cp_split_before_forward, + enable_cp_v2, + prepare_cp_forward, +) +from sglang.srt.model_executor.forward_batch_info import PPProxyTensors +from sglang.srt.model_executor.runner.shape_key import ShapeKey + +if TYPE_CHECKING: + from sglang.srt.model_executor.forward_batch_info import ForwardBatch + from sglang.srt.model_executor.runner.prefill_cuda_graph_runner import ( + PrefillCudaGraphRunner, + ) + from sglang.srt.server_args import ServerArgs + + +def supports_prefill_cp_bcg(server_args: ServerArgs) -> bool: + """Return whether the selected prefill-CP configuration supports BCG.""" + resolved = server_args._resolved() + prefill_attention_backend, _ = server_args._resolved_attention_backends() + return ( + server_args.enable_prefill_cp + and resolved.attn_cp_size == server_args.tp_size + and server_args.cp_strategy == "zigzag" + and prefill_attention_backend == "trtllm_mha" + ) + + +def enable_cp_v2_bcg_capture(server_args: ServerArgs) -> bool: + """Return whether CP-v2 breakable prefill capture is enabled.""" + return enable_cp_v2() and supports_prefill_cp_bcg(server_args) + + +def filter_prefill_cp_bcg_capture_num_tokens( + capture_num_tokens: list[int], server_args: ServerArgs +) -> list[int]: + """Keep only token buckets where the zigzag CP strategy can run.""" + min_num_tokens = server_args._resolved().attn_cp_size * 2 + filtered = [size for size in capture_num_tokens if size >= min_num_tokens] + if not filtered: + raise ValueError( + "Prefill CP breakable CUDA graph requires at least one token bucket " + f">= {min_num_tokens}, but got {capture_num_tokens}." + ) + return filtered + + +def _slice_output_rows(output: Any, num_tokens: int) -> Any: + if output is None: + return None + if torch.is_tensor(output) or isinstance(output, PPProxyTensors): + return output[:num_tokens] + if isinstance(output, tuple): + return tuple(_slice_output_rows(item, num_tokens) for item in output) + if isinstance(output, list): + return [_slice_output_rows(item, num_tokens) for item in output] + raise TypeError(f"Unsupported prefill CP BCG output: {type(output)}") + + +@dataclass +class PrefillCPBCGInput: + """Fixed-address CP-local inputs and per-bucket replay state.""" + + input_embeds: torch.Tensor + positions: torch.Tensor + bucket_local_tokens: Dict[int, int] = field(default_factory=dict) + live_local_tokens: int = 0 + + @classmethod + def create(cls, runner: PrefillCudaGraphRunner) -> PrefillCPBCGInput: + with torch.device(runner.device): + return cls( + input_embeds=torch.zeros( + ( + runner.max_num_tokens, + runner.model_runner.model_config.hidden_size, + ), + dtype=runner.model_runner.dtype, + ), + positions=torch.zeros( + (runner.max_num_tokens,), + dtype=torch.int64, + ), + ) + + def prepare( + self, + runner: PrefillCudaGraphRunner, + forward_batch: ForwardBatch, + *, + static_num_tokens: int, + capture: bool, + ) -> None: + """Shard global prefill inputs into fixed-address CP-local buffers.""" + # Replay batches may reuse a ForwardBatch object whose metadata was + # built for a different request layout. Always rebuild before sharding. + forward_batch.attn_cp_metadata = None + prepare_cp_forward(forward_batch) + + captured_local_tokens = None + if not capture: + try: + captured_local_tokens = self.bucket_local_tokens[static_num_tokens] + except KeyError as exc: + raise RuntimeError( + "Missing CP-local capture capacity for global prefill bucket " + f"{static_num_tokens}" + ) from exc + + # Breakable graph segments retain their captured row geometry when + # a smaller live batch reuses the same global token bucket. + metadata = forward_batch.attn_cp_metadata + if hasattr(metadata, "per_rank_actual_token"): + live_physical_tokens = max(metadata.per_rank_actual_token) + if live_physical_tokens > captured_local_tokens: + raise RuntimeError( + f"Live batch needs {live_physical_tokens} local CP rows, " + f"but global prefill bucket {static_num_tokens} has " + f"captured capacity {captured_local_tokens}" + ) + cp_size = len(metadata.per_rank_actual_token) + metadata.per_rank_actual_token = [captured_local_tokens] * cp_size + metadata.max_rank_len = [captured_local_tokens] * cp_size + + raw_tokens = int(forward_batch.extend_num_tokens) + global_input_ids = forward_batch.input_ids[:raw_tokens] + global_positions = forward_batch.positions[:raw_tokens] + global_input_embeds = runner.model_runner.model.get_input_embeddings()( + global_input_ids + ) + local_input_embeds, local_positions = cp_split_before_forward( + global_input_embeds, + global_positions, + forward_batch, + ) + live_local_tokens = int(local_input_embeds.shape[0]) + + if capture: + captured_local_tokens = live_local_tokens + self.bucket_local_tokens[static_num_tokens] = captured_local_tokens + else: + assert captured_local_tokens is not None + if live_local_tokens > captured_local_tokens: + raise RuntimeError( + f"Live batch needs {live_local_tokens} local CP rows, but global " + f"prefill bucket {static_num_tokens} has captured capacity " + f"{captured_local_tokens}" + ) + + if captured_local_tokens > self.input_embeds.shape[0]: + raise RuntimeError( + f"CP-local capture needs {captured_local_tokens} rows, but the " + f"fixed input buffer has capacity {self.input_embeds.shape[0]}" + ) + + input_embeds = self.input_embeds[:captured_local_tokens] + positions = self.positions[:captured_local_tokens] + input_embeds.zero_() + positions.zero_() + input_embeds[:live_local_tokens].copy_(local_input_embeds) + positions[:live_local_tokens].copy_(local_positions) + forward_batch.input_embeds = input_embeds + forward_batch.positions = positions + self.live_local_tokens = live_local_tokens + + +def execute_prefill_cp_bcg( + runner: PrefillCudaGraphRunner, + forward_batch: ForwardBatch, + static_forward_batch: ForwardBatch, + static_num_tokens: int, + raw_num_tokens: int, + **kwargs, +): + """Replay a CP-local body and run the global gather/logits tail eagerly.""" + cp_input = runner.prefill_cp_bcg_input + assert cp_input is not None + model = runner.model_runner.model + with runner._prefill_forward_context( + static_forward_batch, + num_tokens=static_num_tokens, + raw_num_tokens=raw_num_tokens, + ): + local_output = runner.backend.replay( + ShapeKey(size=static_num_tokens), + static_forward_batch, + **kwargs, + ) + local_output = _slice_output_rows(local_output, cp_input.live_local_tokens) + + capture_aux_hidden_states = getattr(model, "capture_aux_hidden_states", False) + aux_hidden_states = None + if capture_aux_hidden_states: + hidden_states, aux_hidden_states = local_output + else: + hidden_states = local_output + + if not model.pp_group.is_last_rank: + return ( + (hidden_states, aux_hidden_states) + if capture_aux_hidden_states + else hidden_states + ) + + hidden_states = cp_gather_after_forward( + hidden_states, + static_forward_batch, + torch.cuda.current_stream(), + ) + return model.logits_processor( + forward_batch.input_ids, + hidden_states, + model.lm_head, + forward_batch, + aux_hidden_states, + ) diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 880c81b6d..d91c51328 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -76,6 +76,7 @@ from sglang.srt.layers import deep_gemm_wrapper, model_parallel from sglang.srt.layers.attention.dsa.utils import is_dsa_enable_prefill_cp from sglang.srt.layers.cp.utils import ( get_cp_strategy, + is_cp_v2_active, ) from sglang.srt.layers.logits_processor import LogitsProcessorOutput from sglang.srt.layers.sampler import create_sampler @@ -234,6 +235,16 @@ elif current_platform.is_out_of_tree(): logger = logging.getLogger(__name__) +def _prefill_cuda_graph_allows_context_parallel( + prefill_runner, forward_batch: ForwardBatch +) -> bool: + """Allow CP only through a runner that captured the validated CP-v2 body.""" + return get_cp_strategy() is None or ( + bool(getattr(prefill_runner, "enable_cp_v2_bcg_capture", False)) + and is_cp_v2_active(forward_batch) + ) + + @dataclass class ModelRunnerOutput: logits_output: Union[LogitsProcessorOutput, PPProxyTensors] @@ -1520,7 +1531,9 @@ class ModelRunner: and not isinstance(self.prefill_cuda_graph_runner, EagerRunner) and self.prefill_cuda_graph_runner is not None and self.prefill_cuda_graph_runner.can_run_graph(forward_batch) - and get_cp_strategy() is None + and _prefill_cuda_graph_allows_context_parallel( + self.prefill_cuda_graph_runner, forward_batch + ) ): # Prefill cuda graph (piecewise). kwargs = self._extend_forward_kwargs(forward_batch, pp_proxy_tensors) 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 ab21cc9dd..c334f7f27 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 @@ -53,6 +53,16 @@ from sglang.kernels.ops.kvcache.kv_indices import ( from sglang.srt.configs.model_config import is_deepseek_dsa from sglang.srt.distributed.parallel_state import graph_capture from sglang.srt.layers.attention.dsa.utils import is_dsa_enable_prefill_cp +from sglang.srt.layers.cp.bcg import ( + PrefillCPBCGInput, +) +from sglang.srt.layers.cp.bcg import ( + enable_cp_v2_bcg_capture as should_enable_cp_v2_bcg_capture, +) +from sglang.srt.layers.cp.bcg import ( + execute_prefill_cp_bcg, + filter_prefill_cp_bcg_capture_num_tokens, +) from sglang.srt.layers.dp_attention import ( DpPaddingMode, set_dp_buffer_len, @@ -350,6 +360,8 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner): self._is_full_backend = False # Same ordering requirement: capture_prepare reads this. self._capture_lora = False + self.enable_cp_v2_bcg_capture = False + self.prefill_cp_bcg_input: Optional[PrefillCPBCGInput] = None # TcPiecewise does its compile pass during backend construction. # Wrap only that path with the prefill CUDA graph failure hint. try: @@ -447,6 +459,16 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner): for name in _PREFILL_STATIC_FIELDS } + server_args = model_runner.server_args + self.enable_cp_v2_bcg_capture = isinstance( + self.backend, BreakableCudaGraphBackend + ) and should_enable_cp_v2_bcg_capture(server_args) + if self.enable_cp_v2_bcg_capture: + self.capture_num_tokens = filter_prefill_cp_bcg_capture_num_tokens( + self.capture_num_tokens, server_args + ) + self.prefill_cp_bcg_input = PrefillCPBCGInput.create(self) + # Static hidden_states buffer giving the captured graph a stable # address; load_batch refreshes it from live spec_info at replay. # Draft consumes aux-concatenated hidden states from the target @@ -1290,6 +1312,14 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner): """ num_tokens = size forward_batch, attn_backend = self.capture_prepare(num_tokens) + if self.enable_cp_v2_bcg_capture: + assert self.prefill_cp_bcg_input is not None + self.prefill_cp_bcg_input.prepare( + self, + forward_batch, + static_num_tokens=num_tokens, + capture=True, + ) if forward_batch.lora_ids is not None: # Fill the static prefill LoRA batch info the captured kernels # will read (all-None ids: ranks stay 0, kernels no-op). @@ -1526,8 +1556,19 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner): forward_batch.spec_info.hidden_states ) + metadata_forward_batch = forward_batch + if self.enable_cp_v2_bcg_capture: + assert self.prefill_cp_bcg_input is not None + self.prefill_cp_bcg_input.prepare( + self, + static_forward_batch, + static_num_tokens=static_num_tokens, + capture=False, + ) + metadata_forward_batch = static_forward_batch + self._prepare_forward_metadata_for_replay( - forward_batch, static_forward_batch, static_num_tokens + metadata_forward_batch, static_forward_batch, static_num_tokens ) return static_forward_batch @@ -1667,7 +1708,16 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner): if shape_key.variant_label is not None: self._prepare_chunked_prefix_replay(shape_key, forward_batch) - if self._uses_eager_prefill_tail(): + if self.enable_cp_v2_bcg_capture: + output = execute_prefill_cp_bcg( + self, + forward_batch, + static_forward_batch, + static_num_tokens, + raw_num_tokens, + **kwargs, + ) + elif self._uses_eager_prefill_tail(): output = self._execute_body_capture( forward_batch, static_forward_batch, diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 8e57077eb..32aba514f 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -4456,6 +4456,7 @@ class ServerArgs: is_deepseek_v4, is_nemotron_h, ) + from sglang.srt.layers.cp.bcg import supports_prefill_cp_bcg rules = [ # MLA prefill under BCG takes forward_mha, which has no eager @@ -4482,7 +4483,8 @@ class ServerArgs: # CP all_gather replay size mismatch under BCG. ( "context parallel (attn_cp_size > 1)", - lambda: self._resolved().attn_cp_size > 1, + lambda: self._resolved().attn_cp_size > 1 + and not supports_prefill_cp_bcg(self), ), # Capture builds a dummy extend forward with attn_dcp_metadata=None. ( diff --git a/test/registered/cp/test_gpt_oss_4gpu_mxfp4_cp.py b/test/registered/cp/test_gpt_oss_4gpu_mxfp4_cp.py index 6ad3aaaa9..74564f723 100644 --- a/test/registered/cp/test_gpt_oss_4gpu_mxfp4_cp.py +++ b/test/registered/cp/test_gpt_oss_4gpu_mxfp4_cp.py @@ -22,6 +22,8 @@ class TestGptOss4GpuMxfp4CP(BaseTestGptOss): "4", "--cp-strategy", "zigzag", + "--cuda-graph-backend-prefill", + "breakable", "--cuda-graph-max-bs-decode", "200", ],