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
test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py
+109
View File
@@ -3137,6 +3137,7 @@ class UnifiedRadixCacheSuite:
storage_dir,
prefetch_policy: str = "wait_complete",
storage_extra: Optional[dict] = None,
context_length: Optional[int] = None,
):
if self.cfg.has_mamba:
self.skipTest(
@@ -3151,6 +3152,7 @@ class UnifiedRadixCacheSuite:
host_memory_mode="buffer_only",
prefetch_policy=prefetch_policy,
storage_extra=storage_extra,
context_length=context_length,
)
def _pump_hicache_until(self, cache, cond, msg, timeout: float = 10.0):
@@ -3391,6 +3393,111 @@ class UnifiedRadixCacheSuite:
self.assertIn("occupancy_ratio", cons.prefetch_outcome_stats_snapshot())
cons.sanity_check()
def test_buffer_only_storage_prefetch_miss_marker_and_retry(self):
"""A too-early storage query misses; the miss must arm the retry
marker exactly once, a re-issued check must serve once the content
lands, and abort cleanup must drop unserved markers."""
self._skip_unsupported_hicache_test()
# Marker bookkeeping is layout-independent, and each hicache fixture
# retains ~100MiB of device memory for the whole file run. Pin to one
# config so the matrix does not exhaust a small CI GPU.
if self.cfg.page_size != 1 or self.cfg.sliding_window_size != 4:
self.skipTest("requires page_size=1, sliding_window_size=4")
storage_dir = tempfile.mkdtemp()
self.addCleanup(shutil.rmtree, storage_dir, ignore_errors=True)
seq = self._buffer_swa_seq()
cons, cons_alloc, cons_rtp = build_fixture(self.cfg)
self._init_buffer_hicache(cons, storage_dir)
stats = cons._prefetch_outcome_stats
# Query BEFORE any producer wrote the span: full miss -> revoked.
req_id = "early-query-miss"
cons.prefetch_from_storage(
req_id, cons.root_node.id, array("q", seq), None, None
)
self._pump_hicache_until(
cons,
lambda: cons.check_prefetch_progress(req_id),
"miss prefetch did not resolve",
)
self.assertFalse(cons.buffer_pipeline.has_staged(req_id))
self.assertEqual(stats["revoked_full_miss"], 1)
self.assertTrue(cons.pop_storage_prefetch_miss(req_id))
self.assertFalse(cons.pop_storage_prefetch_miss(req_id)) # served once
# Producer commits the span; the re-issued check (paced retry) hits
# and stages what the first, too-early query could not see.
self._produce_buffer_l3(storage_dir, seq)
cons.prefetch_from_storage(
req_id, cons.root_node.id, array("q", seq), None, None
)
self._pump_hicache_until(
cons,
lambda: cons.check_prefetch_progress(req_id)
and cons.buffer_pipeline.has_staged(req_id),
"retried prefetch did not stage",
)
self.assertFalse(cons.pop_storage_prefetch_miss(req_id))
self.assertEqual(cons.pop_prefetch_loaded_tokens(req_id), len(seq))
# Unserved markers must not leak: abort cleanup ...
aborted_rid = "aborted-miss"
cons.prefetch_from_storage(
aborted_rid,
cons.root_node.id,
array("q", self._make_seq(700, 4)),
None,
None,
)
self._pump_hicache_until(
cons,
lambda: cons.check_prefetch_progress(aborted_rid),
"aborted-rid miss did not resolve",
)
cons.release_aborted_request(aborted_rid)
self.assertFalse(cons.pop_storage_prefetch_miss(aborted_rid))
# A fully-device-matched (empty-suffix) decline also arms the retry:
# the device match can evict while the request waits in the queue.
cons.prefetch_from_storage(
"fully-matched", cons.root_node.id, array("q", []), None, None
)
self.assertTrue(cons.pop_storage_prefetch_miss("fully-matched"))
cons.sanity_check()
def test_buffer_only_anchor_lock_cap_clamped_by_context_headroom(self):
"""Deadlock invariant: cap <= pool - context_length (floor 0), so
pinned anchors always leave room to admit the largest request."""
self._skip_unsupported_hicache_test()
# The cap is pool-size arithmetic, not a layout property; pin to one
# SWA config (which also covers the SWAKVPool full_kv_pool branch) so
# the retained per-fixture device memory stays bounded.
if self.cfg.page_size != 1 or self.cfg.sliding_window_size != 4:
self.skipTest("requires page_size=1, sliding_window_size=4")
cm = envs.SGLANG_ENABLE_HICACHE_BUFFER_ANCHOR_LOCK.override(True)
cm.__enter__()
self.addCleanup(cm.__exit__, None, None, None)
storage_dir = tempfile.mkdtemp()
self.addCleanup(shutil.rmtree, storage_dir, ignore_errors=True)
from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool
cache, _alloc, _rtp = build_fixture(self.cfg)
kvcache = cache.token_to_kv_pool_allocator.get_kvcache()
full_pool = kvcache.full_kv_pool if isinstance(kvcache, SWAKVPool) else kvcache
pool = full_pool.size
self._init_buffer_hicache(
cache, storage_dir, context_length=pool - self.cfg.page_size
)
self.assertEqual(
cache.buffer_pipeline.anchor_lock_cap_tokens, self.cfg.page_size
)
# Zero headroom -> zero cap: no pin can ever deadlock such a pool.
cache2, _alloc2, _rtp2 = build_fixture(self.cfg)
self._init_buffer_hicache(cache2, storage_dir, context_length=pool)
self.assertEqual(cache2.buffer_pipeline.anchor_lock_cap_tokens, 0)
def test_buffer_load_back_swa_window_charged_at_admission(self):
"""Admission contract: a request the SWA budget gate accepts must be
allocatable at batch time (_swa_reserved_tokens: "an admitted request
@@ -4057,6 +4164,7 @@ class UnifiedRadixCacheSuite:
prefetch_policy: str = "wait_complete",
host_memory_mode: str = "cache",
storage_extra: Optional[dict] = None,
context_length: Optional[int] = None,
):
storage_extra_config = None
if storage_backend == "file":
@@ -4093,6 +4201,7 @@ class UnifiedRadixCacheSuite:
hicache_storage_backend_extra_config=storage_extra_config,
hicache_storage_prefetch_policy=prefetch_policy,
hicache_host_memory_mode=host_memory_mode,
context_length=context_length,
)
# See build_fixture for why _mamba_cache_chunk_size is preset.
server_args._mamba_cache_chunk_size = max(FLA_CHUNK_SIZE, self.cfg.page_size)