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

Keep CP HiCache free room from collapsing under cache reuse

Browse Source 
Repeated GSM8K/cache-hit traffic can drain a CP owner lane while exact
allocation still succeeds. This changes L1/L2 free-room logic so the trigger
accounts for the pending allocation, then evicts a minimum target chunk or the
exact shortage, whichever is larger. CP load-back also performs best-effort
owner-lane free-room eviction before admitting tiny cache hits that would
otherwise skew one lane.

The commit also adds gated bs>1 prefill timing markers around
recv/bootstrap/batch/inflight polling so future hangs expose the scheduler
boundary instead of only surfacing as a watchdog shutdown. The post-fix repeat
GSM8K failure is recorded as an active regression to continue investigating,
not as old-process noise.

Constraint: Free-room policy must reduce repeated eviction and owner-lane starvation without reserving required+target pages after every allocation
Rejected: Evict to required+target availability | wastes L1/L2 residency under many short cache hits
Rejected: Treat free-room misses as fatal on load-back | exact capacity should remain the strict admission condition
Confidence: medium
Scope-risk: moderate
Directive: Do not remove the gated bs>1 timing markers until repeat-GSM8K hangs have a direct failing boundary
Tested: git diff --check
Tested: python -m py_compile for touched Python files
Tested: remote earlier test_cp_shared_kv_layout.py and test_cp_hicache_metadata.py passed 157 tests after this free-room formula
Not-tested: Two full repeat GSM8K runs after this commit; latest repeat still killed prefill and requires follow-up root-cause work
This commit is contained in:
laoyao0822
2026-06-07 01:23:34 +08:00
parent 68defa3923
commit b17976b60d
9 changed files with 642 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/disaggregation/prefill.py
+178 -4
View File
@@ -20,6 +20,7 @@ Life cycle of a request in the prefill server
from __future__ import annotations
import logging
import time
from collections import deque
from http import HTTPStatus
from typing import TYPE_CHECKING, List, Optional
@@ -86,6 +87,59 @@ def _cp_shared_kv_bs_gt1_prefill_debug(
logger.info("[CP_SHARED_KV_BS_GT1_DEBUG] event=%s " + message, key, *args)
_CP_SHARED_KV_BS_GT1_PREFILL_TIMING_COUNTS = {}
def _cp_shared_kv_bs_gt1_prefill_timing_start() -> Optional[float]:
if not envs.SGLANG_CP_SHARED_KV_BS_GT1_TIMING.get():
return None
return time.perf_counter()
def _cp_shared_kv_bs_gt1_prefill_timing(
key: str,
start_time: Optional[float],
message: str,
*args,
) -> None:
if start_time is None or not envs.SGLANG_CP_SHARED_KV_BS_GT1_TIMING.get():
return
elapsed_ms = (time.perf_counter() - start_time) * 1000.0
slow_ms = float(envs.SGLANG_CP_SHARED_KV_BS_GT1_TIMING_SLOW_MS.get())
if slow_ms > 0 and elapsed_ms < slow_ms:
return
limit = int(envs.SGLANG_CP_SHARED_KV_BS_GT1_TIMING_LIMIT.get())
count = _CP_SHARED_KV_BS_GT1_PREFILL_TIMING_COUNTS.get(key, 0)
if limit > 0 and count >= limit:
return
_CP_SHARED_KV_BS_GT1_PREFILL_TIMING_COUNTS[key] = count + 1
logger.info(
"[CP_SHARED_KV_BS_GT1_TIMING] event=%s elapsed_ms=%.3f " + message,
key,
elapsed_ms,
*args,
)
def _cp_shared_kv_bs_gt1_prefill_marker(
key: str,
message: str,
*args,
) -> None:
if not envs.SGLANG_CP_SHARED_KV_BS_GT1_TIMING.get():
return
limit = int(envs.SGLANG_CP_SHARED_KV_BS_GT1_TIMING_LIMIT.get())
count = _CP_SHARED_KV_BS_GT1_PREFILL_TIMING_COUNTS.get(key, 0)
if limit > 0 and count >= limit:
return
_CP_SHARED_KV_BS_GT1_PREFILL_TIMING_COUNTS[key] = count + 1
logger.info(
"[CP_SHARED_KV_BS_GT1_TIMING] event=%s elapsed_ms=0.000 " + message,
key,
*args,
)
def _seq_summary(values) -> str:
if values is None:
return "None"
@@ -458,11 +512,26 @@ class PrefillBootstrapQueue:
else:
return [], []
poll_start = _cp_shared_kv_bs_gt1_prefill_timing_start()
_cp_shared_kv_bs_gt1_prefill_marker(
"bootstrap_poll_start",
"queue=%s return_failed=%s rids_to_check=%s",
len(self.queue),
return_failed_reqs,
len(rids_to_check) if rids_to_check is not None else None,
)
polls = poll_and_all_reduce_attn_cp_tp_group(
[req.disagg_kv_sender for req in self.queue],
self.scheduler.attn_cp_cpu_group,
self.scheduler.attn_tp_cpu_group,
)
_cp_shared_kv_bs_gt1_prefill_timing(
"bootstrap_poll_done",
poll_start,
"queue=%s polls_head=%s",
len(self.queue),
polls[:8],
)
for i, (req, poll) in enumerate(zip(self.queue, polls)):
if rids_to_check is not None:
@@ -549,24 +618,114 @@ class SchedulerDisaggregationPrefillMixin:
while True:
# Receive requests
recv_reqs = self.recv_requests()
self.process_input_requests(recv_reqs)
self.waiting_queue.extend(
self.disagg_prefill_bootstrap_queue.pop_bootstrapped()
recv_start = _cp_shared_kv_bs_gt1_prefill_timing_start()
_cp_shared_kv_bs_gt1_prefill_marker(
"event_loop_recv_start",
"waiting=%s inflight=%s bootstrap=%s",
len(self.waiting_queue),
len(self.disagg_prefill_inflight_queue),
len(self.disagg_prefill_bootstrap_queue.queue),
)
recv_reqs = self.recv_requests()
_cp_shared_kv_bs_gt1_prefill_timing(
"event_loop_recv_done",
recv_start,
"recv=%s waiting=%s inflight=%s bootstrap=%s",
len(recv_reqs),
len(self.waiting_queue),
len(self.disagg_prefill_inflight_queue),
len(self.disagg_prefill_bootstrap_queue.queue),
)
self.process_input_requests(recv_reqs)
pop_start = _cp_shared_kv_bs_gt1_prefill_timing_start()
bootstrapped_reqs = self.disagg_prefill_bootstrap_queue.pop_bootstrapped()
_cp_shared_kv_bs_gt1_prefill_timing(
"event_loop_pop_bootstrapped_done",
pop_start,
"bootstrapped=%s waiting_before=%s inflight=%s bootstrap_remaining=%s",
len(bootstrapped_reqs),
len(self.waiting_queue),
len(self.disagg_prefill_inflight_queue),
len(self.disagg_prefill_bootstrap_queue.queue),
)
self.waiting_queue.extend(bootstrapped_reqs)
# Get the next batch to run
batch_start = _cp_shared_kv_bs_gt1_prefill_timing_start()
_cp_shared_kv_bs_gt1_prefill_marker(
"event_loop_get_batch_start",
"waiting=%s inflight=%s bootstrap=%s",
len(self.waiting_queue),
len(self.disagg_prefill_inflight_queue),
len(self.disagg_prefill_bootstrap_queue.queue),
)
batch = self.get_next_disagg_prefill_batch_to_run()
_cp_shared_kv_bs_gt1_prefill_timing(
"event_loop_get_batch_done",
batch_start,
"has_batch=%s batch_size=%s waiting=%s inflight=%s bootstrap=%s",
batch is not None,
len(batch.reqs) if batch is not None else 0,
len(self.waiting_queue),
len(self.disagg_prefill_inflight_queue),
len(self.disagg_prefill_bootstrap_queue.queue),
)
self.cur_batch = batch
# Launch the current batch
if batch:
run_start = _cp_shared_kv_bs_gt1_prefill_timing_start()
_cp_shared_kv_bs_gt1_prefill_marker(
"event_loop_run_batch_start",
"batch_size=%s extend_lens=%s prefix_lens=%s inflight=%s",
len(batch.reqs),
_seq_summary(getattr(batch, "extend_lens", None)),
_seq_summary(getattr(batch, "prefix_lens", None)),
len(self.disagg_prefill_inflight_queue),
)
result = self.run_batch(batch)
_cp_shared_kv_bs_gt1_prefill_timing(
"event_loop_run_batch_done",
run_start,
"batch_size=%s inflight=%s",
len(batch.reqs),
len(self.disagg_prefill_inflight_queue),
)
result_start = _cp_shared_kv_bs_gt1_prefill_timing_start()
_cp_shared_kv_bs_gt1_prefill_marker(
"event_loop_process_result_start",
"batch_size=%s inflight_before=%s",
len(batch.reqs),
len(self.disagg_prefill_inflight_queue),
)
self.process_batch_result(batch, result)
_cp_shared_kv_bs_gt1_prefill_timing(
"event_loop_process_result_done",
result_start,
"batch_size=%s inflight_after=%s",
len(batch.reqs),
len(self.disagg_prefill_inflight_queue),
)
else:
self.self_check_during_idle()
inflight_start = _cp_shared_kv_bs_gt1_prefill_timing_start()
_cp_shared_kv_bs_gt1_prefill_marker(
"event_loop_process_inflight_start",
"inflight=%s waiting=%s bootstrap=%s",
len(self.disagg_prefill_inflight_queue),
len(self.waiting_queue),
len(self.disagg_prefill_bootstrap_queue.queue),
)
self.process_disagg_prefill_inflight_queue()
_cp_shared_kv_bs_gt1_prefill_timing(
"event_loop_process_inflight_done",
inflight_start,
"inflight=%s waiting=%s bootstrap=%s",
len(self.disagg_prefill_inflight_queue),
len(self.waiting_queue),
len(self.disagg_prefill_bootstrap_queue.queue),
)
# Update last_batch
self.last_batch = batch
@@ -766,11 +925,26 @@ class SchedulerDisaggregationPrefillMixin:
done_reqs = []
poll_start = _cp_shared_kv_bs_gt1_prefill_timing_start()
_cp_shared_kv_bs_gt1_prefill_marker(
"inflight_poll_start",
"inflight=%s rids_head=%s rids_to_check=%s",
len(self.disagg_prefill_inflight_queue),
[req.rid for req in self.disagg_prefill_inflight_queue[:8]],
len(rids_to_check) if rids_to_check is not None else None,
)
polls = poll_and_all_reduce_attn_cp_tp_group(
[req.disagg_kv_sender for req in self.disagg_prefill_inflight_queue],
self.attn_cp_cpu_group,
self.attn_tp_cpu_group,
)
_cp_shared_kv_bs_gt1_prefill_timing(
"inflight_poll_done",
poll_start,
"inflight=%s polls_head=%s",
len(self.disagg_prefill_inflight_queue),
polls[:8],
)
undone_reqs: List[Req] = []
# Check .poll() for the reqs in disagg_prefill_inflight_queue. If Success, respond to the client and remove it from the queue