From 864b1c808ee925cf5a705fdac2b6c5187e53eb73 Mon Sep 17 00:00:00 2001 From: leavelet Date: Tue, 23 Jun 2026 00:09:38 +0000 Subject: [PATCH] L3 3.2: async disk->L2 reload + request hold (Model B), zero new collectives, EAGLE-correct MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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) --- python/sglang/srt/managers/scheduler.py | 25 ++ python/sglang/srt/mem_cache/hiradix_cache.py | 374 ++++++++++++++++-- .../mem_cache/test_convert_to_bigram_key.py | 11 + 3 files changed, 384 insertions(+), 26 deletions(-) diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 5a56e8fbc..5dccdfe11 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -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: diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index 4270dd58f..bf01368e6 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/python/sglang/srt/mem_cache/hiradix_cache.py @@ -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 diff --git a/test/registered/unit/mem_cache/test_convert_to_bigram_key.py b/test/registered/unit/mem_cache/test_convert_to_bigram_key.py index d34c9f28c..74d68a23d 100644 --- a/test/registered/unit/mem_cache/test_convert_to_bigram_key.py +++ b/test/registered/unit/mem_cache/test_convert_to_bigram_key.py @@ -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()