From 4f65d7a176898faf4834c6f5b95e101b32991e17 Mon Sep 17 00:00:00 2001 From: laoyao0822 Date: Wed, 10 Jun 2026 22:21:46 +0800 Subject: [PATCH] Bound CP prefill batches by cached-token pressure High cache-hit CP shared-KV batches can have small extend tokens while still carrying substantial prefix/load-back work. Add a CP-specific cached-token admission limit so operators can bound that pressure independently from extend-token batching. Constraint: The limit must not deadlock a single high-cache-hit request; it only stops adding additional requests to a non-empty batch. Constraint: Cached tokens are counted after L2 load-back planning via prefix_len, so L1 hits and successful L2 hits share one scheduler budget. Rejected: Reuse max_prefill_tokens | it limits generic input budget and does not represent cached-token work. Rejected: Count only L1 prefix before load-back | would miss L2 hit pressure, which is one of the target cases. Confidence: high Scope-risk: moderate Directive: Keep cached-token and extend-token limits separate; they bound different scheduler costs. Tested: Remote pytest targeted cached-token PrefillAdder cases: 2 passed. Tested: Remote pytest test/registered/unit/managers/test_prefill_adder.py: 18 passed. Tested: Remote ServerArgs CP validation smoke: SERVER_ARGS_CACHED_LIMIT_OK. Not-tested: Full ETE replay with a production cached-token limit value. --- python/sglang/srt/managers/schedule_policy.py | 36 ++++++++++ python/sglang/srt/managers/scheduler.py | 10 ++- python/sglang/srt/server_args.py | 23 +++++++ .../unit/managers/test_prefill_adder.py | 69 +++++++++++++++++++ 4 files changed, 137 insertions(+), 1 deletion(-) diff --git a/python/sglang/srt/managers/schedule_policy.py b/python/sglang/srt/managers/schedule_policy.py index 579a74149..661ef5aba 100644 --- a/python/sglang/srt/managers/schedule_policy.py +++ b/python/sglang/srt/managers/schedule_policy.py @@ -393,6 +393,7 @@ class PrefillAdder: 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, + cp_shared_kv_prefill_max_total_cached_tokens: Optional[int] = None, prefill_delayer_single_pass: Optional[PrefillDelayerSinglePassExecutor] = None, dllm_config: Optional[DllmConfig] = None, ): @@ -451,6 +452,9 @@ class PrefillAdder: self.cp_shared_kv_prefill_max_total_extend_tokens = ( cp_shared_kv_prefill_max_total_extend_tokens ) + self.cp_shared_kv_prefill_max_total_cached_tokens = ( + cp_shared_kv_prefill_max_total_cached_tokens + ) if ( self._is_cp_prefill_context() and self.enable_cp_shared_kv_prefill_bs_gt1 @@ -474,6 +478,7 @@ class PrefillAdder: - mixed_with_decode_tokens ) self.cp_shared_kv_prefill_total_extend_tokens = 0 + self.cp_shared_kv_prefill_total_cached_tokens = 0 def _init_dllm_meta(self, dllm_config: DllmConfig): self.dllm_block_size = dllm_config.block_size @@ -575,6 +580,25 @@ class PrefillAdder: # the maximum legal request size; actual KV capacity remains allocator-owned. return projected > limit and len(self.can_run_list) > 0 + def _cp_prefill_cached_limit_exceeded(self, cached_tokens: 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_cached_tokens + if limit is None: + return False + + projected = self.cp_shared_kv_prefill_total_cached_tokens + self.ceil_paged_tokens( + cached_tokens + ) + + # Do not deadlock a single high-cache-hit request. This limit bounds + # grouping pressure from prefix/load-back work, not legal request size. + 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 ( @@ -632,6 +656,10 @@ class PrefillAdder: self.log_hit_tokens += prefix_len self.log_input_tokens += extend_input_len + if self._is_cp_prefill_context(): + self.cp_shared_kv_prefill_total_cached_tokens += self.ceil_paged_tokens( + prefix_len + ) def _get_dllm_remain_tokens(self) -> int: _rem_tokens = min( @@ -804,6 +832,8 @@ class PrefillAdder: ): if self._cp_prefill_extend_limit_exceeded(req.extend_input_len): return AddReqResult.OTHER + if self._cp_prefill_cached_limit_exceeded(0): + return AddReqResult.OTHER # Non-chunked prefill self.can_run_list.append(req) self._update_prefill_budget( @@ -819,6 +849,8 @@ class PrefillAdder: trunc_len = self.rem_chunk_tokens if self._cp_prefill_extend_limit_exceeded(trunc_len): return AddReqResult.OTHER + if self._cp_prefill_cached_limit_exceeded(0): + return AddReqResult.OTHER req.set_extend_input_len(trunc_len) req.fill_ids = req.fill_ids[:trunc_len] @@ -900,6 +932,8 @@ class PrefillAdder: return AddReqResult.OTHER if self._cp_prefill_extend_limit_exceeded(input_tokens): return AddReqResult.OTHER + if self._cp_prefill_cached_limit_exceeded(prefix_len): + return AddReqResult.OTHER if self.dllm_config is not None: if self.rem_dllm_tokens <= 0: @@ -944,6 +978,8 @@ class PrefillAdder: if self._cp_prefill_extend_limit_exceeded(trunc_len): return AddReqResult.OTHER + if self._cp_prefill_cached_limit_exceeded(prefix_len): + return AddReqResult.OTHER # Chunked prefill req.set_extend_input_len(trunc_len) diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index cf945ffdb..f43709358 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -2419,6 +2419,9 @@ class Scheduler( cp_shared_kv_prefill_max_total_extend_tokens=( self.server_args.cp_shared_kv_prefill_max_total_extend_tokens ), + cp_shared_kv_prefill_max_total_cached_tokens=( + self.server_args.cp_shared_kv_prefill_max_total_cached_tokens + ), prefill_delayer_single_pass=prefill_delayer_single_pass, dllm_config=self.dllm_config, ) @@ -2584,7 +2587,7 @@ class Scheduler( "scheduler_prefill_batch", "bs=%s extend_lens=%s prefix_lens=%s seq_lens=%s " "out_cache_tokens=%s chunked_req=%s enable_bs_gt1=%s " - "max_batch_reqs=%s max_total_extend=%s", + "max_batch_reqs=%s max_total_extend=%s max_total_cached=%s", len(can_run_list), list(getattr(new_batch, "extend_lens", []) or []), list(getattr(new_batch, "prefix_lens", []) or []), @@ -2603,6 +2606,11 @@ class Scheduler( "cp_shared_kv_prefill_max_total_extend_tokens", None, ), + getattr( + self.server_args, + "cp_shared_kv_prefill_max_total_cached_tokens", + None, + ), ) # Record prefill stats for logging after forward diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 6e100e0a2..2a4189890 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -677,6 +677,7 @@ class ServerArgs: 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 + cp_shared_kv_prefill_max_total_cached_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 @@ -993,6 +994,14 @@ class ServerArgs: "cp_shared_kv_prefill_max_total_extend_tokens must be a positive " "integer when specified." ) + if ( + self.cp_shared_kv_prefill_max_total_cached_tokens is not None + and self.cp_shared_kv_prefill_max_total_cached_tokens <= 0 + ): + raise ValueError( + "cp_shared_kv_prefill_max_total_cached_tokens must be a positive " + "integer when specified." + ) def _handle_cp_hicache_layout_validation(self): if not ( @@ -5896,6 +5905,20 @@ class ServerArgs: ) + f"\n\n{human_readable_int.__doc__}", ) + parser.add_argument( + "--cp-shared-kv-prefill-max-total-cached-tokens", + type=human_readable_int, + default=ServerArgs.cp_shared_kv_prefill_max_total_cached_tokens, + help=( + "Maximum page-aligned cached/hit tokens admitted into one NSA " + "in-seq CP shared-KV prefill batch when " + "--enable-cp-shared-kv-prefill-bs-gt1 is set. Cached tokens are " + "counted after L2 HiCache load-back planning, so this bounds " + "prefix/load-back work for high-cache-hit batches. A single " + "request larger than this limit is still allowed to run alone." + ) + + f"\n\n{human_readable_int.__doc__}", + ) parser.add_argument( "--nsa-prefill-cp-mode", type=str, diff --git a/test/registered/unit/managers/test_prefill_adder.py b/test/registered/unit/managers/test_prefill_adder.py index 7880d90c7..7778a7eca 100644 --- a/test/registered/unit/managers/test_prefill_adder.py +++ b/test/registered/unit/managers/test_prefill_adder.py @@ -774,6 +774,75 @@ class TestPrefillAdder(CustomTestCase): ) self.assertEqual([req.rid for req in adder.can_run_list], ["chunked"]) + def test_cp_prefill_total_cached_limit_stops_second_cached_request(self): + set_global_server_args_for_scheduler( + ServerArgs( + model_path="dummy", + enable_nsa_prefill_context_parallel=True, + nsa_prefill_cp_mode="in-seq-split", + ) + ) + self.mock_token_allocator.available_size.return_value = 10000 + adder = self.create_adder( + self.create_running_batch(), + page_size=64, + rem_input_tokens=4096, + enable_cp_shared_kv_prefill_bs_gt1=True, + cp_shared_kv_prefill_max_batch_requests=8, + cp_shared_kv_prefill_max_total_extend_tokens=4096, + cp_shared_kv_prefill_max_total_cached_tokens=4096, + ) + + first = self.create_prefill_req("first", extend_input_len=64) + first.prefix_indices = torch.arange(4096, dtype=torch.int64) + first.fill_ids = list(range(4096 + 64)) + second = self.create_prefill_req("second", extend_input_len=64) + second.prefix_indices = torch.arange(4096, dtype=torch.int64) + second.fill_ids = list(range(4096 + 64)) + + self.assertEqual( + adder.add_one_req(first, has_chunked_req=False, truncation_align_size=None), + AddReqResult.CONTINUE, + ) + self.assertEqual( + adder.add_one_req(second, has_chunked_req=False, truncation_align_size=None), + AddReqResult.OTHER, + ) + self.assertEqual([req.rid for req in adder.can_run_list], ["first"]) + self.assertEqual(adder.cp_shared_kv_prefill_total_cached_tokens, 4096) + + def test_cp_prefill_total_cached_limit_allows_single_oversized_cached_request(self): + set_global_server_args_for_scheduler( + ServerArgs( + model_path="dummy", + enable_nsa_prefill_context_parallel=True, + nsa_prefill_cp_mode="in-seq-split", + ) + ) + self.mock_token_allocator.available_size.return_value = 20000 + adder = self.create_adder( + self.create_running_batch(), + page_size=64, + rem_input_tokens=4096, + enable_cp_shared_kv_prefill_bs_gt1=True, + cp_shared_kv_prefill_max_batch_requests=8, + cp_shared_kv_prefill_max_total_extend_tokens=4096, + cp_shared_kv_prefill_max_total_cached_tokens=4096, + ) + + oversized = self.create_prefill_req("oversized", extend_input_len=64) + oversized.prefix_indices = torch.arange(8192, dtype=torch.int64) + oversized.fill_ids = list(range(8192 + 64)) + + self.assertEqual( + adder.add_one_req( + oversized, has_chunked_req=False, truncation_align_size=None + ), + AddReqResult.CONTINUE, + ) + self.assertEqual([req.rid for req in adder.can_run_list], ["oversized"]) + self.assertEqual(adder.cp_shared_kv_prefill_total_cached_tokens, 8192) + if __name__ == "__main__": unittest.main()