From a328c19c81894b2093211bc089215048faa4320a Mon Sep 17 00:00:00 2001 From: Zhiqiang Xie Date: Sat, 29 Aug 2026 01:51:47 -0700 Subject: [PATCH] [HiCache] buffer mode: decide staged-fetch fate against the live tree (#36834) --- python/sglang/srt/managers/scheduler.py | 12 +- .../srt/mem_cache/buffer_mode/pipeline.py | 201 +++++++++++---- .../srt/mem_cache/unified_radix_cache.py | 54 ++-- .../test_unified_radix_cache_unittest.py | 236 ++++++++++++++++++ 4 files changed, 429 insertions(+), 74 deletions(-) diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 77728521c..e9a42dabf 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -3549,16 +3549,18 @@ class Scheduler( # Buffer mode: surface a staged prefetch as the request's host # hit (consumed through init_load_back) plus its SWA window, # which consumption allocates and the request lock pins — - # uncharged, the batch alloc can OOM. Set AFTER + # uncharged, the batch alloc can OOM. Planned against the same + # live prefix admission uses, so only the splice-able span + # tail is charged and unusable holds are freed. Set AFTER # init_next_round_input (which recomputes host_hit). Mamba # (fenced in init_hicache) will need the same charge via # mamba_host_hit_length. - held_tokens = self.tree_cache.staged_prefetch_tokens(req.rid) + held_tokens, held_swa_tokens = self.tree_cache.plan_staged_splice( + req.rid, len(req.prefix_indices) + ) if held_tokens > 0: req.host_hit_length = held_tokens - req.swa_host_hit_length = ( - self.tree_cache.staged_prefetch_swa_tokens(req.rid) - ) + req.swa_host_hit_length = held_swa_tokens res = adder.add_one_req( req, has_chunked_req=(self.chunked_req is not None), diff --git a/python/sglang/srt/mem_cache/buffer_mode/pipeline.py b/python/sglang/srt/mem_cache/buffer_mode/pipeline.py index 3986893fb..efa532772 100644 --- a/python/sglang/srt/mem_cache/buffer_mode/pipeline.py +++ b/python/sglang/srt/mem_cache/buffer_mode/pipeline.py @@ -148,6 +148,21 @@ def _untrack_content_refs(refs: dict[str, int], hash_values: list[str]) -> None: refs[h] = n +def staged_splice_tokens(f: _StagedPrefetch, device_prefix_len: int) -> int: + """Tokens a staged prefetch can still splice beyond the live device + prefix; 0 = unusable hold (prefix shrunk below the span, span fully + device-resident, or the trim would cut into a staged aux trailing + window — aux pools splice whole or not at all).""" + span_end = f.matched_len + f.num_tokens + if device_prefix_len < f.matched_len or device_prefix_len >= span_end: + return 0 + splice_tokens = span_end - device_prefix_len + for t in f.aux_xfers: + if t.host_indices is not None and t.host_indices.numel() > splice_tokens: + return 0 + return splice_tokens + + def validate_buffer_only_stack( sidecar_pool_specs: list, swa_component: Optional[SWAComponent] ) -> None: @@ -239,7 +254,9 @@ class BufferModePipeline: # prefill admission, and load-backs in flight (keyed by synthetic # negative ack id). self.pending_hit_allocs: deque = deque() - self._prefetch_prefix_ctx: dict[str, list[int]] = {} + self._prefetch_prefix_ctx: dict[ + str, tuple[list[int], Optional[str], Optional[str]] + ] = {} self.staged_prefetches: dict[str, _StagedPrefetch] = {} self.ongoing_buffer_load_back: dict[int, _OngoingBufferLoadBack] = {} # Backup pipeline: FIFO intents awaiting a D2H slot, node ids @@ -622,15 +639,21 @@ class BufferModePipeline: # ---- load back pipeline (storage -> staging -> device) ---- - def try_lock_anchor(self, req_id: str, anchor_node_id: NodeId) -> None: - """Pin the device anchor at IO commit so eviction cannot invalidate - the splice; drift resolves at consumption, and the cap keeps queued - holds from making the pool unevictable (over-cap launches unlocked).""" - if not self.anchor_lock_enabled or req_id in self.anchor_locks: - return - prefix_tokens = self._prefetch_prefix_ctx.get(req_id) - if not prefix_tokens: - return # root anchor: nothing to pin + def try_lock_anchor(self, req_id: str) -> str: + """Pin the staged prefetch's device anchor so eviction cannot + invalidate the splice, finding it by re-matching the live tree + (carried node ids go stale via splits and eviction; the walk is + O(prefix path)). Returns "locked", "no_anchor" (nothing to pin), + "cap_skip" (over cap; launches unlocked), or "anchor_lost" (splice + base gone — the caller cancels the storage IO).""" + if not self.anchor_lock_enabled: + return "no_anchor" + if req_id in self.anchor_locks: + return "locked" + prefix_ctx = self._prefetch_prefix_ctx.get(req_id) + if not prefix_ctx or not prefix_ctx[0]: + return "no_anchor" # root anchor: nothing to pin + prefix_tokens, extra_key, cache_salt = prefix_ctx matched_len = len(prefix_tokens) if self.anchor_locked_tokens_ + matched_len > self.anchor_lock_cap_tokens: self._anchor_lock_cap_skips += 1 @@ -646,23 +669,28 @@ class BufferModePipeline: matched_len, self.anchor_lock_cap_tokens, ) - return + return "cap_skip" cache = self._cache - try: - node = cache.tree_core.node_by_id(anchor_node_id) - except KeyError: - return # anchor deleted; fetch unlocked - if node.component_data[BASE_COMPONENT_TYPE].value is None: - # Evicted since enqueue; fetch unlocked. - logger.warning("HiCache anchor evicted before IO commit req=%s", req_id) - return - lock_params = cache.inc_lock_ref(anchor_node_id).to_dec_params() + match = cache.match_prefix( + MatchPrefixParams( + key=RadixKey( + array("q", prefix_tokens), + extra_key=extra_key, + is_bigram=cache.tree_core.is_eagle, + cache_salt=cache_salt, + ) + ) + ) + if len(match.device_indices) < matched_len: + return "anchor_lost" + lock_params = cache.inc_lock_ref(match.last_device_node).to_dec_params() self.anchor_locks[req_id] = _AnchorLock( - node_id=anchor_node_id, + node_id=match.last_device_node, lock_params=lock_params, tokens=matched_len, ) self.anchor_locked_tokens_ += matched_len + return "locked" def release_anchor_lock(self, req_id: str) -> None: """Drop a staged prefetch's anchor lock (idempotent; called at every @@ -677,10 +705,42 @@ class BufferModePipeline: f"after releasing {req_id}" ) - def set_prefix_ctx(self, req_id: str, matched_prefix_tokens) -> None: - """Record the device-matched prefix at prefetch enqueue; consumed at - staging commit to build the full-span tree key.""" - self._prefetch_prefix_ctx[req_id] = list(matched_prefix_tokens or []) + def staged_span_covered(self, req_id: str, span_tokens: int) -> bool: + """True when the live device tree already covers the fetch's whole + would-be span (prefix + the storage-hit tokens): nothing would be + left to splice at consumption, so the IO-commit caller cancels + before the bounce alloc and the storage read.""" + info = self._cache.ongoing_prefetch.get(req_id) + if info is None or span_tokens <= 0: + return False + prefix_tokens, _, _ = self._prefetch_prefix_ctx[req_id] + span_key = info.prefetch_key + full_tokens = array("q", prefix_tokens) + full_tokens.extend(span_key[:span_tokens].token_ids) + key = RadixKey( + full_tokens, + extra_key=span_key.extra_key, + is_bigram=self._cache.tree_core.is_eagle, + cache_salt=span_key.cache_salt, + ) + match = self._cache.match_prefix(MatchPrefixParams(key=key)) + return len(match.device_indices) >= len(key) + + def set_prefix_ctx( + self, + req_id: str, + matched_prefix_tokens, + extra_key: Optional[str] = None, + cache_salt: Optional[str] = None, + ) -> None: + """Record the device-matched prefix (and its tree-key namespace) at + prefetch enqueue; consumed at staging commit to build the full-span + tree key, and by try_lock_anchor to re-match a stale anchor.""" + self._prefetch_prefix_ctx[req_id] = ( + list(matched_prefix_tokens or []), + extra_key, + cache_salt, + ) def pop_prefix_ctx(self, req_id: str) -> None: self._prefetch_prefix_ctx.pop(req_id, None) @@ -713,7 +773,8 @@ class BufferModePipeline: comp_xfers, ) = cache.ongoing_prefetch.pop(req_id) cc = cache.cache_controller - prefix_tokens = self._prefetch_prefix_ctx.pop(req_id, None) + prefix_ctx = self._prefetch_prefix_ctx.pop(req_id, None) + prefix_tokens = prefix_ctx[0] if prefix_ctx is not None else None aux_xfers = [x for xfers in comp_xfers.values() for x in xfers] if num_tokens == 0 or prefix_tokens is None: @@ -751,11 +812,29 @@ class BufferModePipeline: cache.prefetch_loaded_tokens_by_reqid[req_id] = num_tokens return True - def staged_prefetch_tokens(self, req_id: str) -> int: - """Tokens a staged prefetch would splice (0 = no hold); surfaced by the - scheduler as the request's host_hit_length.""" + def plan_staged_splice( + self, req_id: str, device_prefix_len: int + ) -> tuple[int, int]: + """(kv, swa) host-hit tokens consumption will splice given the + request's live device prefix, so admission charges no phantom + tokens. Frees a hold that can no longer splice: surfaced as 0 but + kept, it would leak — the adder only consumes surfaced host hits.""" f = self.staged_prefetches.get(req_id) - return f.num_tokens if f is not None else 0 + if f is None: + return 0, 0 + splice_tokens = staged_splice_tokens(f, device_prefix_len) + if splice_tokens == 0: + logger.info( + "HiCache staged prefetch released req=%s matched=%d " + "device_prefix=%d tokens=%d", + req_id, + f.matched_len, + device_prefix_len, + f.num_tokens, + ) + self.release_staged_hold(req_id) + return 0, 0 + return splice_tokens, self.staged_prefetch_swa_tokens(req_id) def staged_prefetch_swa_tokens(self, req_id: str) -> int: """SWA device tokens consuming this staged prefetch will allocate (the @@ -773,7 +852,9 @@ class BufferModePipeline: def init_load_back(self, params: InitLoadBackParams) -> tuple[torch.Tensor, NodeId]: """Consume the staged prefetch at prefill admission: device alloc, layer-gated H2D, and a plain insert so downstream sees ordinary tree - state; invalid holds drop and the request recomputes. + state. The splice base is the request's live device prefix — growth + trims to the span tail beyond it; unusable holds drop and the + request recomputes. Ownership contract: cc.load queues the H2D before insert adjudicates ownership, so the live pre-checks below must prove the insert can @@ -794,6 +875,9 @@ class BufferModePipeline: self.release_anchor_lock(req.rid) self._free_staging_now(f.host_indices, f.aux_xfers) cc.prefetch_tokens_occupied -= f.occupied_tokens + # Nothing spliced: keep the surfaced host-hit fields truthful. + req.host_hit_length = 0 + req.swa_host_hit_length = 0 return unchanged # A hold staged under a different namespace than the consuming request @@ -809,20 +893,24 @@ class BufferModePipeline: ) return _drop() - # Splice-validity: the span only fits if the device prefix still - # ends exactly at the enqueue-time matched_len. - if len(req.prefix_indices) != f.matched_len: + splice_base = len(req.prefix_indices) + splice_tokens = staged_splice_tokens(f, splice_base) + if splice_tokens == 0: logger.warning( - "HiCache staged prefetch dropped req=%s reason=%s matched=%d " - "now=%d tokens_wasted=%d locked=%s", + "HiCache staged prefetch dropped req=%s matched=%d now=%d " + "tokens_wasted=%d locked=%s", req.rid, - "growth" if len(req.prefix_indices) > f.matched_len else "shrink", f.matched_len, - len(req.prefix_indices), + splice_base, f.num_tokens, req.rid in self.anchor_locks, ) return _drop() + trim_tokens = splice_base - f.matched_len + assert trim_tokens % cache.page_size == 0, ( + f"staged splice trim not page-aligned req={req.rid}: " + f"matched={f.matched_len} splice_base={splice_base}" + ) key = RadixKey( array("q", f.key_tokens), @@ -832,20 +920,21 @@ class BufferModePipeline: ).page_aligned(cache.page_size) span_end = f.matched_len + f.num_tokens - # Live ownership pre-check: the unified length detects anchor drift, + # Live ownership pre-check at the splice base: the unified length + # detects a stale request view (req matched before a later publish), # full_kv_hit_length detects FULL overlap the insert would dedup-free # (an SWA tombstone can mask live FULL from the unified match alone). live = cache.match_prefix(MatchPrefixParams(key=key)) if ( - len(live.device_indices) != f.matched_len - or live.full_kv_hit_length != f.matched_len + len(live.device_indices) != splice_base + or live.full_kv_hit_length != splice_base ): logger.warning( "HiCache staged prefetch dropped req=%s reason=overlap " - "matched=%d live_unified=%d live_full=%d tokens_wasted=%d " + "splice_base=%d live_unified=%d live_full=%d tokens_wasted=%d " "locked=%s", req.rid, - f.matched_len, + splice_base, len(live.device_indices), live.full_kv_hit_length, f.num_tokens, @@ -859,20 +948,20 @@ class BufferModePipeline: avail = cache.token_to_kv_pool_allocator.full_available_size() else: avail = cache.token_to_kv_pool_allocator.available_size() - if avail < f.num_tokens: - needed = f.num_tokens - avail + if avail < splice_tokens: + needed = splice_tokens - avail cache.evict_for_alloc(EvictParams(num_tokens=needed)) if cache.supports_swa(): avail = cache.token_to_kv_pool_allocator.full_available_size() else: avail = cache.token_to_kv_pool_allocator.available_size() - if avail < f.num_tokens: + if avail < splice_tokens: # Genuinely no room (locked pages): recompute. return _drop() load_back_id = -(f.operation_id) - 1 device_indices = cc.load( - host_indices=f.host_indices, + host_indices=f.host_indices[trim_tokens:], node_id=load_back_id, extra_pools=f.aux_xfers or None, ) @@ -905,7 +994,7 @@ class BufferModePipeline: InsertParams( key=key, value=torch.cat([req.prefix_indices, device_indices]), - prev_prefix_len=f.matched_len, + prev_prefix_len=splice_base, swa_evicted_seqlen=( max(0, span_end - len(swa_dev)) if swa_dev is not None else 0 ), @@ -913,15 +1002,17 @@ class BufferModePipeline: ) self.ongoing_buffer_load_back[load_back_id] = _OngoingBufferLoadBack( req_id=f.req_id, - num_tokens=f.num_tokens, + num_tokens=splice_tokens, occupied_tokens=f.occupied_tokens, aux_xfers=f.aux_xfers, + # The full staged bounce (not the trimmed H2D source): the ack + # frees it whole, trimmed head included. host_indices=f.host_indices, hash_values=f.hash_values, ) m = cache.match_prefix(MatchPrefixParams(key=key)) self.release_anchor_lock(req.rid) - canonical = m.device_indices[f.matched_len : span_end] + canonical = m.device_indices[splice_base:span_end] if len(m.device_indices) < span_end or not torch.equal( canonical, device_indices ): @@ -930,7 +1021,7 @@ class BufferModePipeline: raise RuntimeError( f"HiCache buffer load-back ownership violation req={f.req_id}: " f"insert prefix_len={insert_result.prefix_len} " - f"expected={f.matched_len}, adopted={len(m.device_indices)} " + f"expected={splice_base}, adopted={len(m.device_indices)} " f"span_end={span_end}, canonical_matches_incoming=" f"{len(m.device_indices) >= span_end and torch.equal(canonical, device_indices)}; " f"in-flight H2D targets freed slots" @@ -965,9 +1056,11 @@ class BufferModePipeline: cache.storage_metrics_collector.log_prefetched_tokens(f.num_tokens) return True - def release_aborted_staged(self, rid: str) -> bool: - """Free an aborted request's staged prefetch (nothing device-side - exists yet — only the bounce). Returns True when a hold existed.""" + def release_staged_hold(self, rid: str) -> bool: + """Free a staged hold outright — anchor pin, host bounce (KV + aux), + occupancy grant; nothing device-side exists yet. Called for aborts + and for holds that can no longer splice. Returns True when a hold + existed.""" self.release_anchor_lock(rid) staged = self.staged_prefetches.pop(rid, None) if staged is None: diff --git a/python/sglang/srt/mem_cache/unified_radix_cache.py b/python/sglang/srt/mem_cache/unified_radix_cache.py index 72d891957..8c26a8b24 100644 --- a/python/sglang/srt/mem_cache/unified_radix_cache.py +++ b/python/sglang/srt/mem_cache/unified_radix_cache.py @@ -254,6 +254,8 @@ class UnifiedRadixCache(BasePrefixCache): "issued": 0, "declined_too_short": 0, "declined_rate_limited": 0, + "declined_anchor_lost": 0, + "declined_device_covered": 0, "revoked_insufficient": 0, "revoked_full_miss": 0, "l3_demand_requests": 0, @@ -1788,11 +1790,16 @@ class UnifiedRadixCache(BasePrefixCache): comp_xfers, ) 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) + self.buffer_pipeline.set_prefix_ctx( + req_id, + matched_prefix_tokens, + extra_key=extra_key, + cache_salt=cache_salt, + ) + # Pin the just-matched anchor now: deferred to IO commit it is + # often already deleted under churn. The IO-commit call remains + # as the second chance that decides the fetch's fate. + self.buffer_pipeline.try_lock_anchor(req_id) else: # Cache mode reserves the requested span up front; buffer mode # grants occupancy later at hit-alloc time, sized to the hit. @@ -2054,12 +2061,14 @@ class UnifiedRadixCache(BasePrefixCache): 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.""" + def plan_staged_splice( + self, req_id: str, device_prefix_len: int + ) -> tuple[int, int]: + """(kv, swa) host-hit tokens a staged buffer-mode prefetch will splice + given the request's live device prefix; frees unusable holds.""" if self.buffer_pipeline is None: - return 0 - return self.buffer_pipeline.staged_prefetch_tokens(req_id) + return 0, 0 + return self.buffer_pipeline.plan_staged_splice(req_id, device_prefix_len) def staged_prefetch_swa_tokens(self, req_id: str) -> int: """SWA device tokens consuming a staged buffer-mode prefetch will @@ -2074,7 +2083,7 @@ class UnifiedRadixCache(BasePrefixCache): self._storage_prefetch_missed_rids.discard(rid) if ( self.buffer_pipeline is not None - and self.buffer_pipeline.release_aborted_staged(rid) + and self.buffer_pipeline.release_staged_hold(rid) ): return if rid not in self.ongoing_prefetch: @@ -2246,6 +2255,25 @@ class UnifiedRadixCache(BasePrefixCache): # prefetches ahead of us are consumed. The op stays in # ongoing_prefetch, so wait_complete keeps gating admission. return False + if buffer_mode: + # IO commit: pin before the bounce alloc so a cancel is a + # plain revoke and a parked op keeps its pin; a fetch whose + # splice base is gone is not worth its storage read. + if self.buffer_pipeline.try_lock_anchor(req_id) == "anchor_lost": + self._prefetch_outcome_stats["declined_anchor_lost"] += 1 + # Span still L3-resident: arm the paced retry to re-fetch + # from the shorter post-loss match. + self._storage_prefetch_missed_rids.add(req_id) + self.revoke_pending_prefetch(req_id) + return True + if self.buffer_pipeline.staged_span_covered( + req_id, operation.storage_hit_count + ): + # Live tree already covers the span: nothing left to + # splice, so skip the storage read. + self._prefetch_outcome_stats["declined_device_covered"] += 1 + self.revoke_pending_prefetch(req_id) + return True alloc_len = operation.storage_hit_count host_indices = cc.mem_pool_host.alloc(alloc_len) if host_indices is None: @@ -2273,10 +2301,6 @@ class UnifiedRadixCache(BasePrefixCache): self.ongoing_prefetch[req_id] = info._replace(host_indices=host_indices) if buffer_mode: cc.prefetch_tokens_occupied += alloc_len - # IO commit: pin the anchor until consumption. Do not read - # attributes off `operation` here — alternative cache - # controllers may expose a narrower surface. - self.buffer_pipeline.try_lock_anchor(req_id, info.anchor_node_id) cc.prefetch_buffer.put(operation) return True 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 9d3b4fd6c..12baceea5 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 @@ -3782,6 +3782,172 @@ class UnifiedRadixCacheSuite: ) ) + def test_buffer_only_load_back_trims_head_published_by_sibling(self): + """Growth-waste regression: a sibling publishing the span HEAD after + enqueue used to invalidate the whole staged fetch at consumption + (splice base moved past matched_len -> full drop, every fetched byte + wasted). Consumption must instead splice the tail beyond the live + prefix: sibling head slots stay untouched (add-only insert), the + tail carries the producer's bytes, and the ack frees the entire + bounce including the trimmed head.""" + self._skip_unsupported_hicache_test() + # Buffer-mode plan/commit logic is layout-independent, and each + # hicache fixture retains its pools 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() + _, (expected_k, expected_v) = self._produce_buffer_l3( + storage_dir, seq, marker=9 + ) + + cons, cons_alloc, cons_rtp = build_fixture(self.cfg) + self._init_buffer_hicache(cons, storage_dir) + avail0 = self._host_avail_sizes(cons) + + req_id = "growth-trim" + 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), + "prefetch did not stage", + ) + cons.pop_prefetch_loaded_tokens(req_id) + + # Sibling publishes the first page of the span while the hold parks. + head = seq[: self.cfg.page_size] + self._insert(cons, cons_alloc, cons_rtp, head) + sib = cons.match_prefix(MatchPrefixParams(key=RadixKey(array("q", head)))) + self.assertEqual(len(sib.device_indices), len(head)) + self._fill_full_kv(cons_alloc, sib.device_indices, marker=3) + head_k, head_v = self._snapshot_full_kv(cons_alloc, sib.device_indices) + + # The surfaced host hit is the splice-able tail, not the full span. + kv_tokens, swa_tokens = cons.plan_staged_splice(req_id, len(head)) + self.assertEqual(kv_tokens, len(seq) - len(head)) + self.assertEqual(swa_tokens, cons.staged_prefetch_swa_tokens(req_id)) + self.assertTrue(cons.buffer_pipeline.has_staged(req_id)) + + spliced = self._consume_staged_prefetch( + cons, req_id, prefix_len=len(head), prefix_indices=sib.device_indices + ) + self.assertEqual(int(spliced.numel()), len(seq) - len(head)) + + m = cons.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq)))) + self.assertEqual(len(m.device_indices), len(seq)) + self.assertTrue(torch.equal(m.device_indices[: len(head)], sib.device_indices)) + k, v = self._snapshot_full_kv(cons_alloc, m.device_indices[len(head) :]) + self.assertTrue(torch.equal(k, expected_k[len(head) :])) + self.assertTrue(torch.equal(v, expected_v[len(head) :])) + hk, hv = self._snapshot_full_kv(cons_alloc, m.device_indices[: len(head)]) + self.assertTrue(torch.equal(hk, head_k)) + self.assertTrue(torch.equal(hv, head_v)) + + # The whole bounce (trimmed head included) frees at the ack. + self.assertEqual(self._host_avail_sizes(cons), avail0) + self.assertEqual(cons.cache_controller.prefetch_tokens_occupied, 0) + cons.sanity_check() + + def test_buffer_only_plan_frees_covered_hold(self): + """A hold whose whole span became device-resident can never splice: + the surface-time plan must report (0, 0) and free it (bounce, anchor + pin, occupancy) — a kept hold would leak, since admission without a + host hit never calls init_load_back.""" + self._skip_unsupported_hicache_test() + # Buffer-mode plan/commit logic is layout-independent, and each + # hicache fixture retains its pools 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() + self._produce_buffer_l3(storage_dir, seq) + + cons, cons_alloc, cons_rtp = build_fixture(self.cfg) + self._init_buffer_hicache(cons, storage_dir) + avail0 = self._host_avail_sizes(cons) + + req_id = "covered-hold" + 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), + "prefetch did not stage", + ) + cons.pop_prefetch_loaded_tokens(req_id) + self._insert(cons, cons_alloc, cons_rtp, seq) + + self.assertEqual(cons.plan_staged_splice(req_id, len(seq)), (0, 0)) + self.assertFalse(cons.buffer_pipeline.has_staged(req_id)) + self.assertEqual(cons.buffer_pipeline.anchor_locks, {}) + self.assertEqual(cons.cache_controller.prefetch_tokens_occupied, 0) + self.assertEqual(self._host_avail_sizes(cons), avail0) + # Idempotent once freed. + self.assertEqual(cons.plan_staged_splice(req_id, len(seq)), (0, 0)) + cons.sanity_check() + + def test_buffer_only_hit_commit_cancels_device_covered_fetch(self): + """A sibling that publishes the span while the storage hit query is + in flight makes the fetch unconsumable; the IO-commit gate must + cancel it BEFORE the bounce alloc and the storage read (counted as + declined_device_covered), leaving no staging or occupancy behind.""" + self._skip_unsupported_hicache_test() + # Buffer-mode plan/commit logic is layout-independent, and each + # hicache fixture retains its pools 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() + self._produce_buffer_l3(storage_dir, seq) + + cons, cons_alloc, cons_rtp = build_fixture(self.cfg) + self._init_buffer_hicache(cons, storage_dir) + avail0 = self._host_avail_sizes(cons) + stats = cons._prefetch_outcome_stats + + req_id = "covered-at-commit" + cons.prefetch_from_storage( + req_id, cons.root_node.id, array("q", seq), None, None + ) + # Wait for the hit verdict WITHOUT draining it (the drain is the + # scheduler-thread IO commit under test). + deadline = time.time() + 10.0 + while ( + cons.cache_controller.prefetch_hit_queue.qsize() == 0 + and time.time() < deadline + ): + time.sleep(0.01) + self.assertGreater(cons.cache_controller.prefetch_hit_queue.qsize(), 0) + + self._insert(cons, cons_alloc, cons_rtp, seq) + cons.drain_storage_control_queues() + + self.assertEqual(stats["declined_device_covered"], 1) + self.assertNotIn(req_id, cons.ongoing_prefetch) + self.assertFalse(cons.buffer_pipeline.has_staged(req_id)) + self.assertTrue(cons.check_prefetch_progress(req_id)) + self.assertEqual(cons.cache_controller.prefetch_tokens_occupied, 0) + self.assertFalse(cons.pop_storage_prefetch_miss(req_id)) + # Aux staging released during the drain lands on queues sized before + # it; a second drain flushes them. + cons.drain_storage_control_queues() + self.assertEqual(self._host_avail_sizes(cons), avail0) + cons.sanity_check() + def test_buffer_only_swa_window_semantics(self): """SWA window handling across the three partial-window cases: root-anchored sub-window sequence (the sequence IS its window), @@ -8046,5 +8212,75 @@ class TestUnifiedRadixCacheStorageAttachBackfill(CustomTestCase): ) +class TestAnchorLockOutcomePolicy(CustomTestCase): + """try_lock_anchor finds the anchor by re-matching the live tree (no + carried node id to go stale): prefix intact -> lock the live node; + prefix shrunk -> anchor_lost so the caller cancels the storage IO + instead of gambling the read; cap_skip over budget (checked before the + match walk).""" + + _REQ = "req-1" + _PREFIX = list(range(100, 100 + 8)) + + def _make_pipeline(self, cache, cap_tokens=10_000): + from sglang.srt.mem_cache.buffer_mode.pipeline import BufferModePipeline + + pipeline = BufferModePipeline.__new__(BufferModePipeline) + pipeline.anchor_lock_enabled = True + pipeline.anchor_locks = {} + pipeline.anchor_locked_tokens_ = 0 + pipeline.anchor_lock_cap_tokens = cap_tokens + pipeline._anchor_lock_cap_skips = 0 + pipeline._prefetch_prefix_ctx = {self._REQ: (list(self._PREFIX), None, None)} + pipeline._cache = cache + return pipeline + + def _make_cache(self, live_match_len): + from types import SimpleNamespace + + cache = mock.MagicMock() + cache.tree_core.is_eagle = False + cache.match_prefix.return_value = SimpleNamespace( + device_indices=list(range(live_match_len)), last_device_node=99 + ) + return cache + + def test_intact_prefix_locks_live_node(self): + cache = self._make_cache(live_match_len=len(self._PREFIX)) + pipeline = self._make_pipeline(cache) + self.assertEqual(pipeline.try_lock_anchor(self._REQ), "locked") + self.assertEqual(pipeline.anchor_locks[self._REQ].node_id, 99) + self.assertEqual(pipeline.anchor_locked_tokens_, len(self._PREFIX)) + + def test_shrunk_prefix_reports_anchor_lost(self): + cache = self._make_cache(live_match_len=len(self._PREFIX) - 2) + pipeline = self._make_pipeline(cache) + self.assertEqual(pipeline.try_lock_anchor(self._REQ), "anchor_lost") + self.assertEqual(pipeline.anchor_locks, {}) + self.assertEqual(pipeline.anchor_locked_tokens_, 0) + + def test_over_cap_reports_cap_skip_before_matching(self): + cache = self._make_cache(live_match_len=len(self._PREFIX)) + pipeline = self._make_pipeline(cache, cap_tokens=len(self._PREFIX) - 1) + self.assertEqual(pipeline.try_lock_anchor(self._REQ), "cap_skip") + self.assertEqual(pipeline.anchor_locks, {}) + cache.match_prefix.assert_not_called() + + def test_root_anchor_reports_no_anchor(self): + cache = self._make_cache(live_match_len=0) + pipeline = self._make_pipeline(cache) + pipeline._prefetch_prefix_ctx[self._REQ] = ([], None, None) + self.assertEqual(pipeline.try_lock_anchor(self._REQ), "no_anchor") + cache.match_prefix.assert_not_called() + + def test_already_locked_is_idempotent(self): + cache = self._make_cache(live_match_len=len(self._PREFIX)) + pipeline = self._make_pipeline(cache) + self.assertEqual(pipeline.try_lock_anchor(self._REQ), "locked") + self.assertEqual(pipeline.try_lock_anchor(self._REQ), "locked") + self.assertEqual(pipeline.anchor_locked_tokens_, len(self._PREFIX)) + cache.match_prefix.assert_called_once() + + if __name__ == "__main__": unittest.main()