chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

Default CP bs>1 extend admission to chunk budget

Browse Source 
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>
This commit is contained in:
laoyao0822
2026-06-11 03:51:41 +08:00
parent d696039092
commit adf357b02c
2 changed files with 36 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

20
python/sglang/srt/managers/schedule_policy.py
View File

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

25
test/registered/unit/managers/test_prefill_adder.py
View File

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