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

Bound CP prefill batches by cached-token pressure

Browse Source 
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.
This commit is contained in:
laoyao0822
2026-06-10 22:21:46 +08:00
parent 342c552ab3
commit 4f65d7a176
4 changed files with 137 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

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

10
python/sglang/srt/managers/scheduler.py
View File

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

23
python/sglang/srt/server_args.py
View File

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

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

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