Let tail chunks co-batch with cache-hit requests (SGLANG_CP_PREFILL_MIX_CHUNKED)

专题 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 <noreply@anthropic.com>

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)

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
 

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)

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 (

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(