L3 3.2: async disk->L2 reload + request hold (Model B), zero new collectives, EAGLE-correct
On a radix miss, reload the evicted-but-L3-durable suffix from disk into L2 + insert a fresh radix node, holding the triggering request until then (mirrors the storage-prefetch hold, which is disabled under CP). The reloaded prefix re-enters the radix as a normal L2 node, so the existing load_back serves L2->L1 -- no new device path. Design (docs_internal/cp_hicache_l3_phase3_impl_design.md §9, refined for minimal sync): - ENTRY (scheduler _prefetch_kvcache, enable_cp_l3 branch): mark a miss-with-suffix as a reload candidate (rank-uniform: requests are broadcast). cp_l3_reload_lookup computes the suffix's full-page content hashes -- the SAME SHA chain spill wrote (compute_node_hash_values). - RESERVE + ADMIT fold into the existing per-tick _drain_l3_control_queues MINs -- ZERO new collectives: the candidates' per-rank exists_prefix counts ride the qsize MIN vector -> a rank-uniform agreed count C (this reconciles the async-LMDB read-skew); reload-ack oks ride the durable ok-AND. Reserve is then collective-free (deterministic evict-to-fit mirroring _reserve_write_cp_shared_l2_evict_to_fit); the request waits non-blocking in waiting_queue (check_cp_l3_reload_progress skip) and is admitted + re-matched the SAME tick the reload lands (check_hicache_events runs before batch formation). CP-aware insert only on a still-clean attach (else abort+recompute). Capped (cp_l3_reload_max_inflight) + content-key deduped (piggyback). EAGLE/bigram (opus-studied, PROVABLY correct -- not deferred): the radix key is bigram-converted over the whole request, and convert_to_bigram_key(fill_ids[M:]) == bigrams(fill_ids)[M:] exactly (the boundary bigram belongs to the matched prefix), so the bigram-converted suffix hashes reproduce node.hash_value. Verified vs source + a new slice-identity regression test. Opus review (FIX-THEN-SHIP, no CRITICAL; the 4 crash surfaces -- MIN-vector shape, read-skew, attach-anchor, placement-digest -- traced clean) + independently re-verified; fixes folded: - M1: cp_l3_release_request clears reload state on request abort/timeout/preempt (no leak). - H1: rank-uniform per-op TTL bound releases a held request if its reload op never acks (the default SGLANG_REQ_WAITING_TIMEOUT is off); the reservation frees safely at the (late) ack. - L1: fail-soft anchor guard at insert (re-derive the first suffix hash from lhn's live hash). - empty-waiters reload aborts (frees L2) instead of inserting an unrequested node. Imports hoisted to top-level (get_hash_str -- hicache_storage is a leaf module, no cycle; convert_to_bigram_key already top-level), no inline imports. py_compile clean; 54 L3 unit tests green. Reload triggering (submit_reload/exists_prefix/free_object) now wired (was write-only). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
@@ -2053,6 +2053,18 @@ class Scheduler(
|
||||
last_hash,
|
||||
prefix_keys,
|
||||
)
|
||||
elif getattr(self.tree_cache, "enable_cp_l3", False):
|
||||
# CP L3 reload (3.2): on a radix miss, mark the unmatched suffix as an L3 reload candidate. The
|
||||
# rank-uniform reserve + hold decision happens in check_hicache_events (runs before the waiting
|
||||
# queue forms a batch); a non-L3 miss falls through to normal prefill. Mirrors the storage prefetch
|
||||
# hold, which is disabled under CP.
|
||||
req.init_next_round_input(self.tree_cache, cow_mamba=False)
|
||||
matched_len = len(req.prefix_indices) + req.host_hit_length
|
||||
new_input_tokens = req.fill_ids[matched_len:]
|
||||
if new_input_tokens:
|
||||
self.tree_cache.cp_l3_reload_lookup(
|
||||
req, req.last_host_node, new_input_tokens
|
||||
)
|
||||
|
||||
def _add_request_to_queue(self, req: Req, is_retracted: bool = False):
|
||||
if self.disaggregation_mode == DisaggregationMode.NULL:
|
||||
@@ -2134,6 +2146,8 @@ class Scheduler(
|
||||
self.tree_cache.release_aborted_request(candidate_req.rid)
|
||||
elif self.enable_hierarchical_cache:
|
||||
self.tree_cache.terminate_prefetch(candidate_req.rid)
|
||||
if getattr(self.tree_cache, "enable_cp_l3", False):
|
||||
self.tree_cache.cp_l3_release_request(candidate_req.rid)
|
||||
self.waiting_queue.pop(idx)
|
||||
req_to_abort = candidate_req
|
||||
message = "The request is aborted by a higher priority request."
|
||||
@@ -2164,6 +2178,8 @@ class Scheduler(
|
||||
if self.enable_hicache_storage:
|
||||
# Release prefetch events associated with the request
|
||||
self.tree_cache.release_aborted_request(req.rid)
|
||||
if getattr(self.tree_cache, "enable_cp_l3", False):
|
||||
self.tree_cache.cp_l3_release_request(req.rid)
|
||||
self.send_to_tokenizer.send_output(
|
||||
AbortReq(
|
||||
finished_reason={
|
||||
@@ -2639,6 +2655,13 @@ class Scheduler(
|
||||
req.storage_hit_length = self.tree_cache.pop_prefetch_loaded_tokens(
|
||||
req.rid
|
||||
)
|
||||
elif getattr(self.tree_cache, "enable_cp_l3", False):
|
||||
# Hold a request whose unmatched suffix is being reloaded from the CP L3 disk tier; on
|
||||
# admission the reloaded prefix is an L2-resident radix node, so the re-match below picks it
|
||||
# up as a host hit (then the normal load_back serves L2->L1). No collective here (the done
|
||||
# flag was set rank-uniformly in check_hicache_events).
|
||||
if not self.tree_cache.check_cp_l3_reload_progress(req.rid):
|
||||
continue
|
||||
|
||||
req.init_next_round_input(self.tree_cache)
|
||||
|
||||
@@ -3618,6 +3641,8 @@ class Scheduler(
|
||||
if self.enable_hicache_storage:
|
||||
# to release prefetch events associated with the request
|
||||
self.tree_cache.release_aborted_request(req.rid)
|
||||
if getattr(self.tree_cache, "enable_cp_l3", False):
|
||||
self.tree_cache.cp_l3_release_request(req.rid)
|
||||
self.send_to_tokenizer.send_output(AbortReq(rid=req.rid), req)
|
||||
# For disaggregation decode mode, the request in the waiting queue has KV cache allocated.
|
||||
if self.disaggregation_mode == DisaggregationMode.DECODE:
|
||||
|
||||
@@ -42,6 +42,7 @@ from sglang.srt.mem_cache.cp_shared_kv_layout import (
|
||||
CpSharedKVLayout,
|
||||
pad_token_locs_to_page_boundary,
|
||||
)
|
||||
from sglang.srt.mem_cache.hicache_storage import get_hash_str
|
||||
from sglang.srt.mem_cache.memory_pool import (
|
||||
MHATokenToKVPool,
|
||||
MLATokenToKVPool,
|
||||
@@ -1059,6 +1060,17 @@ class HiRadixCache(RadixCache):
|
||||
self._cp_l3_spill_dropped = 0
|
||||
# backpressure: cap concurrent background spills (bounds the bg queue + staging RAM, rank-uniform)
|
||||
self.cp_l3_spill_max_inflight = 32
|
||||
# L3 RELOAD (3.2): disk->L2 + radix re-insert, holding the triggering request (mirrors the storage
|
||||
# prefetch hold, which is dead under CP). All cross-rank agreement folds into the existing per-tick
|
||||
# _drain_l3_control_queues MIN -- no new collective. ongoing keyed by reload op_id (like spill).
|
||||
self.cp_l3_reload_candidates = [] # [(rid, last_host_node, suffix_key, suffix_hashes)] marked at entry
|
||||
self.ongoing_l3_reload = {} # op_id -> dict(reload_key,last_host_node,suffix_key,suffix_hashes,n_pages,waiters)
|
||||
self.cp_l3_reload_inflight_keys = {} # reload_key -> op_id (same-suffix dedup: a 2nd req piggybacks)
|
||||
self.cp_l3_reload_done = {} # rid -> bool; read (no collective) by the scheduler waiting-queue skip
|
||||
self.cp_l3_reload_max_inflight = 32 # cap concurrent reloads (bounds reserved-but-pending L2 pages)
|
||||
# Bound the request hold (rank-uniform tick countdown): the default SGLANG_REQ_WAITING_TIMEOUT is -1
|
||||
# (off), so a lost/hung reload op (e.g. a dead disk) could otherwise wedge a held request forever.
|
||||
self.cp_l3_reload_ttl_ticks = 4096
|
||||
|
||||
(
|
||||
extra_config,
|
||||
@@ -1281,45 +1293,351 @@ class HiRadixCache(RadixCache):
|
||||
ok-AND (phase 2), reload-ack -> 3.2. Idle early-return: has_inflight() is a replicated quantity (submits
|
||||
are lockstep on the replicated plan), so an idle tick issues no collective."""
|
||||
store = self.cp_l3_store
|
||||
if store is None or not store.has_inflight():
|
||||
if store is None:
|
||||
return
|
||||
qsizes = torch.tensor(
|
||||
[store.ack_gather_qsize(), store.ack_durable_qsize(), store.ack_reload_qsize()], dtype=torch.int
|
||||
marked = self.cp_l3_reload_candidates
|
||||
if not store.has_inflight() and not marked:
|
||||
return # idle (replicated quantity) -> no collective
|
||||
# Re-validate marked reload candidates against the REPLICATED radix -> rank-uniform valid set + order;
|
||||
# each contributes its LOCAL exists_prefix count (skew-prone value, reconciled by the MIN below).
|
||||
self.cp_l3_reload_candidates = []
|
||||
payloads = store.payloads
|
||||
valid_cands = [] # (rid, last_host_node, suffix_key, suffix_hashes)
|
||||
cand_counts = []
|
||||
for (rid, lhn, skey, shashes) in marked:
|
||||
if rid in self.cp_l3_reload_done or not self._cp_l3_reload_attach_ok(lhn, skey):
|
||||
continue # already in flight, or lhn evicted / suffix no longer a clean miss -> recompute
|
||||
cand_counts.append(int(store.exists_prefix(shashes, payloads)))
|
||||
valid_cands.append((rid, lhn, skey, shashes))
|
||||
# (1) ONE MIN reduce: the three ack qsizes + the per-candidate counts (rank-uniform vector shape).
|
||||
vec = torch.tensor(
|
||||
[store.ack_gather_qsize(), store.ack_durable_qsize(), store.ack_reload_qsize()] + cand_counts,
|
||||
dtype=torch.int,
|
||||
)
|
||||
if self.tp_world_size > 1:
|
||||
self._cp_hicache_all_reduce(
|
||||
qsizes, op=torch.distributed.ReduceOp.MIN, group=self.tp_group, tag="l3_queue_min"
|
||||
vec, op=torch.distributed.ReduceOp.MIN, group=self.tp_group, tag="l3_queue_min"
|
||||
)
|
||||
n_gather, n_durable, n_reload = map(int, qsizes.tolist())
|
||||
# Phase 1 -- GATHER done: release the eviction pin. The staged copy is slab-independent, so the object
|
||||
# may be evicted from L2 now; the disk write still completes off the staged copy. The node stays in
|
||||
# ongoing_l3_spill (inflight, counts toward budget) until the durable ack. Pin release is MIN-gated, so
|
||||
# it is rank-uniform -- a divergent unpin would drift the shared-slab eviction order (placement_digest).
|
||||
vals = [int(v) for v in vec.tolist()]
|
||||
n_gather, n_durable, n_reload = vals[0], vals[1], vals[2]
|
||||
agreed_counts = vals[3:]
|
||||
# Phase 1 -- GATHER done: release the eviction pin (MIN-gated -> rank-uniform). The staged copy is
|
||||
# slab-independent, so the object may be evicted now; the disk write still completes off it.
|
||||
for op_id in store.drain_gather_acks(n_gather):
|
||||
node = self.ongoing_l3_spill.get(op_id)
|
||||
if node is not None and int(getattr(node, "host_ref_counter", 0) or 0) > 0:
|
||||
node.release_host()
|
||||
# Phase 2 -- DURABLE: mark L3-durable, but ONLY if EVERY rank's write succeeded (rank-uniform durability
|
||||
# via an ok-AND, i.e. ReduceOp.MIN over the per-op ok bits). A write failure on any rank (e.g. slab full
|
||||
# pre-3.3) => no rank marks it durable => fail-soft recompute, never a divergent l3_durable that would
|
||||
# split eviction eligibility. The acks are FIFO in replicated submit order, so the ok tensor aligns
|
||||
# per-op across ranks. No allocator release here (approach D: it was a copy, not a move).
|
||||
# A failed object is deliberately NOT re-enqueued: ok=False means slab-full (until 3.3 disk-eviction)
|
||||
# or a disk I/O error -- an immediate retry would just re-gather + re-fail every tick (churn strictly
|
||||
# worse than the benign abandon). It is retried on its next natural re-commit; 3.3 removes the slab-full
|
||||
# cause. Correctness is unaffected (a non-durable object simply recomputes on a future miss).
|
||||
# (2) ONE ok-AND (MIN over ok bits) covering BOTH spill-durable and reload acks: an op is marked
|
||||
# durable / admitted only if EVERY rank succeeded -> rank-uniform (a per-rank disk failure -> uniform
|
||||
# fail-soft, never a divergent l3_durable or divergent insert that would crash placement_digest). Acks
|
||||
# are FIFO in replicated submit order, so the ok vector aligns per-op. (Spill ok=False is NOT re-enqueued
|
||||
# -- it retries on natural re-commit; 3.3 removes the slab-full cause; recompute is correctness-safe.)
|
||||
durable = store.drain_durable_acks(n_durable)
|
||||
if durable:
|
||||
oks = torch.tensor([1 if ok else 0 for (_id, ok) in durable], dtype=torch.int)
|
||||
reload_acks = store.drain_reload_acks(n_reload)
|
||||
ok_bits = [1 if ok else 0 for (_i, ok) in durable] + [1 if ok else 0 for (_i, ok) in reload_acks]
|
||||
if ok_bits:
|
||||
ok_t = torch.tensor(ok_bits, dtype=torch.int)
|
||||
if self.tp_world_size > 1:
|
||||
self._cp_hicache_all_reduce(
|
||||
oks, op=torch.distributed.ReduceOp.MIN, group=self.tp_group, tag="l3_durable_ok"
|
||||
ok_t, op=torch.distributed.ReduceOp.MIN, group=self.tp_group, tag="l3_ok_and"
|
||||
)
|
||||
for (op_id, _ok), agreed in zip(durable, oks.tolist()):
|
||||
node = self.ongoing_l3_spill.pop(op_id, None)
|
||||
if node is not None and int(agreed) == 1:
|
||||
node.l3_durable = True # now reloadable from L3; maintainer skips it (O(1))
|
||||
store.drain_reload_acks(n_reload) # 3.2 wires reload-ack admission + pin release
|
||||
agreed_ok = [int(a) for a in ok_t.tolist()]
|
||||
else:
|
||||
agreed_ok = []
|
||||
nd = len(durable)
|
||||
for (op_id, _ok), a in zip(durable, agreed_ok[:nd]): # spill durable
|
||||
node = self.ongoing_l3_spill.pop(op_id, None)
|
||||
if node is not None and a == 1:
|
||||
node.l3_durable = True
|
||||
for (op_id, _ok), a in zip(reload_acks, agreed_ok[nd:]): # reload admission
|
||||
self._cp_l3_admit_reload(op_id, a == 1)
|
||||
# Reserve NEW reloads from the rank-uniform agreed counts (collective-free deterministic reserve).
|
||||
self._cp_l3_reserve_reloads(valid_cands, agreed_counts)
|
||||
# Bound the hold: release waiters of any reload op that hasn't acked within ttl ticks (rank-uniform).
|
||||
self._cp_l3_expire_stale_reloads()
|
||||
|
||||
def cp_l3_reload_lookup(self, req, last_host_node: TreeNode, suffix_tokens) -> None:
|
||||
"""Entry hook (scheduler _prefetch_kvcache, rank-uniform: requests are broadcast). Mark a radix-miss
|
||||
request as an L3 reload candidate. Cheap + LOCAL: compute the suffix's FULL-page content hashes (the
|
||||
same SHA chain spill wrote, radix_cache.compute_node_hash_values) and queue the candidate. NO
|
||||
exists_prefix / collective / reserve here -- the rank-uniform reserve decision (a MIN over the per-rank
|
||||
exists_prefix counts) happens in _drain_l3_control_queues, which runs before the waiting-queue batch
|
||||
forms. A non-L3 miss simply falls out there (agreed count 0) and the request proceeds to recompute.
|
||||
|
||||
EAGLE/bigram (verified correct): the radix key is bigram-converted over the WHOLE request
|
||||
(maybe_bigram_convert), so slicing raw fill_ids at the bigram-unit matched_len then re-converting
|
||||
reproduces EXACTLY the suffix node's bigrams -- the identity convert_to_bigram_key(f[M:]) == bigrams(f)[M:]
|
||||
holds because the boundary bigram (f[M-1],f[M]) belongs to the matched prefix (index M-1), not the suffix.
|
||||
So get_hash_str over the bigram pages, chained from the prefix's last hash, reproduces the spilled
|
||||
node.hash_value (this is the proven storage-prefetch recipe). The content hash is over token_ids only
|
||||
(compute_node_hash_values ignores extra_key); the RadixKey still carries extra_key + is_bigram so the
|
||||
attach/insert child-key routing matches the radix. page_size is in bigram units under eagle, so the page
|
||||
chunking is unchanged."""
|
||||
if not self.enable_cp_l3 or self.cp_l3_store is None or last_host_node is None:
|
||||
return
|
||||
if self.is_eagle:
|
||||
suffix_tokens = convert_to_bigram_key(list(suffix_tokens))
|
||||
page = self.page_size
|
||||
n_full = len(suffix_tokens) // page # only full pages are content pages (spill wrote full pages)
|
||||
if n_full <= 0:
|
||||
return
|
||||
parent_hash = (
|
||||
last_host_node.get_last_hash_value()
|
||||
if last_host_node is not self.root_node
|
||||
else None
|
||||
)
|
||||
suffix_hashes: List[str] = []
|
||||
suffix_key_tokens = [] # bigram tuples under eagle, raw ints otherwise
|
||||
for i in range(n_full):
|
||||
pg = list(suffix_tokens[i * page : (i + 1) * page])
|
||||
h = get_hash_str(pg, prior_hash=parent_hash) # get_hash_str hashes bigram tuples too (compute_node_hash_values)
|
||||
suffix_hashes.append(h)
|
||||
parent_hash = h
|
||||
suffix_key_tokens.extend(pg)
|
||||
suffix_key = RadixKey(
|
||||
token_ids=suffix_key_tokens,
|
||||
extra_key=getattr(req, "extra_key", None),
|
||||
is_bigram=self.is_eagle,
|
||||
)
|
||||
self.cp_l3_reload_candidates.append(
|
||||
(req.rid, last_host_node, suffix_key, suffix_hashes)
|
||||
)
|
||||
|
||||
def check_cp_l3_reload_progress(self, rid: str) -> bool:
|
||||
"""Scheduler waiting-queue skip (mirrors check_prefetch_progress). True = the request may proceed
|
||||
(not an L3 reload candidate, or its reload is admitted); False = an L3 reload is in flight -> hold.
|
||||
Reads the rank-uniformly-set done flag -- no collective. On a True (admitted) it consumes the flag;
|
||||
the request then re-matches (init_next_round_input) and hits the inserted L2 node."""
|
||||
done = self.cp_l3_reload_done.get(rid)
|
||||
if done is None:
|
||||
return True # not held for L3 reload
|
||||
if done:
|
||||
self.cp_l3_reload_done.pop(rid, None)
|
||||
return True
|
||||
return False
|
||||
|
||||
def cp_l3_release_request(self, rid: str) -> None:
|
||||
"""Abort/cleanup hook (scheduler): drop a request from all L3 reload state so an aborted held request
|
||||
leaves no dangling done-flag / waiter. The reservation is freed when its op acks (a then-empty-waiters
|
||||
op aborts in _cp_l3_admit_reload). Rank-uniform: aborts are broadcast."""
|
||||
self.cp_l3_reload_done.pop(rid, None)
|
||||
for info in self.ongoing_l3_reload.values():
|
||||
if rid in info["waiters"]:
|
||||
info["waiters"].remove(rid)
|
||||
|
||||
def _cp_l3_expire_stale_reloads(self) -> None:
|
||||
"""Bound the request hold (rank-uniform tick countdown, called per tick): if a reload op hasn't acked
|
||||
within ttl ticks (a lost/hung op -- e.g. a dead disk; the default waiting-timeout is off), RELEASE its
|
||||
waiters to recompute. The op + its reservation stay until the (eventual) ack frees them -- we must NOT
|
||||
free the reserved pages while the bg thread might still scatter into them. A real disk hang wedges all
|
||||
of L3 anyway; this just keeps the held request from waiting forever."""
|
||||
for info in self.ongoing_l3_reload.values():
|
||||
info["ttl"] -= 1
|
||||
if info["ttl"] <= 0 and info["waiters"]:
|
||||
for rid in info["waiters"]:
|
||||
self.cp_l3_reload_done.pop(rid, None) # release -> recompute
|
||||
info["waiters"] = []
|
||||
|
||||
def _cp_l3_reload_attach_ok(self, lhn: TreeNode, suffix_key) -> bool:
|
||||
"""Rank-uniform (replicated radix): lhn is still attached to the tree AND the suffix's first page is a
|
||||
clean miss under it (no existing child) -> a single new child can attach. Else the reload is dropped
|
||||
(recompute). Checked at reserve AND at admission (the tree may change in between)."""
|
||||
if lhn is None or len(suffix_key) == 0:
|
||||
return False
|
||||
if lhn is not self.root_node:
|
||||
parent = getattr(lhn, "parent", None)
|
||||
if parent is None:
|
||||
return False
|
||||
try:
|
||||
pkey = self.get_child_key_fn(lhn.key)
|
||||
except Exception:
|
||||
return False
|
||||
if getattr(parent, "children", {}).get(pkey) is not lhn:
|
||||
return False # lhn detached from the tree since the candidate was marked
|
||||
try:
|
||||
child_key = self.get_child_key_fn(suffix_key)
|
||||
except Exception:
|
||||
return False
|
||||
return child_key not in getattr(lhn, "children", {})
|
||||
|
||||
def _cp_l3_reserve_reloads(self, valid_cands, agreed_counts) -> None:
|
||||
"""Collective-free reserve (rank-uniform inputs => identical mutation on every rank): for each
|
||||
candidate whose AGREED reloadable page count C >= the load-back floor, reserve a NEW (uncommitted) L2
|
||||
object (owner-aware, evict-to-fit) + submit the disk->slab reload + hold the request. Same-suffix dedup
|
||||
(a 2nd request piggybacks on the in-flight reload). Bounded by cp_l3_reload_max_inflight."""
|
||||
store = self.cp_l3_store
|
||||
allocator = getattr(self.cache_controller, "cp_shared_l2_page_allocator", None)
|
||||
if allocator is None:
|
||||
return
|
||||
min_pages = max(1, self.load_back_threshold // self.page_size)
|
||||
reserved_this_tick = {}
|
||||
for (rid, lhn, skey, shashes), c in zip(valid_cands, agreed_counts):
|
||||
c = int(c)
|
||||
if c < min_pages:
|
||||
continue # below the load-back floor (or 0 = not in L3) -> recompute
|
||||
reload_key = f"l3rl:{shashes[c - 1]}" # rank-uniform + dedups identical (content,length) reloads
|
||||
existing = self.cp_l3_reload_inflight_keys.get(reload_key)
|
||||
if existing is None:
|
||||
existing = reserved_this_tick.get(reload_key)
|
||||
if existing is not None:
|
||||
self.ongoing_l3_reload[existing]["waiters"].append(rid) # piggyback on the in-flight reload
|
||||
self.cp_l3_reload_done[rid] = False
|
||||
continue
|
||||
if len(self.ongoing_l3_reload) >= self.cp_l3_reload_max_inflight:
|
||||
continue # cap reached (len is replicated) -> recompute
|
||||
ranges = self._cp_l3_reserve_reload_object(reload_key, store.payloads, c, allocator)
|
||||
if ranges is None:
|
||||
continue # could not fit even after evict-to-fit -> recompute
|
||||
owned_pages = self._cp_l3_build_reload_owned_pages(ranges, shashes[:c], store)
|
||||
op_id = store.submit_reload(reload_key, owned_pages)
|
||||
self.ongoing_l3_reload[op_id] = {
|
||||
"reload_key": reload_key, "last_host_node": lhn, "suffix_key": skey,
|
||||
"suffix_hashes": list(shashes[:c]), "n_pages": c, "waiters": [rid],
|
||||
"ttl": self.cp_l3_reload_ttl_ticks,
|
||||
}
|
||||
self.cp_l3_reload_inflight_keys[reload_key] = op_id
|
||||
reserved_this_tick[reload_key] = op_id
|
||||
self.cp_l3_reload_done[rid] = False # hold the request until admission
|
||||
|
||||
def _cp_l3_reserve_reload_object(self, reload_key, payloads, n_pages, allocator):
|
||||
"""Reserve n_pages per payload for a reload as a NEW uncommitted object, evicting cold committed
|
||||
shared-L2 leaves (replicated SLRU order -> collective-free, mirrors _reserve_write_cp_shared_l2_evict_to_fit)
|
||||
to open contiguous room. Returns the ranges dict, or None if it cannot fit (abort + recompute)."""
|
||||
victims = None
|
||||
vi = 0
|
||||
try:
|
||||
for pk in payloads:
|
||||
while True:
|
||||
try:
|
||||
allocator.reserve(reload_key, pk, n_pages)
|
||||
break
|
||||
except ValueError as e:
|
||||
if "insufficient contiguous" not in str(e):
|
||||
raise # a real programming/validation error -> fail loud
|
||||
if victims is None:
|
||||
victims = sorted(
|
||||
(
|
||||
nd
|
||||
for nd in getattr(self, "evictable_host_leaves", set())
|
||||
if self._cp_host_leaf_is_plannable_victim(nd)
|
||||
and self._is_cp_shared_l2_metadata(getattr(nd, "cp_hicache", None))
|
||||
),
|
||||
key=self._cp_host_evict_key,
|
||||
)
|
||||
if vi >= len(victims):
|
||||
allocator.abort(reload_key) # free any payload already reserved under this key
|
||||
return None
|
||||
victim = victims[vi]
|
||||
vi += 1
|
||||
if not self._cp_host_leaf_is_plannable_victim(victim):
|
||||
raise RuntimeError(
|
||||
"CP shared-L2 reload evict-to-fit plan diverged before reservation: "
|
||||
f"reload_key={reload_key} victim_id={getattr(victim, 'id', None)}"
|
||||
)
|
||||
self._record_remove_event(victim)
|
||||
self._cp_account_leave_node_evict(victim, victim.cp_hicache)
|
||||
self.cache_controller.evict_cp_host(victim.cp_hicache)
|
||||
victim.host_len = 0
|
||||
victim.cp_hicache = None
|
||||
victim.host_value = None
|
||||
self._remove_host_leaf(victim)
|
||||
except Exception:
|
||||
allocator.abort(reload_key)
|
||||
raise
|
||||
return allocator.object_ranges(reload_key)
|
||||
|
||||
def _cp_l3_build_reload_owned_pages(self, ranges, page_hashes, store):
|
||||
"""This rank's owned (dest_slab_page, content_hash) per payload (owner = object-local i % cp_size,
|
||||
identical to spill's _cp_l3_build_owned_pages -> each rank reloads exactly the pages it spilled)."""
|
||||
cp_size, cp_rank = store.cp_size, store.cp_rank
|
||||
n = len(page_hashes)
|
||||
owned = {}
|
||||
for pk, rng in ranges.items():
|
||||
base = rng.base_page
|
||||
owned[pk] = [(base + i, page_hashes[i]) for i in range(n) if i % cp_size == cp_rank]
|
||||
return owned
|
||||
|
||||
def _cp_l3_admit_reload(self, op_id: int, ok: bool) -> None:
|
||||
"""Reload op done (rank-uniform ok via the ok-AND). ok + still-clean-attach -> commit the reserved L2
|
||||
object + insert ONE CP-aware radix node + mark waiters done (they re-match -> L2 hit -> load_back). Else
|
||||
-> abort the reservation + release waiters to recompute. Rank-uniform: same op, same ok, replicated radix."""
|
||||
info = self.ongoing_l3_reload.pop(op_id, None)
|
||||
if info is None:
|
||||
return
|
||||
self.cp_l3_reload_inflight_keys.pop(info["reload_key"], None)
|
||||
allocator = getattr(self.cache_controller, "cp_shared_l2_page_allocator", None)
|
||||
inserted = False
|
||||
if not info["waiters"]:
|
||||
# every waiter was released (aborted, or the ttl bound fired) -> nobody needs this reload now;
|
||||
# free the reservation rather than insert an unrequested node (frees L2 for active work).
|
||||
if allocator is not None:
|
||||
allocator.abort(info["reload_key"])
|
||||
return
|
||||
if (
|
||||
ok
|
||||
and allocator is not None
|
||||
and self._cp_l3_reload_attach_ok(info["last_host_node"], info["suffix_key"])
|
||||
):
|
||||
allocator.mark_object_committed(info["reload_key"])
|
||||
inserted = self._cp_l3_insert_reloaded_node(info)
|
||||
if not inserted:
|
||||
allocator.abort(info["reload_key"])
|
||||
elif allocator is not None:
|
||||
allocator.abort(info["reload_key"]) # failed reload / stale attach -> free the reservation
|
||||
for rid in info["waiters"]:
|
||||
if inserted:
|
||||
self.cp_l3_reload_done[rid] = True # admitted -> re-match hits the inserted node
|
||||
else:
|
||||
self.cp_l3_reload_done.pop(rid, None) # released -> proceed (recompute)
|
||||
|
||||
def _cp_l3_insert_reloaded_node(self, info) -> bool:
|
||||
"""Insert ONE CP-aware radix node for the reloaded prefix under last_host_node (clean attach
|
||||
re-validated here against the live tree). cp_hicache = a committed CpSharedL2NodeMetadata over the
|
||||
reserved ranges; value=None (device-evicted, L2-resident) so the normal load_back serves L2->L1; the
|
||||
node is L3-durable (just read from L3). Rank-uniform (replicated radix + rank-uniform inputs)."""
|
||||
from sglang.srt.mem_cache.cp_shared_l2_pool import CpSharedL2NodeMetadata
|
||||
|
||||
lhn = info["last_host_node"]
|
||||
skey = info["suffix_key"]
|
||||
allocator = getattr(self.cache_controller, "cp_shared_l2_page_allocator", None)
|
||||
ranges = allocator.object_ranges(info["reload_key"]) if allocator is not None else {}
|
||||
if not ranges:
|
||||
return False
|
||||
child_key = self.get_child_key_fn(skey)
|
||||
if child_key in getattr(lhn, "children", {}):
|
||||
return False # raced to non-clean since the attach_ok check
|
||||
# Fail-soft anchor guard: re-derive the first suffix hash from lhn's LIVE last hash; if it no longer
|
||||
# matches the candidate's recorded hash the anchor moved (lhn's chain changed since the mark) -> abort
|
||||
# (recompute). Traced unreachable today (a retained device-evicted node is provably backuped + its hash
|
||||
# stable), but L3 reload is the first CP consumer of this match boundary, so guard explicitly.
|
||||
parent_basis = lhn.get_last_hash_value() if lhn is not self.root_node else None
|
||||
first_page = list(skey.token_ids[: self.page_size])
|
||||
if not first_page or get_hash_str(first_page, prior_hash=parent_basis) != info["suffix_hashes"][0]:
|
||||
return False
|
||||
n_pages = info["n_pages"]
|
||||
meta = CpSharedL2NodeMetadata(
|
||||
logical_len=n_pages * self.page_size,
|
||||
padded_len=n_pages * self.page_size,
|
||||
page_size=self.page_size,
|
||||
object_ranges=ranges,
|
||||
required_payloads=tuple(ranges.keys()),
|
||||
object_key=info["reload_key"],
|
||||
)
|
||||
new_node = TreeNode(priority=lhn.priority)
|
||||
new_node.parent = lhn
|
||||
new_node.key = skey.compacted()
|
||||
new_node.value = None # device-evicted; load_back fills L1 from L2
|
||||
new_node.host_value = None
|
||||
new_node.host_len = n_pages * self.page_size
|
||||
new_node.cp_hicache = meta
|
||||
new_node.hash_value = list(info["suffix_hashes"])
|
||||
new_node.l3_durable = True # just reloaded from L3 -> the spill maintainer skips it (O(1))
|
||||
new_node.last_access_time = TreeNode.next_access_time()
|
||||
lhn.children[child_key] = new_node
|
||||
self._cp_account_enter_committed(new_node.cp_hicache)
|
||||
self._update_host_leaf_status(new_node)
|
||||
self._update_host_leaf_status(lhn)
|
||||
return True
|
||||
|
||||
def _cp_l3_enqueue_spill(self, object_key: str, node: TreeNode) -> None:
|
||||
"""Enqueue a committed object for proactive spill (rank-uniform: called only from the replicated
|
||||
@@ -2935,6 +3253,10 @@ class HiRadixCache(RadixCache):
|
||||
self.cp_l3_store.clear()
|
||||
self.cp_l3_spill_pending.clear() # drop queued (now-flushed) spill candidates
|
||||
self.ongoing_l3_spill.clear()
|
||||
self.cp_l3_reload_candidates.clear()
|
||||
self.ongoing_l3_reload.clear()
|
||||
self.cp_l3_reload_inflight_keys.clear()
|
||||
self.cp_l3_reload_done.clear()
|
||||
self.token_to_kv_pool_host.clear()
|
||||
self.cache_controller.clear_draft_host_pool()
|
||||
# Clear per-request tracking dicts
|
||||
|
||||
@@ -82,6 +82,17 @@ class TestConvertToBigramKey(unittest.TestCase):
|
||||
self.assertEqual(got, [(7, 7), (7, 7)])
|
||||
self.assertIsNot(got[0], got[1])
|
||||
|
||||
def test_suffix_slice_identity_for_l3_reload(self):
|
||||
# The CP L3 reload (cp_l3_reload_lookup) computes a radix-miss suffix's content hashes by slicing the
|
||||
# RAW fill_ids at the bigram-unit matched_len, then re-converting. For the hashes to match the spilled
|
||||
# node.hash_value (computed over the full-sequence bigram key), this MUST hold:
|
||||
# convert_to_bigram_key(fill_ids[M:]) == convert_to_bigram_key(fill_ids)[M:]
|
||||
# i.e. the boundary bigram (f[M-1], f[M]) belongs to the matched prefix (index M-1), not the suffix.
|
||||
fill_ids = list(range(100))
|
||||
full = convert_to_bigram_key(fill_ids)
|
||||
for m in (0, 1, 7, 50, 98): # M <= len-1
|
||||
self.assertEqual(convert_to_bigram_key(fill_ids[m:]), full[m:])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
Reference in New Issue
Block a user