From ee843a946b9b566a3ec429c6fd8de5f22f3a85dd Mon Sep 17 00:00:00 2001 From: laoyao0822 Date: Sat, 13 Jun 2026 01:28:35 +0800 Subject: [PATCH] Keep chunked CP prefills solo during bs>1 admission 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. --- python/sglang/srt/environ.py | 6 -- .../attention/nsa/cp_shared_kv_compose.py | 18 +---- python/sglang/srt/managers/schedule_policy.py | 41 +--------- .../test_cp_shared_kv_compose_v2_8rank.py | 14 +--- .../unit/managers/test_prefill_adder.py | 77 ------------------- 5 files changed, 9 insertions(+), 147 deletions(-) diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index ea0d4d0f5..28708d39a 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -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. diff --git a/python/sglang/srt/layers/attention/nsa/cp_shared_kv_compose.py b/python/sglang/srt/layers/attention/nsa/cp_shared_kv_compose.py index 531e5d229..05cd1138f 100644 --- a/python/sglang/srt/layers/attention/nsa/cp_shared_kv_compose.py +++ b/python/sglang/srt/layers/attention/nsa/cp_shared_kv_compose.py @@ -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 diff --git a/python/sglang/srt/managers/schedule_policy.py b/python/sglang/srt/managers/schedule_policy.py index c77bafb6c..abae4051a 100644 --- a/python/sglang/srt/managers/schedule_policy.py +++ b/python/sglang/srt/managers/schedule_policy.py @@ -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) diff --git a/test/manual/test_cp_shared_kv_compose_v2_8rank.py b/test/manual/test_cp_shared_kv_compose_v2_8rank.py index 90de69b69..47ea1cda1 100644 --- a/test/manual/test_cp_shared_kv_compose_v2_8rank.py +++ b/test/manual/test_cp_shared_kv_compose_v2_8rank.py @@ -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 ( diff --git a/test/registered/unit/managers/test_prefill_adder.py b/test/registered/unit/managers/test_prefill_adder.py index 2d6914f4f..d72e40b15 100644 --- a/test/registered/unit/managers/test_prefill_adder.py +++ b/test/registered/unit/managers/test_prefill_adder.py @@ -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(