[HiCache] buffer mode: decide staged-fetch fate against the live tree (#36834)
This commit is contained in:
@@ -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()
|
||||
|
||||
Reference in New Issue
Block a user