[HiCache] buffer mode: decide staged-fetch fate against the live tree (#36834)

This commit is contained in:
Zhiqiang Xie
2026-08-29 01:51:47 -07:00
committed by GitHub
parent 0a585d5bb1
commit a328c19c81
4 changed files with 429 additions and 74 deletions
+7 -5
View File
@@ -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),
@@ -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:
@@ -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
@@ -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()