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:15 +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

182
python/sglang/srt/disaggregation/prefill.py
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

5
python/sglang/srt/mem_cache/allocator.py
View File

@@ -57,10 +57,11 @@ def compute_owner_lane_free_room_deficits(
if capacity > 0 and trigger_ratio > 0
else 0
)
if int(avail) >= int(req) + trigger_room:
trigger_available = int(req) + trigger_room
if int(avail) >= trigger_available:
deficits.append(0)
else:
deficits.append(max(0, int(req) + target_room - int(avail)))
deficits.append(max(max(0, int(req) - int(avail)), target_room))
return deficits

26
python/sglang/srt/mem_cache/common.py
View File

@@ -444,6 +444,18 @@ def _evict_for_compute_owner_lanes(
return
def _cp_owner_lane_free_room_enabled(tree_cache: BasePrefixCache | None) -> bool:
if tree_cache is None:
return False
return (
float(getattr(tree_cache, "hicache_l1_free_room_ratio", 0.0) or 0.0) > 0.0
or float(
getattr(tree_cache, "hicache_l1_free_room_trigger_ratio", 0.0) or 0.0
)
> 0.0
)
def alloc_paged_token_slots_extend(
tree_cache: BasePrefixCache,
prefix_lens: torch.Tensor,
@@ -540,6 +552,20 @@ def alloc_paged_token_slots_extend(
if backup_state:
state = allocator.backup_state()
# Maintain the configured L1 owner-lane free room on successful
# allocations as well. Previously this eviction path only ran after
# exact allocation failure, so cache hits and short extends could drain
# one CP owner lane down to a handful of pages while still "succeeding",
# leaving later batches to stall under skewed owner-lane pressure.
if _cp_owner_lane_free_room_enabled(tree_cache):
_evict_for_compute_owner_lanes(
tree_cache=tree_cache,
allocator=allocator,
page_compute_owners=page_compute_owners,
)
if backup_state:
state = allocator.backup_state()
out_cache_loc = alloc_extend_compute_owner(
prefix_lens,
prefix_lens_cpu,

33
python/sglang/srt/mem_cache/hiradix_cache.py
View File

@@ -152,7 +152,7 @@ def _free_room_deficit(
trigger_room = _page_aligned_room_from_ratio(capacity, trigger_ratio, page_size)
if int(available) >= int(required) + trigger_room:
return 0
return max(0, int(required) + target_room - int(available))
return max(max(0, int(required) - int(available)), target_room)
@dataclass
@@ -3438,6 +3438,37 @@ class HiRadixCache(RadixCache):
)
record_stage("owner_lane_evict")
if not any(v > 0 for v in load_back_plan.deficit_by_owner) and any(
v > 0 for v in load_back_plan.free_room_deficit_by_owner
):
# Best-effort L1 free-room maintenance for CP load-back. Exact
# admission may still succeed with a nearly empty owner lane
# (common for one-page host hits that all map to owner 0), but
# letting those loads proceed without replenishing free room can
# starve the next extend batch. Free-room target misses are not
# fatal; exact deficits remain the strict condition below.
free_room_deficits = [
int(v) for v in load_back_plan.free_room_deficit_by_owner
]
free_room_pages = sum(max(0, v) for v in free_room_deficits)
if free_room_pages > 0:
evict_result = self.evict(
EvictParams(
num_tokens=free_room_pages * self.page_size,
owner_lane_deficits=free_room_deficits,
)
)
logger.debug(
"[HiCache-load] owner-lane free-room eviction before CP "
"load-back: node_id=%d free_room_deficit_by_owner=%s "
"num_tokens_evicted=%d",
last_hit_node.id,
free_room_deficits,
getattr(evict_result, "num_tokens_evicted", 0),
)
load_back_plan = self._refresh_cp_load_back_plan(load_back_plan)
record_stage("owner_lane_free_room_evict")
if any(v > 0 for v in load_back_plan.deficit_by_owner):
self.dec_lock_ref(ancester_node)
logger.warning(