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

Fix prefill CP graph overflow with larger bucket search (#33906)

Browse Source
This commit is contained in:
Baizhou Zhang
2026-08-07 01:33:14 -07:00
committed by GitHub
parent 7395ee833e
commit 5e60363960
3 changed files with 263 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/cp/bcg.py
+65 -1
View File
@@ -17,16 +17,19 @@
from __future__ import annotations
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any, Dict
from typing import TYPE_CHECKING, Any, Dict, Optional
import torch
from sglang.srt.layers.cp.base import get_cp_strategy
from sglang.srt.layers.cp.padding import get_cp_padding_align_size
from sglang.srt.layers.cp.utils import (
cp_gather_after_forward,
cp_split_before_forward,
enable_cp_v2,
prepare_cp_forward,
)
from sglang.srt.layers.cp.zigzag import ZigzagCPStrategy
from sglang.srt.model_executor.forward_batch_info import PPProxyTensors
if TYPE_CHECKING:
@@ -106,6 +109,67 @@ class PrefillCPBCGInput:
),
)
def required_local_tokens(self, extend_seq_lens: Any) -> Optional[int]:
"""Return the aligned CP-local rows required by a live zigzag layout."""
strategy = get_cp_strategy()
if not isinstance(strategy, ZigzagCPStrategy) or extend_seq_lens is None:
return None
cp_segment_num = strategy.cp_size * 2
per_rank_logical_tokens = [0] * strategy.cp_size
for raw_length in extend_seq_lens:
base, remainder = divmod(int(raw_length), cp_segment_num)
for rank in range(strategy.cp_size):
opposite_rank = cp_segment_num - 1 - rank
per_rank_logical_tokens[rank] += (
base * 2 + int(rank < remainder) + int(opposite_rank < remainder)
)
align_size = get_cp_padding_align_size()
return (
(max(per_rank_logical_tokens) + align_size - 1) // align_size * align_size
)
def select_replay_bucket(
self,
*,
num_tokens: int,
required_local_tokens: int,
capture_num_tokens: list[int],
max_padding_factor: int,
) -> Optional[int]:
"""Return the smallest global capture whose CP-local rows fit."""
max_num_tokens = num_tokens * max_padding_factor
for bucket in capture_num_tokens:
if bucket < num_tokens:
continue
if bucket > max_num_tokens:
break
captured_local_tokens = self.bucket_local_tokens.get(bucket)
if (
captured_local_tokens is not None
and required_local_tokens <= captured_local_tokens
):
return bucket
return None
def select_replay_bucket_for_batch(
self,
*,
num_tokens: int,
extend_seq_lens: Any,
capture_num_tokens: list[int],
max_padding_factor: int,
) -> Optional[int]:
required_local_tokens = self.required_local_tokens(extend_seq_lens)
if required_local_tokens is None:
return None
return self.select_replay_bucket(
num_tokens=num_tokens,
required_local_tokens=required_local_tokens,
capture_num_tokens=capture_num_tokens,
max_padding_factor=max_padding_factor,
)
def prepare(
self,
runner: PrefillCudaGraphRunner,
python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py
+31
View File
@@ -63,6 +63,7 @@ from sglang.srt.layers.cp.bcg import (
execute_prefill_cp_bcg,
filter_prefill_cp_bcg_capture_num_tokens,
)
from sglang.srt.layers.cp.utils import is_cp_v2_active
from sglang.srt.layers.dp_attention import (
DpPaddingMode,
set_dp_buffer_len,
@@ -1137,6 +1138,20 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
),
):
return False
if getattr(self, "enable_cp_v2_bcg_capture", False) and is_cp_v2_active(
forward_batch
):
assert self.prefill_cp_bcg_input is not None
if (
self.prefill_cp_bcg_input.select_replay_bucket_for_batch(
num_tokens=len(forward_batch.input_ids),
extend_seq_lens=forward_batch.extend_seq_lens_cpu,
capture_num_tokens=self.capture_num_tokens,
max_padding_factor=_MAX_PREFILL_CUDA_GRAPH_PADDING_FACTOR,
)
is None
):
return False
# Multi-req replay is supported by body-capture backends via the
# layer_model.forward monkey-patch in replay(): the captured graph runs
# the transformer stack, then the outer model.forward runs
@@ -1418,6 +1433,22 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
"""
num_tokens = len(forward_batch.input_ids)
static_num_tokens = self._pad_to_bucket(num_tokens, self.capture_num_tokens)
if getattr(self, "enable_cp_v2_bcg_capture", False) and is_cp_v2_active(
forward_batch
):
assert self.prefill_cp_bcg_input is not None
static_num_tokens = (
self.prefill_cp_bcg_input.select_replay_bucket_for_batch(
num_tokens=num_tokens,
extend_seq_lens=forward_batch.extend_seq_lens_cpu,
capture_num_tokens=self.capture_num_tokens,
max_padding_factor=_MAX_PREFILL_CUDA_GRAPH_PADDING_FACTOR,
)
)
if static_num_tokens is None:
raise RuntimeError(
"Prefill CUDA graph replay was admitted without a fitting bucket"
)
self.raw_num_tokens = num_tokens
bs = forward_batch.batch_size
test/registered/cp/test_cp_strategy_unit.py
+167
View File
@@ -14,6 +14,7 @@ from sglang.srt.layers.cp.base import (
is_interleave,
is_zigzag,
)
from sglang.srt.layers.cp.bcg import PrefillCPBCGInput
from sglang.srt.layers.cp.interleave import InterleaveCPStrategy
from sglang.srt.layers.cp.padding import (
get_cp_padding_align_size,
@@ -28,6 +29,14 @@ from sglang.srt.layers.cp.utils import (
)
from sglang.srt.layers.cp.zigzag import ZigzagCPStrategy
from sglang.srt.mem_cache.memory_pool import KVWriteLoc
from sglang.srt.model_executor.cuda_graph_config import Backend
from sglang.srt.model_executor.forward_batch_info import (
CaptureHiddenMode,
ForwardMode,
)
from sglang.srt.model_executor.runner.prefill_cuda_graph_runner import (
PrefillCudaGraphRunner,
)
from sglang.srt.runtime_context import get_parallel
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase
@@ -102,6 +111,164 @@ class TestCPStrategyUnit(CustomTestCase):
self.assertIsNotNone(get_cp_strategy())
class TestPrefillCPBCGReplay(CustomTestCase):
def tearDown(self):
init_cp_strategy(SimpleNamespace(enable_prefill_cp=False))
def _make_runner(self):
runner = PrefillCudaGraphRunner.__new__(PrefillCudaGraphRunner)
runner._is_full_backend = False
runner.enable_lora = False
runner._capture_chunked_prefix = False
runner.prefill_backend_name = Backend.TC_PIECEWISE
runner.has_mha_companion_layers = False
runner.capture_hidden_mode = CaptureHiddenMode.NULL
runner.capture_num_tokens = [2048, 2304]
runner.max_num_tokens = 2304
runner.enable_cp_v2_bcg_capture = True
return runner
def _make_forward_batch(self):
return SimpleNamespace(
batch_size=3,
input_embeds=None,
replace_embeds=None,
mm_inputs=None,
forward_mode=ForwardMode.EXTEND,
capture_hidden_mode=CaptureHiddenMode.NULL,
global_num_tokens_cpu=None,
return_logprob=False,
input_ids=list(range(2048)),
seq_lens_cpu=[1534, 161, 353],
extend_seq_lens_cpu=[1534, 161, 353],
extend_prefix_lens_cpu=[0, 0, 0],
)
def _enable_zigzag(self):
init_cp_strategy(
SimpleNamespace(
enable_prefill_cp=True,
cp_strategy="zigzag",
attn_cp_size=4,
)
)
def test_local_capacity_overflow_uses_next_capture_bucket(self):
runner = self._make_runner()
runner.capture_num_tokens.append(2560)
runner.max_num_tokens = 2560
runner.prefill_cp_bcg_input = PrefillCPBCGInput(
input_embeds=torch.empty(0),
positions=torch.empty(0),
bucket_local_tokens={2048: 512, 2304: 576, 2560: 640},
)
forward_batch = self._make_forward_batch()
self._enable_zigzag()
with (
patch(
"sglang.srt.environ.envs.SGLANG_ENABLE_CP_V2.get",
return_value=True,
),
patch(
"sglang.srt.layers.cp.bcg.get_cp_padding_align_size",
return_value=8,
),
):
selected_buckets = []
for cp_rank in range(4):
with get_parallel().override(attn_cp_rank=cp_rank, attn_cp_size=4):
selected_buckets.append(
runner.prefill_cp_bcg_input.select_replay_bucket_for_batch(
num_tokens=2048,
extend_seq_lens=[1534, 161, 353],
capture_num_tokens=runner.capture_num_tokens,
max_padding_factor=2,
)
)
self.assertEqual(selected_buckets, [2304, 2304, 2304, 2304])
with get_parallel().override(attn_cp_rank=0, attn_cp_size=4):
self.assertTrue(runner.can_run_graph(forward_batch))
def test_bucket_search_preserves_two_x_padding_limit(self):
runner = self._make_runner()
runner.prefill_cp_bcg_input = PrefillCPBCGInput(
input_embeds=torch.empty(0),
positions=torch.empty(0),
bucket_local_tokens={2048: 512, 2304: 576},
)
self.assertIsNone(
runner.prefill_cp_bcg_input.select_replay_bucket(
num_tokens=1024,
required_local_tokens=520,
capture_num_tokens=runner.capture_num_tokens,
max_padding_factor=2,
)
)
def test_bucket_search_falls_back_when_no_capture_has_capacity(self):
runner = self._make_runner()
runner.prefill_cp_bcg_input = PrefillCPBCGInput(
input_embeds=torch.empty(0),
positions=torch.empty(0),
bucket_local_tokens={2048: 512, 2304: 516},
)
forward_batch = self._make_forward_batch()
self._enable_zigzag()
with (
get_parallel().override(attn_cp_rank=0, attn_cp_size=4),
patch(
"sglang.srt.environ.envs.SGLANG_ENABLE_CP_V2.get",
return_value=True,
),
patch(
"sglang.srt.layers.cp.bcg.get_cp_padding_align_size",
return_value=8,
),
):
self.assertFalse(runner.can_run_graph(forward_batch))
def test_load_batch_uses_selected_larger_bucket(self):
class StopAfterRecordingFill(Exception):
pass
class RecordingRegistry:
padded_num_tokens = None
def fill_from(self, _source, **kwargs):
self.padded_num_tokens = kwargs["padded_num_tokens"]
raise StopAfterRecordingFill
runner = self._make_runner()
runner.prefill_cp_bcg_input = PrefillCPBCGInput(
input_embeds=torch.empty(0),
positions=torch.empty(0),
bucket_local_tokens={2048: 512, 2304: 576},
)
runner.buffer_registry = RecordingRegistry()
forward_batch = self._make_forward_batch()
self._enable_zigzag()
with (
get_parallel().override(attn_cp_rank=0, attn_cp_size=4),
patch(
"sglang.srt.environ.envs.SGLANG_ENABLE_CP_V2.get",
return_value=True,
),
patch(
"sglang.srt.layers.cp.bcg.get_cp_padding_align_size",
return_value=8,
),
self.assertRaises(StopAfterRecordingFill),
):
runner.load_batch(forward_batch)
self.assertEqual(runner.buffer_registry.padded_num_tokens, 2304)
class TestCPZigzagStrategy(CustomTestCase):
def setUp(self):
init_cp_strategy(