Decouple CP prefill batching from generic token cap

CP shared-KV bs>1 uses cp_shared_kv_prefill_max_total_extend_tokens as its grouping admission limit, but the generic max_prefill_tokens budget could still stop batching earlier. Raise only the legacy input-token admission budget for this CP path while keeping allocator-owned capacity checks unchanged.

Constraint: CP shared-KV bs>1 needs large cache-hit batches without relying on the generic max_prefill_tokens default.

Constraint: Allocator capacity must remain enforced by rem_total_tokens, cur_rem_tokens, and prepare_for_extend().

Rejected: Increase server-wide max_prefill_tokens | would change generic scheduler behavior and non-CP paths.

Confidence: high

Scope-risk: narrow

Directive: Do not use max_prefill_tokens as the CP shared-KV bs>1 grouping limit; use the CP-specific total extend token knob.

Tested: Local py_compile for schedule_policy.py and test_prefill_adder.py.

Tested: Remote pytest targeted CP PrefillAdder cases: 5 passed.

Tested: Remote pytest test/registered/unit/managers/test_prefill_adder.py: 16 passed.

Not-tested: Full ETE replay after this scheduler-only change.

python/sglang/srt/managers/schedule_policy.py

@@ -417,6 +417,7 @@ class PrefillAdder:
         self.can_run_list = []
         self.preempt_list = []
         self.new_chunked_req = None
+        self.has_chunked_req_in_batch = False
         self.log_hit_tokens = 0
         # TODO(lsyin): report the real input tokens excluding page alignment
         self.log_input_tokens = 0
@@ -450,6 +451,28 @@ class PrefillAdder:
         self.cp_shared_kv_prefill_max_total_extend_tokens = (
             cp_shared_kv_prefill_max_total_extend_tokens
         )
+        if (
+            self._is_cp_prefill_context()
+            and self.enable_cp_shared_kv_prefill_bs_gt1
+            and self.cp_shared_kv_prefill_max_total_extend_tokens is not None
+        ):
+            # CP shared-KV bs>1 owns its batch-size/token admission via
+            # cp_shared_kv_prefill_max_total_extend_tokens.  Do not let the
+            # generic max_prefill_tokens default silently cap this path below
+            # the CP-specific limit.  Keep a larger generic budget intact when
+            # it is configured, because the CP-specific checker below is the
+            # actual CP grouping limit.  The extra page keeps exact-fill CP
+            # batches from tripping the legacy generic >= budget check.
+            # Allocator capacity is still enforced by rem_total_tokens,
+            # cur_rem_tokens, and prepare_for_extend().
+            self.rem_input_tokens = (
+                max(
+                    rem_input_tokens,
+                    self.cp_shared_kv_prefill_max_total_extend_tokens
+                    + self.page_size,
+                )
+                - mixed_with_decode_tokens
+            )
         self.cp_shared_kv_prefill_total_extend_tokens = 0
 
     def _init_dllm_meta(self, dllm_config: DllmConfig):
@@ -685,6 +708,7 @@ class PrefillAdder:
         req.set_extend_input_len(min(req.extend_input_len, _rem_tokens))
         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._update_prefill_budget(
             0,
             req.extend_input_len,
@@ -819,6 +843,12 @@ class PrefillAdder:
             )
         ):
             return AddReqResult.OTHER
+        if (
+            self._is_cp_prefill_context()
+            and len(self.can_run_list) > 0
+            and (has_chunked_req or self.has_chunked_req_in_batch)
+        ):
+            return AddReqResult.OTHER
         if self._cp_prefill_multi_request_disabled():
             return AddReqResult.OTHER
         if self._cp_prefill_request_limit_reached():
@@ -921,6 +951,7 @@ class PrefillAdder:
 
                 self.can_run_list.append(req)
                 self.new_chunked_req = req
+                self.has_chunked_req_in_batch = True
 
                 self._req_inc_lock_ref(req)
                 self._update_prefill_budget(prefix_len, trunc_len, 0)