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(