minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[HiCache] Retry L3 storage prefetch after a missed attempt (#36227)

Browse Source
This commit is contained in:
Zhiqiang Xie
2026-08-28 01:35:21 -07:00
committed by GitHub
parent 3f1031d697
commit 1061e34785
7 changed files with 213 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/managers/schedule_batch.py
+5
View File
@@ -1029,6 +1029,11 @@ class Req(ReqDllmMixin):
self.num_matched_prefix_tokens = 0
# Tokens loaded from storage backend (L3) during prefetch for this request
self.storage_hit_length = 0
# Storage prefetch retry state while queued
# (see Scheduler._retry_missed_storage_prefetches).
self.storage_prefetch_retry_pending = False
self.storage_prefetch_retry_wait_polls = 0
self.storage_prefetch_retry_attempts = 0
# The node to lock until for swa radix tree lock ref
self.swa_uuid_for_lock: Optional[int] = None
# Whether the prefill-time SWA tree lock has been released early
python/sglang/srt/managers/scheduler.py
+32
View File
@@ -2847,6 +2847,36 @@ class Scheduler(
matched_prefix_tokens=req.full_untruncated_fill_ids[:matched_len],
)
def _retry_missed_storage_prefetches(self):
"""Re-issue the availability check for queued requests whose prefetch
missed. Pacing counts scheduling passes so TP ranks re-issue on the
same pass; a sweep (the admission loop stops at the first
unschedulable request) covers the whole queue."""
interval = get_memory().hicache_storage_prefetch_retry_poll_interval
if interval <= 0 or not self.waiting_queue:
return
max_attempts = get_memory().hicache_storage_prefetch_retry_max_attempts
for req in self.waiting_queue:
if self.tree_cache.pop_storage_prefetch_miss(req.rid):
req.storage_prefetch_retry_pending = True
req.storage_prefetch_retry_wait_polls = 0
if (
not req.storage_prefetch_retry_pending
or req.storage_prefetch_retry_attempts >= max_attempts
):
continue
req.storage_prefetch_retry_wait_polls += 1
if req.storage_prefetch_retry_wait_polls <= interval:
continue
req.storage_prefetch_retry_pending = False
req.storage_prefetch_retry_attempts += 1
logger.debug(
"HiCache storage prefetch retry req=%s attempt=%d",
req.rid,
req.storage_prefetch_retry_attempts,
)
self._prefetch_kvcache(req)
def _add_request_to_queue(self, req: Req, is_retracted: bool = False):
if not self._set_or_validate_priority(req):
return
@@ -3353,6 +3383,8 @@ class Scheduler(
if self.enable_hierarchical_cache or get_memory().enable_flexkv:
self.tree_cache.check_hicache_events()
if self.enable_hicache_storage:
self._retry_missed_storage_prefetches()
if self.enable_priority_preemption or self.is_hybrid_swa:
# Reset batch_is_full to try preemption with a prefill adder.
python/sglang/srt/mem_cache/buffer_mode/pipeline.py
+15 -2
View File
@@ -199,6 +199,7 @@ class BufferModePipeline:
cache: UnifiedRadixCache,
swa_window_pages: int,
write_backlog_cap: int,
max_context_len: int = 0,
):
self._cache = cache
# SWA window size in KV pages when the SWA component stages through
@@ -214,8 +215,20 @@ class BufferModePipeline:
kvcache = cache.token_to_kv_pool_allocator.get_kvcache()
full_pool = kvcache.full_kv_pool if isinstance(kvcache, SWAKVPool) else kvcache
self.anchor_lock_cap_tokens = int(
envs.SGLANG_HICACHE_BUFFER_ANCHOR_LOCK_CAP.get() * full_pool.size
# Clamp by admission headroom: pins must leave room for the largest
# allowed request, else a queued hold can wedge admission permanently
# (pool full, nothing retractable). No-headroom pools take no pins.
self.anchor_lock_cap_tokens = max(
0,
min(
int(envs.SGLANG_HICACHE_BUFFER_ANCHOR_LOCK_CAP.get() * full_pool.size),
full_pool.size - max_context_len,
),
)
logger.info(
"BufferModePipeline anchor_lock_enabled=%s cap_tokens=%d",
self.anchor_lock_enabled,
self.anchor_lock_cap_tokens,
)
self.reset()
python/sglang/srt/mem_cache/hiradix_cache.py
+5
View File
@@ -1733,6 +1733,11 @@ class HiRadixCache(RadixCache):
"""
return self.prefetch_loaded_tokens_by_reqid.pop(req_id, 0)
def pop_storage_prefetch_miss(self, req_id: str) -> bool:
"""Storage prefetch miss markers are not tracked on the dense path;
the scheduler's paced availability-check retry is inert here."""
return False
def match_prefix(self, params: MatchPrefixParams):
if self.disable:
return self._empty_match_result
python/sglang/srt/mem_cache/unified_radix_cache.py
+30 -1
View File
@@ -351,6 +351,9 @@ class UnifiedRadixCache(BasePrefixCache):
self.enable_storage = False
self.prefetch_loaded_tokens_by_reqid: dict[str, int] = {}
self.ongoing_prefetch: dict[str, _OngoingPrefetch] = {}
# Rids whose storage prefetch resolved without a usable result;
# popped by the scheduler to pace availability-check retries.
self._storage_prefetch_missed_rids: set[str] = set()
self.ongoing_backup: dict[int, tuple[NodeId, DecLockRefParams]] = {}
if self.buffer_pipeline is not None:
self.buffer_pipeline.reset()
@@ -427,6 +430,7 @@ class UnifiedRadixCache(BasePrefixCache):
)
self.buffer_pipeline = BufferModePipeline(
cache=self,
max_context_len=server_args.context_length or 0,
swa_window_pages=(
swa.full_window_pages
if swa is not None and self.tree_core.has_swa_host_pool
@@ -1673,9 +1677,13 @@ class UnifiedRadixCache(BasePrefixCache):
if prefetch_length < self.prefetch_threshold:
if prefetch_length > 0:
stats["declined_too_short"] += 1
# A too-short/fully-matched suffix can become a full recompute if
# the device match evicts while queued; arm the paced retry.
self._storage_prefetch_missed_rids.add(req_id)
return
if not buffer_mode and self.cache_controller.prefetch_rate_limited():
stats["declined_rate_limited"] += 1
self._storage_prefetch_missed_rids.add(req_id)
return
if req_id in self.ongoing_prefetch or (
buffer_mode and self.buffer_pipeline.has_staged(req_id)
@@ -1736,6 +1744,8 @@ class UnifiedRadixCache(BasePrefixCache):
)
if anchor_lock_params is not None:
self.dec_host_lock_ref(last_host_node_id, anchor_lock_params)
# Forfeited over transient staging pressure; retryable.
self._storage_prefetch_missed_rids.add(req_id)
return
aux_xfers = [x for xfers in comp_xfers.values() for x in xfers]
@@ -1764,6 +1774,10 @@ class UnifiedRadixCache(BasePrefixCache):
)
if buffer_mode:
self.buffer_pipeline.set_prefix_ctx(req_id, matched_prefix_tokens)
# Pin the just-matched anchor now: deferred to hit-alloc it is
# often already deleted under churn (silent unlocked launch).
# The hit-alloc call remains as an idempotent second chance.
self.buffer_pipeline.try_lock_anchor(req_id, last_host_node_id)
else:
# Cache mode reserves the requested span up front; buffer mode
# grants occupancy later at hit-alloc time, sized to the hit.
@@ -1811,6 +1825,7 @@ class UnifiedRadixCache(BasePrefixCache):
self.cache_controller.terminate_prefetch(operation)
if operation.host_indices is None:
self._storage_prefetch_missed_rids.add(req_id)
self.revoke_pending_prefetch(req_id)
else:
self._handle_prefetch_result(operation)
@@ -2006,8 +2021,18 @@ class UnifiedRadixCache(BasePrefixCache):
return True
def pop_prefetch_loaded_tokens(self, req_id: str) -> int:
# The request is being scheduled; a still-unserved miss marker is moot.
self._storage_prefetch_missed_rids.discard(req_id)
return self.prefetch_loaded_tokens_by_reqid.pop(req_id, 0)
def pop_storage_prefetch_miss(self, req_id: str) -> bool:
"""True once per resolved storage-prefetch miss for a live request;
the scheduler uses it to arm the paced availability-check retry."""
if req_id in self._storage_prefetch_missed_rids:
self._storage_prefetch_missed_rids.discard(req_id)
return True
return False
def staged_prefetch_tokens(self, req_id: str) -> int:
"""Tokens a staged buffer-mode prefetch would splice (0 = no hold);
surfaced by the scheduler as the request's host_hit_length."""
@@ -2025,6 +2050,7 @@ class UnifiedRadixCache(BasePrefixCache):
@rank_consensus(same_params=True)
def release_aborted_request(self, rid: str) -> None:
self.prefetch_loaded_tokens_by_reqid.pop(rid, None)
self._storage_prefetch_missed_rids.discard(rid)
if (
self.buffer_pipeline is not None
and self.buffer_pipeline.release_aborted_staged(rid)
@@ -2250,13 +2276,16 @@ class UnifiedRadixCache(BasePrefixCache):
self._invalidate_absent_from_hit_query(operation)
continue
if operation.is_terminated():
# Aborted while the storage query was in flight.
# Controller-side miss termination (retryable) or an abort
# race (abort cleanup discards the marker).
self._storage_prefetch_missed_rids.add(req_id)
self.revoke_pending_prefetch(req_id)
continue
if operation.storage_hit_count < self.prefetch_threshold:
# Below-threshold hit: classify + feed the L3 miss
# accounting, then revoke (not enough benefit).
self._account_prefetch_outcome(operation, revoked=True)
self._storage_prefetch_missed_rids.add(req_id)
self.revoke_pending_prefetch(req_id)
continue
self._invalidate_absent_from_hit_query(operation)
python/sglang/srt/server_args.py
+17
View File
@@ -2820,6 +2820,23 @@ class ServerArgs:
"A dictionary in JSON string format, or a string starting with a leading '@' and a config file in JSON/YAML/TOML format, containing extra configuration for the storage backend.",
NS("memory"),
] = None
hicache_storage_prefetch_retry_poll_interval: A[
int,
Arg(
help=(
"Scheduling passes a queued request waits after a storage "
"prefetch miss before the availability check is retried "
"(under load the first check can run before the needed "
"backup commits). 0 disables retries."
),
),
NS("memory"),
] = 0
hicache_storage_prefetch_retry_max_attempts: A[
int,
"Maximum storage prefetch retries per request when --hicache-storage-prefetch-retry-poll-interval is set.",
NS("memory"),
] = 4
# -------------------------------------------------------------------------
# Hierarchical sparse attention