From 54c056af833dab75f6c16282909bca5b12ee28f9 Mon Sep 17 00:00:00 2001 From: leavelet Date: Fri, 12 Jun 2026 06:38:47 +0000 Subject: [PATCH] Let tail chunks co-batch with cache-hit requests (SGLANG_CP_PREFILL_MIX_CHUNKED) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 专题 S1 (design: docs_internal/perf/prefill-compute-intensity-plan.md S1.0-S1.11). A batch containing a chunked prefill request has been forced to bs=1 by the CP gate, so every chunk of a long prompt monopolizes a forward while short-extend cache-hit continuations queue — the direct cause of the replay TTFT tail (p90 19.3s / p99 49.2s at 91.8% cache hit). Yet mixed chunk batches already occur today (a freshly-chunked request keeps earlier-admitted small requests), proving the CP forward path is mixed-chunk-safe; only admission was asymmetric. Three changes, the first flag-independent: - add_chunked_req now seeds the budget with the chunk's TRUE prefix (was 0), so the CP cached tally and the buffer estimator's mqa_logits k_rows see the chunk's footprint before any later request is admitted (landmine D1). - New SGLANG_CP_PREFILL_MIX_CHUNKED (default OFF): with it on, the gate admits requests after a chunked one and lets the existing CP caps (extend / cached / buffer, now correctly seeded) bound the batch — a FULL chunk still ends the scan by consuming the chunk-clamped extend cap; only a tail chunk leaves headroom. A chunked prefix that is not page-aligned (rare sub-page final-chunk tail) keeps its batch solo (the CP page-aligned split would fail-fast otherwise). - The symm staging capacity identity (admission extend cap + request slack == staging pages) is asserted when the flag is on, locking the coupling the design relies on (plan doc S1.4 I2). Tests: 4 new adder units (budget seeding; tail chunk admits followers; full chunk solo by budget; non-aligned prefix solo); the 8-rank byte-exactness scenario gains a chunk-shaped request (2048-token page-aligned carried prefix + 512 extend) — all four phases (legacy/v2/symm/prefetch) byte-identical on g0033. Known pre-existing cross-file pollution noted in problems.md P16. Co-Authored-By: Claude Fable 5 --- python/sglang/srt/environ.py | 6 + .../attention/nsa/cp_shared_kv_compose.py | 18 ++- python/sglang/srt/managers/schedule_policy.py | 39 ++++++- .../test_cp_shared_kv_compose_v2_8rank.py | 14 ++- .../unit/managers/test_prefill_adder.py | 103 ++++++++++++++++++ 5 files changed, 174 insertions(+), 6 deletions(-) diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index a27a2ace9..41b12bf40 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -236,6 +236,12 @@ 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) # NSA MQA logits are materialized as fp32 [q, k] buffers inside DeepGEMM. # Lower values split query rows more aggressively to cap peak temporary memory. SGLANG_NSA_MQA_LOGITS_FREE_MEM_FRACTION = EnvFloat(0.2) 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 05cd1138f..531e5d229 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,7 +581,23 @@ def compute_staging_capacity_pages( getattr(server_args, "cp_shared_kv_prefill_max_batch_requests", None) or 256 ) - return -(-int(max_extend_tokens) // int(page_size)) + int(max_requests) + 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(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 a4dbd4ef0..6efdbfb46 100644 --- a/python/sglang/srt/managers/schedule_policy.py +++ b/python/sglang/srt/managers/schedule_policy.py @@ -447,6 +447,7 @@ 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 @@ -608,6 +609,21 @@ 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() @@ -871,8 +887,16 @@ 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. self._update_prefill_budget( - 0, + len(req.prefix_indices), req.extend_input_len, ( min(req.sampling_params.max_new_tokens, CLIP_MAX_NEW_TOKENS) @@ -1013,7 +1037,19 @@ 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 @@ -1130,6 +1166,7 @@ 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 47ea1cda1..90de69b69 100644 --- a/test/manual/test_cp_shared_kv_compose_v2_8rank.py +++ b/test/manual/test_cp_shared_kv_compose_v2_8rank.py @@ -42,11 +42,17 @@ class _CpGroupShim: def _build_scenario(rank: int, cp_size: int, device: torch.device): - """bs=4 with mixed prefix/extend lengths, extends crossing page bounds.""" + """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. + """ page_size = 64 - batch_size = 4 - prefix_lens = [640, 1280, 320, 640] - extend_lens = [95, 130, 64, 200] + prefix_lens = [640, 1280, 320, 640, 2048] + extend_lens = [95, 130, 64, 200, 512] + batch_size = len(prefix_lens) 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 c09fcc9a9..e28547c05 100644 --- a/test/registered/unit/managers/test_prefill_adder.py +++ b/test/registered/unit/managers/test_prefill_adder.py @@ -6,6 +6,8 @@ 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" @@ -845,6 +847,107 @@ class TestPrefillAdder(CustomTestCase): ) self.assertEqual([req.rid for req in adder.can_run_list], ["chunked"]) + def test_add_chunked_req_seeds_true_prefix_into_cp_budget(self): + # C1 (plan doc S1.1-1a): the chunk's carried prefix must count toward + # the CP cached tally so later admission gates see its footprint. + 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 = 10000 + adder = self.create_adder( + self.create_running_batch(), + page_size=64, + rem_input_tokens=4096, + rem_chunk_tokens=4096, + enable_cp_shared_kv_prefill_bs_gt1=True, + cp_shared_kv_prefill_max_batch_requests=8, + cp_shared_kv_prefill_max_total_extend_tokens=4096, + ) + chunked = self.create_prefill_req("chunked", extend_input_len=128) + chunked.prefix_indices = torch.zeros((256,), dtype=torch.int64) + adder.add_chunked_req(chunked) + 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(