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

Keep chunked CP prefills solo during bs>1 admission

Browse Source 
Revert the tail-chunk co-batching gate from 54c056af because allowing a continued chunk to share the next CP bs>1 batch reopens the mixed chunk/page-tail scheduler risks we are currently avoiding. Keep the independent real-prefix budget accounting so chunked requests still contribute their carried prefix to CP cached-token and buffer estimates.\n\nConstraint: Chunked-prefill requests must remain solo until the CP split/page-tail contract is revalidated for mixed batches.\nRejected: Full revert of 54c056af | it would also drop true-prefix budget accounting and under-estimate cache/buffer pressure for admitted chunks.\nConfidence: high\nScope-risk: moderate\nDirective: Do not reintroduce tail-chunk co-batching without tests covering page-tail split, CP buffer admission, and ETE chunked+cache-hit replay.\nTested: Local py_compile for environ.py, schedule_policy.py, cp_shared_kv_compose.py, test_prefill_adder.py, test_cp_shared_kv_compose_v2_8rank.py.\nTested: Remote cjy-glm5-new PYTHONPATH=python pytest -q test/registered/unit/managers/test_prefill_adder.py -> 25 passed.\nTested: Remote cjy-glm5-new PYTHONPATH=python pytest -q test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py -> 157 passed, 2 subtests passed.\nNot-tested: Full mixed replay with chunked-prefill traffic after service restart.
This commit is contained in:
laoyao0822
2026-06-13 01:28:35 +08:00
parent f290e1e35c
commit ee843a946b
5 changed files with 9 additions and 147 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
python/sglang/srt/environ.py
View File

@@ -236,12 +236,6 @@ class Envs:
# Total symm slab size override in MB (both halves). 0 = pool-derived
# hard bound (logical pages x dense page unit, see design doc).
SGLANG_CP_SHARED_KV_SYMM_HEAP_MB = EnvInt(0)
# Allow a chunked prefill request (page-aligned carried prefix) to
# co-batch with other requests under CP shared-KV bs>1: a tail chunk
# leaves extend-cap headroom for short-extend cache-hit requests; a full
# chunk still ends the scan by budget. See
# docs_internal/perf/prefill-compute-intensity-plan.md S1.
SGLANG_CP_PREFILL_MIX_CHUNKED = EnvBool(False)
# Cache-affinity batch formation (plan doc S4): keep FCFS, but skip a
# COLD candidate while the batch is warm-led so cache-hit requests group
# into one dense forward instead of being split by a cold 60K chunk.

18
python/sglang/srt/layers/attention/nsa/cp_shared_kv_compose.py
View File

@@ -581,23 +581,7 @@ def compute_staging_capacity_pages(
getattr(server_args, "cp_shared_kv_prefill_max_batch_requests", None)
or 256
)
capacity = -(-int(max_extend_tokens) // int(page_size)) + int(max_requests)
# I2 identity (mix-chunked co-batching, plan doc S1.4): admission
# enforces sum(extend) <= max_extend_tokens and bs <= max_requests,
# which bounds current pages by ceil(cap/page) + bs <= capacity.
# The two sides derive from the SAME numbers today; this assertion
# locks the coupling against either formula changing alone.
if envs.SGLANG_CP_PREFILL_MIX_CHUNKED.get():
admission_bound = -(-int(max_extend_tokens) // int(page_size)) + int(
max_requests
)
assert admission_bound <= capacity, (
"[CP_SHARED_KV_FAIL_FAST][mix_chunked] symm staging capacity "
f"({capacity} pages) no longer covers the admission bound "
f"({admission_bound} pages) — the extend-cap/staging identity "
"was broken; revisit plan doc S1.4 (I2)."
)
return capacity
return -(-int(max_extend_tokens) // int(page_size)) + int(max_requests)
return (int(kv_pool_tokens) // int(page_size)) * int(cp_size) + 512

41
python/sglang/srt/managers/schedule_policy.py
View File

@@ -501,7 +501,6 @@ class PrefillAdder:
self.preempt_list = []
self.new_chunked_req = None
self.has_chunked_req_in_batch = False
self._chunked_req_in_batch = None
self.log_hit_tokens = 0
# TODO(lsyin): report the real input tokens excluding page alignment
self.log_input_tokens = 0
@@ -663,21 +662,6 @@ class PrefillAdder:
and len(self.can_run_list) >= 1
)
def _cp_chunked_prefix_page_aligned(self) -> bool:
"""I1 guard for chunked co-batching (no chunked req -> vacuously True).
cache_unfinished_req(chunked=True) floors the carried prefix to a
page multiple, so chunks normally pass. The rare exception is a
sub-page tail carried into the FINAL chunk: a non-aligned prefix in a
multi-request batch would trip the CP page-aligned split fail-fast,
so that chunk must stay solo.
"""
req = self._chunked_req_in_batch
if req is None:
return True
return len(req.prefix_indices) % self.page_size == 0
def _cp_prefill_request_limit_reached(self) -> bool:
if not (
self._is_cp_prefill_context()
@@ -941,14 +925,10 @@ class PrefillAdder:
req.fill_ids = req.fill_ids[: len(req.prefix_indices) + req.extend_input_len]
self.can_run_list.append(req)
self.has_chunked_req_in_batch = True
self._chunked_req_in_batch = req
# Seed the budget with the chunk's TRUE prefix (the page-floored pages
# computed by earlier chunks, carried in prefix_indices). With 0 the
# CP shared-KV admission state (cached-token tally and the buffer
# estimator's k_rows for the mqa_logits term) under-counts the chunk's
# footprint for every request admitted after it. Do NOT add fail-fast
# gates here — the chunked req must always be admitted (see the
# memory-leak note above); the CP caps bind on the FOLLOWING requests.
# A continued chunk carries a real page-floored prefix. Keep that
# footprint in the CP cached/buffer admission state even though the
# chunked request itself remains solo; otherwise later capacity checks
# underestimate the batch that was just admitted.
self._update_prefill_budget(
len(req.prefix_indices),
req.extend_input_len,
@@ -1091,19 +1071,7 @@ class PrefillAdder:
self._is_cp_prefill_context()
and len(self.can_run_list) > 0
and (has_chunked_req or self.has_chunked_req_in_batch)
and (
not envs.SGLANG_CP_PREFILL_MIX_CHUNKED.get()
or not self._cp_chunked_prefix_page_aligned()
)
):
# Without the mix flag a chunked request keeps its batch solo.
# With it, requests may co-batch with the chunk: the CP caps
# (extend / cached / buffer, seeded with the chunk's footprint by
# add_chunked_req) gate each one, so a FULL chunk still ends the
# scan by budget while a tail chunk leaves room. A chunked
# prefix that is not page-aligned (rare sub-page final-chunk
# tail) would fail the CP page-aligned split fail-fast in a
# multi-request batch, so it stays solo.
return AddReqResult.OTHER
if self._cp_prefill_multi_request_disabled():
return AddReqResult.OTHER
@@ -1220,7 +1188,6 @@ class PrefillAdder:
self.can_run_list.append(req)
self.new_chunked_req = req
self.has_chunked_req_in_batch = True
self._chunked_req_in_batch = req
self._req_inc_lock_ref(req)
self._update_prefill_budget(prefix_len, trunc_len, 0)

14
test/manual/test_cp_shared_kv_compose_v2_8rank.py
View File

@@ -42,17 +42,11 @@ class _CpGroupShim:
def _build_scenario(rank: int, cp_size: int, device: torch.device):
"""bs=5 with mixed prefix/extend lengths, extends crossing page bounds.
The last request is CHUNK-SHAPED (large page-aligned carried prefix +
a tail-chunk-sized extend): with mix-chunked co-batching (plan doc S1)
a continued chunk is exactly such a (prefix, extend) pair to the CP
compose machinery, so this scenario byte-validates the mixed batch.
"""
"""bs=4 with mixed prefix/extend lengths, extends crossing page bounds."""
page_size = 64
prefix_lens = [640, 1280, 320, 640, 2048]
extend_lens = [95, 130, 64, 200, 512]
batch_size = len(prefix_lens)
batch_size = 4
prefix_lens = [640, 1280, 320, 640]
extend_lens = [95, 130, 64, 200]
kv_dim = 656 # fp8 layout bytes/token
from sglang.srt.mem_cache.cp_shared_kv_compute_owner import (

77
test/registered/unit/managers/test_prefill_adder.py
View File

@@ -6,8 +6,6 @@ from unittest.mock import MagicMock
import torch
from sglang.srt.environ import envs
if "sgl_kernel" not in sys.modules:
sys.modules["sgl_kernel"] = types.ModuleType("sgl_kernel")
sys.modules["sgl_kernel"].__file__ = "sgl_kernel_stub.py"
@@ -987,81 +985,6 @@ class TestPrefillAdder(CustomTestCase):
self.assertEqual(adder.cp_shared_kv_prefill_total_cached_tokens, 256)
self.assertEqual(adder.cp_shared_kv_prefill_total_extend_tokens, 128)
def _mix_chunked_adder(self, *, rem_chunk_tokens, extend_cap):
set_global_server_args_for_scheduler(
ServerArgs(
model_path="dummy",
enable_nsa_prefill_context_parallel=True,
nsa_prefill_cp_mode="in-seq-split",
)
)
self.mock_token_allocator.available_size.return_value = 100000
return self.create_adder(
self.create_running_batch(),
page_size=64,
rem_input_tokens=8192,
rem_chunk_tokens=rem_chunk_tokens,
enable_cp_shared_kv_prefill_bs_gt1=True,
cp_shared_kv_prefill_max_batch_requests=8,
cp_shared_kv_prefill_max_total_extend_tokens=extend_cap,
)
def test_cp_prefill_mix_chunked_tail_chunk_admits_following_requests(self):
# Plan doc S1: with the flag on, a TAIL chunk (extend below the chunk
# budget, page-aligned carried prefix) leaves extend-cap headroom and
# the following short-extend request co-batches with it.
with envs.SGLANG_CP_PREFILL_MIX_CHUNKED.override(True):
adder = self._mix_chunked_adder(rem_chunk_tokens=4096, extend_cap=4096)
chunked = self.create_prefill_req("chunked", extend_input_len=1280)
chunked.prefix_indices = torch.zeros((4096,), dtype=torch.int64)
self.assertIsNone(adder.add_chunked_req(chunked)) # tail chunk
normal = self.create_prefill_req("normal", extend_input_len=128)
self.assertEqual(
adder.add_one_req(
normal, has_chunked_req=True, truncation_align_size=None
),
AddReqResult.CONTINUE,
)
self.assertEqual(
[req.rid for req in adder.can_run_list], ["chunked", "normal"]
)
def test_cp_prefill_mix_chunked_full_chunk_stays_solo_by_budget(self):
# A FULL chunk consumes the whole (chunk-clamped) extend cap, so the
# first following request is rejected by the extend gate — no special
# code, the budget arithmetic ends the scan.
with envs.SGLANG_CP_PREFILL_MIX_CHUNKED.override(True):
adder = self._mix_chunked_adder(rem_chunk_tokens=256, extend_cap=4096)
chunked = self.create_prefill_req("chunked", extend_input_len=512)
chunked.prefix_indices = torch.zeros((4096,), dtype=torch.int64)
self.assertIs(adder.add_chunked_req(chunked), chunked) # truncated
self.assertEqual(chunked.extend_input_len, 256)
normal = self.create_prefill_req("normal", extend_input_len=128)
self.assertEqual(
adder.add_one_req(
normal, has_chunked_req=True, truncation_align_size=None
),
AddReqResult.OTHER,
)
self.assertEqual([req.rid for req in adder.can_run_list], ["chunked"])
def test_cp_prefill_mix_chunked_non_aligned_prefix_stays_solo(self):
# I1 guard: a chunked prefix that is not a page multiple would break
# the CP page-aligned split in a multi-request batch — keep it solo.
with envs.SGLANG_CP_PREFILL_MIX_CHUNKED.override(True):
adder = self._mix_chunked_adder(rem_chunk_tokens=4096, extend_cap=4096)
chunked = self.create_prefill_req("chunked", extend_input_len=1280)
chunked.prefix_indices = torch.zeros((100,), dtype=torch.int64)
self.assertIsNone(adder.add_chunked_req(chunked))
normal = self.create_prefill_req("normal", extend_input_len=128)
self.assertEqual(
adder.add_one_req(
normal, has_chunked_req=True, truncation_align_size=None
),
AddReqResult.OTHER,
)
self.assertEqual([req.rid for req in adder.can_run_list], ["chunked"])
def test_cp_prefill_total_cached_limit_stops_second_cached_request(self):
set_global_server_args_for_scheduler(
ServerArgs(