From adf357b02ce3096bf062259dc215ef4ea6e6856d Mon Sep 17 00:00:00 2001
From: laoyao0822 <laoyao0822@qq.com>
Date: Thu, 11 Jun 2026 03:51:41 +0800
Subject: [PATCH] Default CP bs>1 extend admission to chunk budget

When chunked prefill is active, CP shared-KV bs>1 cannot consume more extend
tokens than the current chunk budget. If the CP-specific extend-token limit is
omitted, default it to rem_chunk_tokens so scheduler admission reflects the
reachable chunk capacity. The request-count and cached-token knobs keep their
None-as-unlimited behavior.

Constraint: CP bs>1 batching must not advertise a larger extend batch than chunked prefill can execute.
Rejected: Require users to always set --cp-shared-kv-prefill-max-total-extend-tokens | the safe default is already available from chunked prefill state.
Rejected: Default batch request or cached-token limits | those are policy knobs and None should remain unlimited.
Confidence: high
Scope-risk: narrow
Directive: Keep --cp-shared-kv-prefill-max-total-extend-tokens as min(user_limit, chunk_budget) when both exist.
Tested: Local py_compile for schedule_policy.py and test_prefill_adder.py.
Tested: Remote g0034 cjy-glm5-new targeted prefill_adder tests: 2 passed.
Not-tested: Full ETE scheduler batching distribution after defaulting the extend limit.

Co-authored-by: OmX <omx@oh-my-codex.dev>
---
 python/sglang/srt/managers/schedule_policy.py | 20 ++++++++-------
 .../unit/managers/test_prefill_adder.py       | 25 +++++++++++++++++++
 2 files changed, 36 insertions(+), 9 deletions(-)

diff --git a/python/sglang/srt/managers/schedule_policy.py b/python/sglang/srt/managers/schedule_policy.py
index 2752c6b0e..a4dbd4ef0 100644
--- a/python/sglang/srt/managers/schedule_policy.py
+++ b/python/sglang/srt/managers/schedule_policy.py
@@ -483,18 +483,20 @@ class PrefillAdder:
         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
             and self.rem_chunk_tokens is not None
         ):
             # When chunked prefill is active, a batch cannot consume more
-            # extend tokens than the current chunk budget.  Use the smaller
-            # value as the effective CP-specific grouping limit so the generic
-            # max_prefill_tokens lift below does not advertise unreachable
-            # capacity.
-            self.cp_shared_kv_prefill_max_total_extend_tokens = min(
-                self.cp_shared_kv_prefill_max_total_extend_tokens,
-                max(self.rem_chunk_tokens, 0),
-            )
+            # extend tokens than the current chunk budget.  If the CP-specific
+            # limit is omitted, default it to the chunk budget; other CP bs>1
+            # admission knobs keep their None == unlimited semantics.
+            chunk_extend_limit = max(self.rem_chunk_tokens, 0)
+            if self.cp_shared_kv_prefill_max_total_extend_tokens is None:
+                self.cp_shared_kv_prefill_max_total_extend_tokens = chunk_extend_limit
+            else:
+                self.cp_shared_kv_prefill_max_total_extend_tokens = min(
+                    self.cp_shared_kv_prefill_max_total_extend_tokens,
+                    chunk_extend_limit,
+                )
         self.cp_shared_kv_prefill_max_total_cached_tokens = (
             cp_shared_kv_prefill_max_total_cached_tokens
         )
diff --git a/test/registered/unit/managers/test_prefill_adder.py b/test/registered/unit/managers/test_prefill_adder.py
index b22cb0a4b..c09fcc9a9 100644
--- a/test/registered/unit/managers/test_prefill_adder.py
+++ b/test/registered/unit/managers/test_prefill_adder.py
@@ -754,6 +754,31 @@ class TestPrefillAdder(CustomTestCase):
         # limit, not the raw 256-token CP limit.
         self.assertEqual(adder.rem_input_tokens, 192)
 
+    def test_cp_prefill_total_extend_limit_defaults_to_chunk_budget(self):
+        set_global_server_args_for_scheduler(
+            ServerArgs(
+                model_path="dummy",
+                enable_nsa_prefill_context_parallel=True,
+                nsa_prefill_cp_mode="in-seq-split",
+            )
+        )
+        adder = self.create_adder(
+            self.create_running_batch(),
+            page_size=64,
+            rem_input_tokens=192,
+            rem_chunk_tokens=128,
+            enable_cp_shared_kv_prefill_bs_gt1=True,
+            cp_shared_kv_prefill_max_batch_requests=None,
+            cp_shared_kv_prefill_max_total_extend_tokens=None,
+            cp_shared_kv_prefill_max_total_cached_tokens=None,
+        )
+
+        self.assertIsNone(adder.cp_shared_kv_prefill_max_batch_requests)
+        self.assertEqual(adder.cp_shared_kv_prefill_max_total_extend_tokens, 128)
+        self.assertIsNone(adder.cp_shared_kv_prefill_max_total_cached_tokens)
+        # The generic budget lift uses the effective defaulted extend limit.
+        self.assertEqual(adder.rem_input_tokens, 192)
+
     def test_cp_prefill_total_extend_limit_does_not_bypass_allocator_capacity(self):
         set_global_server_args_for_scheduler(
             ServerArgs(