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

@@ -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.

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
 

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)