From 1061e34785ae0264c09d98b50690903dbf3be7aa Mon Sep 17 00:00:00 2001 From: Zhiqiang Xie Date: Fri, 28 Aug 2026 01:35:21 -0700 Subject: [PATCH] [HiCache] Retry L3 storage prefetch after a missed attempt (#36227) --- python/sglang/srt/managers/schedule_batch.py | 5 + python/sglang/srt/managers/scheduler.py | 32 +++++ .../srt/mem_cache/buffer_mode/pipeline.py | 17 ++- python/sglang/srt/mem_cache/hiradix_cache.py | 5 + .../srt/mem_cache/unified_radix_cache.py | 31 ++++- python/sglang/srt/server_args.py | 17 +++ .../test_unified_radix_cache_unittest.py | 109 ++++++++++++++++++ 7 files changed, 213 insertions(+), 3 deletions(-) diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 40a3c4d1b..81ba9403a 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -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 diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 38f13e95a..4250cbe93 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -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. diff --git a/python/sglang/srt/mem_cache/buffer_mode/pipeline.py b/python/sglang/srt/mem_cache/buffer_mode/pipeline.py index a2facd33c..91bcc9295 100644 --- a/python/sglang/srt/mem_cache/buffer_mode/pipeline.py +++ b/python/sglang/srt/mem_cache/buffer_mode/pipeline.py @@ -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() diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index 343571225..6f82bc810 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/python/sglang/srt/mem_cache/hiradix_cache.py @@ -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 diff --git a/python/sglang/srt/mem_cache/unified_radix_cache.py b/python/sglang/srt/mem_cache/unified_radix_cache.py index 1c91adfed..657ba1c8c 100644 --- a/python/sglang/srt/mem_cache/unified_radix_cache.py +++ b/python/sglang/srt/mem_cache/unified_radix_cache.py @@ -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) diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 56ded1a80..124aec7ed 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -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 diff --git a/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py b/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py index ff6c9ee53..d0a85179c 100644 --- a/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py +++ b/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py @@ -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)