Gate CP shared-KV prefill batching behind explicit limits

The scheduler can now admit multi-request NSA in-seq CP shared-KV prefill batches only when the shared-KV bs>1 flag is explicitly enabled. The gate is still disabled by default and is scoped to CP shared-KV so ordinary CP is not widened accidentally.

Batch admission is bounded by optional request-count and page-aligned extend-token limits while real memory capacity remains allocator-owned. This keeps bf16 and fp8 on the same scheduler path because dtype differences are already reflected in KV pool token/page capacity.

Constraint: bs>1 runtime paths remain guarded by existing CP shared-KV fail-fast checks.

Constraint: Scheduler must not duplicate bf16/fp8 byte-level capacity estimation.

Rejected: Open the old CP gate unconditionally | ordinary CP would inherit an unverified shared-KV-specific batching path.

Rejected: Treat the extend-token cap as a hard per-request limit | a single large request could deadlock the scheduler.

Confidence: medium

Scope-risk: moderate

Directive: Keep CP shared-KV batching gated until ETE validates EAGLE accept length, output length, and HiCache load/backup behavior under real traffic.

Tested: local py_compile for server_args, schedule_policy, scheduler, prefill_adder tests, and server_args tests.

Tested: remote g0034 py_compile for the same files.

Tested: remote g0034 pytest target set: 5 passed for parser, parameter validation, default single-request CP gate, enabled bs>1 gate, and page-aligned extend cap.

Tested: remote g0034 pytest test_prefill_adder.py => 13 passed.

Not-tested: full server_args test file has an unrelated HuggingFace DNS/config-download failure in TestPrepareServerArgs.test_prepare_server_args.

Not-tested: ETE production traffic with --enable-cp-shared-kv-prefill-bs-gt1.

python/sglang/srt/managers/schedule_policy.py

@@ -390,6 +390,9 @@ class PrefillAdder:
         max_prefill_bs: int = 0,
         max_running_requests: Optional[int] = None,
         prefill_max_requests: Optional[int] = None,
+        enable_cp_shared_kv_prefill_bs_gt1: bool = False,
+        cp_shared_kv_prefill_max_batch_requests: Optional[int] = None,
+        cp_shared_kv_prefill_max_total_extend_tokens: Optional[int] = None,
         prefill_delayer_single_pass: Optional[PrefillDelayerSinglePassExecutor] = None,
         dllm_config: Optional[DllmConfig] = None,
     ):
@@ -440,6 +443,14 @@ class PrefillAdder:
         self.prefill_max_requests = prefill_max_requests
         self.prefill_delayer_single_pass = prefill_delayer_single_pass
         self.max_prefill_bs = max_prefill_bs
+        self.enable_cp_shared_kv_prefill_bs_gt1 = enable_cp_shared_kv_prefill_bs_gt1
+        self.cp_shared_kv_prefill_max_batch_requests = (
+            cp_shared_kv_prefill_max_batch_requests
+        )
+        self.cp_shared_kv_prefill_max_total_extend_tokens = (
+            cp_shared_kv_prefill_max_total_extend_tokens
+        )
+        self.cp_shared_kv_prefill_total_extend_tokens = 0
 
     def _init_dllm_meta(self, dllm_config: DllmConfig):
         self.dllm_block_size = dllm_config.block_size
@@ -502,6 +513,45 @@ class PrefillAdder:
     def ceil_paged_tokens(self, tokens: int) -> int:
         return -(-tokens // self.page_size) * self.page_size
 
+    def _is_cp_prefill_context(self) -> bool:
+        return self.nsa_prefill_cp_in_seq_split or self.prefill_context_parallel_enabled
+
+    def _cp_prefill_multi_request_disabled(self) -> bool:
+        return (
+            self._is_cp_prefill_context()
+            and not self.enable_cp_shared_kv_prefill_bs_gt1
+            and len(self.can_run_list) >= 1
+        )
+
+    def _cp_prefill_request_limit_reached(self) -> bool:
+        if not (
+            self._is_cp_prefill_context()
+            and self.enable_cp_shared_kv_prefill_bs_gt1
+        ):
+            return False
+
+        limit = self.cp_shared_kv_prefill_max_batch_requests
+        return limit is not None and len(self.can_run_list) >= limit
+
+    def _cp_prefill_extend_limit_exceeded(self, extend_input_len: int) -> bool:
+        if not (
+            self._is_cp_prefill_context()
+            and self.enable_cp_shared_kv_prefill_bs_gt1
+        ):
+            return False
+
+        limit = self.cp_shared_kv_prefill_max_total_extend_tokens
+        if limit is None:
+            return False
+
+        projected = self.cp_shared_kv_prefill_total_extend_tokens + self.ceil_paged_tokens(
+            extend_input_len
+        )
+
+        # Do not deadlock a single large request. The limit bounds grouping, not
+        # the maximum legal request size; actual KV capacity remains allocator-owned.
+        return projected > limit and len(self.can_run_list) > 0
+
     def _get_available_device_tokens_for_load_back(self) -> int:
         if self.is_hybrid_swa:
             return (
@@ -549,6 +599,8 @@ class PrefillAdder:
         self.rem_total_token_offset += extend_input_len + max_new_tokens + page_overhead
         self.cur_rem_token_offset += extend_input_len + page_overhead
         self.rem_input_tokens -= extend_input_len
+        if self._is_cp_prefill_context():
+            self.cp_shared_kv_prefill_total_extend_tokens += extend_input_len
 
         if self.dllm_config is not None:
             self.rem_dllm_tokens -= extend_input_len
@@ -726,6 +778,8 @@ class PrefillAdder:
             self.rem_chunk_tokens is None  # chunked prefill is disabled
             or req.extend_input_len <= self.rem_chunk_tokens  # it is the last chunk
         ):
+            if self._cp_prefill_extend_limit_exceeded(req.extend_input_len):
+                return AddReqResult.OTHER
             # Non-chunked prefill
             self.can_run_list.append(req)
             self._update_prefill_budget(
@@ -739,6 +793,8 @@ class PrefillAdder:
 
             # Chunked prefill
             trunc_len = self.rem_chunk_tokens
+            if self._cp_prefill_extend_limit_exceeded(trunc_len):
+                return AddReqResult.OTHER
 
             req.set_extend_input_len(trunc_len)
             req.fill_ids = req.fill_ids[:trunc_len]
@@ -763,12 +819,9 @@ class PrefillAdder:
             )
         ):
             return AddReqResult.OTHER
-        # TODO support cp with multiple requests
-        # Enabling context parallelism currently presents precision issues;
-        # therefore, the prefill-batch setting is temporarily set to 1.
-        if (
-            self.nsa_prefill_cp_in_seq_split or self.prefill_context_parallel_enabled
-        ) and len(self.can_run_list) >= 1:
+        if self._cp_prefill_multi_request_disabled():
+            return AddReqResult.OTHER
+        if self._cp_prefill_request_limit_reached():
             return AddReqResult.OTHER
 
         if (x := self.prefill_max_requests) is not None and len(self.can_run_list) >= x:
@@ -815,6 +868,8 @@ class PrefillAdder:
 
             if input_tokens >= self.rem_input_tokens and len(self.can_run_list) != 0:
                 return AddReqResult.OTHER
+            if self._cp_prefill_extend_limit_exceeded(input_tokens):
+                return AddReqResult.OTHER
 
             if self.dllm_config is not None:
                 if self.rem_dllm_tokens <= 0:
@@ -857,6 +912,9 @@ class PrefillAdder:
                             trunc_len // truncation_align_size
                         )
 
+                if self._cp_prefill_extend_limit_exceeded(trunc_len):
+                    return AddReqResult.OTHER
+
                 # Chunked prefill
                 req.set_extend_input_len(trunc_len)
                 req.fill_ids = req.fill_ids[: len(req.prefix_indices) + trunc_len]

python/sglang/srt/managers/scheduler.py

@@ -2337,6 +2337,16 @@ class Scheduler(
             max_prefill_bs=self.max_prefill_bs,
             max_running_requests=self.max_running_requests,
             prefill_max_requests=self.server_args.prefill_max_requests,
+            enable_cp_shared_kv_prefill_bs_gt1=(
+                self.server_args.enable_cp_shared_kv_prefill_bs_gt1
+                and self.server_args.enable_nsa_prefill_cp_shared_kv
+            ),
+            cp_shared_kv_prefill_max_batch_requests=(
+                self.server_args.cp_shared_kv_prefill_max_batch_requests
+            ),
+            cp_shared_kv_prefill_max_total_extend_tokens=(
+                self.server_args.cp_shared_kv_prefill_max_total_extend_tokens
+            ),
             prefill_delayer_single_pass=prefill_delayer_single_pass,
             dllm_config=self.dllm_config,
         )

python/sglang/srt/server_args.py

@@ -672,6 +672,9 @@ class ServerArgs:
     enable_nsa_prefill_context_parallel: bool = False
     nsa_prefill_cp_mode: str = "round-robin-split"
     enable_nsa_prefill_cp_shared_kv: bool = False
+    enable_cp_shared_kv_prefill_bs_gt1: bool = False
+    cp_shared_kv_prefill_max_batch_requests: Optional[int] = None
+    cp_shared_kv_prefill_max_total_extend_tokens: Optional[int] = None
     enable_fused_qk_norm_rope: bool = False
     enable_precise_embedding_interpolation: bool = False
     enable_fused_moe_sum_all_reduce: bool = False
@@ -931,6 +934,22 @@ class ServerArgs:
             "Other backends do not implement CP shared-KV logical-page-position "
             "transfer mapping yet."
         )
+        if (
+            self.cp_shared_kv_prefill_max_batch_requests is not None
+            and self.cp_shared_kv_prefill_max_batch_requests <= 0
+        ):
+            raise ValueError(
+                "cp_shared_kv_prefill_max_batch_requests must be a positive "
+                "integer when specified."
+            )
+        if (
+            self.cp_shared_kv_prefill_max_total_extend_tokens is not None
+            and self.cp_shared_kv_prefill_max_total_extend_tokens <= 0
+        ):
+            raise ValueError(
+                "cp_shared_kv_prefill_max_total_extend_tokens must be a positive "
+                "integer when specified."
+            )
 
     def _handle_cp_hicache_layout_validation(self):
         if not (
@@ -5710,6 +5729,40 @@ class ServerArgs:
                 "Only prefill CP with NSA+MLA is supported; decode CP remains disabled."
             ),
         )
+        parser.add_argument(
+            "--enable-cp-shared-kv-prefill-bs-gt1",
+            action="store_true",
+            default=ServerArgs.enable_cp_shared_kv_prefill_bs_gt1,
+            help=(
+                "Enable experimental multi-request prefill batching for NSA "
+                "in-seq CP shared KV. Capacity is still enforced by the KV "
+                "pool allocator; use the CP shared-KV batch limits to bound "
+                "scheduler admission."
+            ),
+        )
+        parser.add_argument(
+            "--cp-shared-kv-prefill-max-batch-requests",
+            type=int,
+            default=ServerArgs.cp_shared_kv_prefill_max_batch_requests,
+            help=(
+                "Maximum number of requests admitted into one NSA in-seq CP "
+                "shared-KV prefill batch when --enable-cp-shared-kv-prefill-bs-gt1 "
+                "is set. If unset, only the generic --prefill-max-requests limit "
+                "applies."
+            ),
+        )
+        parser.add_argument(
+            "--cp-shared-kv-prefill-max-total-extend-tokens",
+            type=human_readable_int,
+            default=ServerArgs.cp_shared_kv_prefill_max_total_extend_tokens,
+            help=(
+                "Maximum page-aligned extend tokens admitted into one NSA in-seq "
+                "CP shared-KV prefill batch when --enable-cp-shared-kv-prefill-bs-gt1 "
+                "is set. A single request larger than this limit is still allowed "
+                "to run alone to avoid scheduler deadlock."
+            )
+            + f"\n\n{human_readable_int.__doc__}",
+        )
         parser.add_argument(
             "--nsa-prefill-cp-mode",
             type=str,