diff --git a/python/sglang/srt/mem_cache/cp_l3_store.py b/python/sglang/srt/mem_cache/cp_l3_store.py index 790328e12..d4260e750 100644 --- a/python/sglang/srt/mem_cache/cp_l3_store.py +++ b/python/sglang/srt/mem_cache/cp_l3_store.py @@ -17,6 +17,7 @@ index's replicated clock. The store is rank-local for I/O (own disk); only the L from __future__ import annotations +import heapq import logging import os import queue @@ -137,6 +138,17 @@ class CpL3Store: self._ongoing_spill: set = set() # submitted .. durable-acked (the inflight-budget basis + idle gate) self._ongoing_reload: set = set() self._write_err_count = 0 # rate-limit the write-error log (slab-full pre-3.3 streams one/commit) + # ---- 3.3 GC: continuous, background, page-level LRU, ON THE WRITE THREAD (sole pool owner -> no lock). + # Per payload: a lazy-deletion min-heap of (last_access, seq, hash) + an authoritative {hash -> (last_access, + # slot)} map. Add on write, touch (bump last_access) via _touch_q from reload-admit, reclaim coldest when the + # rank-local pool is over the start watermark. Rank-local + owner-partitioned -> zero collectives. (ยง9.6) + self._gc_heap: Dict[str, list] = {pk: [] for pk in self.payloads} + self._gc_current: Dict[str, dict] = {pk: {} for pk in self.payloads} + self._gc_seq = 0 + self._touch_q: "queue.Queue[Tuple[str, str, int]]" = queue.Queue() # (hash_hex, payload_kind, last_access) + self._gc_start_frac = 0.90 # begin reclaiming when a pool exceeds this fraction (tune at impl) + self._gc_stop_frac = 0.85 # reclaim coldest down to this fraction (hysteresis avoids per-write churn) + self._gc_budget = 4096 # max reclaims per pass (bounds a single GC pass so it never starves writes) self._stop = threading.Event() self._gather_thread: Optional[threading.Thread] = None self._write_thread: Optional[threading.Thread] = None @@ -305,6 +317,10 @@ class CpL3Store: non-durable (fail-soft: recompute on a future miss); the thread survives. The staged buffers are freed here (success or fail).""" while not self._stop.is_set(): + # 3.3 GC (continuous, background): apply reload touches, then reclaim coldest pages if any pool is + # over the start watermark -- BEFORE the next write, so spill never hits a full pool (no evict burst). + self._gc_drain_touches() + self._gc_collect() try: staged = self._write_q.get(block=True, timeout=0.2) except queue.Empty: @@ -408,6 +424,68 @@ class CpL3Store: for pk in touched: self.slabs[pk].fdatasync() # (write_batch __exit__ committed the index durably here) + # Register the now-durable pages in the GC LRU (write-thread-owned; last_access = the spilled object's + # replicated clock). Touched later via _touch_q on reload-admit; reclaimed coldest-first by _gc_collect. + for (sp, slot) in to_write: + self._gc_add(sp.pk, sp.h, slot, staged.last_access) + + # ---------- 3.3 GC: continuous, background, page-level LRU (write-thread-owned; rank-local; no collective) ---- + def submit_touch(self, touches: Sequence[Tuple[str, str, int]]) -> None: + """Scheduler (reload-admit) -> write thread: bump the LRU last_access for reloaded pages so a reloaded + prefix stays warm in L3. Non-blocking; applied on the write thread (the sole GC owner).""" + for t in touches: + self._touch_q.put(t) + + def _gc_add(self, pk: str, h: bytes, slot: int, last_access: int) -> None: + self._gc_seq += 1 + self._gc_current[pk][h] = (last_access, slot) + heapq.heappush(self._gc_heap[pk], (last_access, self._gc_seq, h)) + + def _gc_drain_touches(self) -> None: + while True: + try: + h_hex, pk, last_access = self._touch_q.get_nowait() + except queue.Empty: + break + cur = self._gc_current.get(pk) + if cur is None: + continue + h = content_hash_to_bytes(h_hex) + entry = cur.get(h) + if entry is None: + continue # not in L3 here (never spilled on this rank, or already GC'd) -> nothing to bump + self._gc_seq += 1 + cur[h] = (last_access, entry[1]) # keep slot, bump last_access + heapq.heappush(self._gc_heap[pk], (last_access, self._gc_seq, h)) + + def _gc_collect(self) -> None: + """Reclaim coldest owned pages of any pool over the start watermark, down to the stop watermark. Rank- + local (this rank's own pools) -> the owner +/-1 imbalance is self-managed, no cross-rank trigger.""" + for pk in self.payloads: + pool = self.pools[pk] + if pool.num_allocated > int(self._gc_start_frac * pool.num_slots): + self._gc_reclaim(pk, int(self._gc_stop_frac * pool.num_slots)) + + def _gc_reclaim(self, pk: str, target: int) -> None: + """Pop the coldest LIVE owned pages of `pk` down to `target` (or the per-pass budget) and reclaim them: + delete the index entry (durable) BEFORE freeing the slot, so a racing reload reads None (miss), never a + reused slot. Stale heap entries (touched/evicted since pushed) are discarded. Write-thread only -> no lock.""" + pool, heap, cur = self.pools[pk], self._gc_heap[pk], self._gc_current[pk] + to_free: List[Tuple[bytes, int]] = [] + while pool.num_allocated - len(to_free) > target and heap and len(to_free) < self._gc_budget: + last_access, _seq, h = heapq.heappop(heap) + entry = cur.get(h) + if entry is None or entry[0] != last_access: + continue # stale: already reclaimed, or superseded by a newer touch + to_free.append((h, entry[1])) + del cur[h] + if not to_free: + return + with self.index.write_batch() as wb: # delete entries durably FIRST (a racing reload then misses cleanly) + for (h, _slot) in to_free: + wb.delete(h, pk) + for (_h, slot) in to_free: + pool.free(slot) def _reload_object(self, op: L3ReloadOp) -> bool: """Read this rank's owned pages from L3 and scatter into the reserved L2 dest pages. Returns False if @@ -458,12 +536,16 @@ class CpL3Store: break for q in ( self._gather_q, self._write_q, self._ack_gather_q, self._ack_durable_q, - self._reload_q, self._ack_reload_q, + self._reload_q, self._ack_reload_q, self._touch_q, ): with q.mutex: q.queue.clear() self._ongoing_spill.clear() self._ongoing_reload.clear() + for pk in self.payloads: # reset the GC LRU (the slot pools + index are reset below) + self._gc_heap[pk].clear() + self._gc_current[pk].clear() + self._gc_seq = 0 for pool in self.pools.values(): pool.reset() self.index.clear() diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index bf01368e6..b79eb497a 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/python/sglang/srt/mem_cache/hiradix_cache.py @@ -1579,8 +1579,24 @@ class HiRadixCache(RadixCache): 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: + node = self._cp_l3_insert_reloaded_node(info) + inserted = node is not None + if inserted: + # 3.3 GC: bump the L3 last_access of THIS rank's owned reloaded pages (owner = i%cp_size), so + # the reloaded prefix stays warm in L3 and isn't GC'd as cold. Touch ALL reloaded pages (not + # just the tail) so the prefix ages uniformly. last_access = the inserted node's replicated clock. + store = self.cp_l3_store + la = int(node.last_access_time) + shashes = info["suffix_hashes"] + touches = [ + (shashes[i], pk, la) + for pk in store.payloads + for i in range(len(shashes)) + if i % store.cp_size == store.cp_rank + ] + if touches: + store.submit_touch(touches) + else: allocator.abort(info["reload_key"]) elif allocator is not None: allocator.abort(info["reload_key"]) # failed reload / stale attach -> free the reservation @@ -1590,11 +1606,12 @@ class HiRadixCache(RadixCache): else: self.cp_l3_reload_done.pop(rid, None) # released -> proceed (recompute) - def _cp_l3_insert_reloaded_node(self, info) -> bool: + def _cp_l3_insert_reloaded_node(self, info) -> Optional[TreeNode]: """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).""" + node is L3-durable (just read from L3). Returns the new node (its last_access_time seeds the GC touch), + or None if the attach raced to non-clean. Rank-uniform (replicated radix + rank-uniform inputs).""" from sglang.srt.mem_cache.cp_shared_l2_pool import CpSharedL2NodeMetadata lhn = info["last_host_node"] @@ -1602,10 +1619,10 @@ class HiRadixCache(RadixCache): 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 + return None 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 + return None # 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 @@ -1613,7 +1630,7 @@ class HiRadixCache(RadixCache): 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 + return None n_pages = info["n_pages"] meta = CpSharedL2NodeMetadata( logical_len=n_pages * self.page_size, @@ -1637,7 +1654,7 @@ class HiRadixCache(RadixCache): self._cp_account_enter_committed(new_node.cp_hicache) self._update_host_leaf_status(new_node) self._update_host_leaf_status(lhn) - return True + return new_node 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 diff --git a/test/registered/unit/mem_cache/test_cp_l3_store.py b/test/registered/unit/mem_cache/test_cp_l3_store.py index e09bfe13e..93c063e68 100644 --- a/test/registered/unit/mem_cache/test_cp_l3_store.py +++ b/test/registered/unit/mem_cache/test_cp_l3_store.py @@ -93,12 +93,13 @@ class TestCpL3Store(unittest.TestCase): self.store.close() self._td.cleanup() - def _finish_spill(self): + def _finish_spill(self, store=None): """Drive the 2-phase spill to completion (store-level; no cross-rank MIN). Returns durable acks as [(op_id, ok), ...]. Draining the gather-ack is what releases the eviction pin in production.""" - self.assertTrue(_wait_ack(self.store.ack_durable_qsize, 1)) - self.store.drain_gather_acks(self.store.ack_gather_qsize()) - return self.store.drain_durable_acks(self.store.ack_durable_qsize()) + st = store if store is not None else self.store + self.assertTrue(_wait_ack(st.ack_durable_qsize, 1)) + st.drain_gather_acks(st.ack_gather_qsize()) + return st.drain_durable_acks(st.ack_durable_qsize()) def test_two_phase_gather_precedes_durable(self): # the core 3.1 fix: gather (pin release) acks BEFORE the slow durable write; inflight until durable. @@ -200,6 +201,77 @@ class TestCpL3Store(unittest.TestCase): self.assertEqual(pool.num_free, 1) # no orphan: the partial alloc was rolled back self.assertFalse(store.has_inflight()) + def test_gc_reclaims_coldest_to_stop_watermark(self): + # 3.3 continuous-background GC: filling a pool past the start watermark triggers the write-thread GC to + # reclaim the COLDEST (lowest last_access) owned pages down to the stop watermark, deleting their index + # entries, with no slot leak. Small pool so the watermark trips deterministically. + td = tempfile.TemporaryDirectory() + self.addCleanup(td.cleanup) + cfg = cfg_mod.CpL3Config.from_dict({ + "backend": "posix", "require_plp": False, "index_map_gb": 0.05, + "disks": [{"path": os.path.join(td.name, "d"), "budget_gb": 0.0001}], + }) + store = store_mod.CpL3Store.from_config( + cfg, cp_rank=0, cp_size=1, accessors={"target_kv": self.acc}) + store.connect(cfg) + self.addCleanup(store.close) + pool = store.pools["target_kv"] + n = pool.num_slots + self.assertGreaterEqual(n, 8) + start = int(store._gc_start_frac * n) + stop = int(store._gc_stop_frac * n) + # spill n distinct pages, page i with last_access=i (so smaller i == colder) + for i in range(n): + store.submit_spill(f"o{i}", {"target_kv": [(i % PAGE_NUM, _h(3000 + i))]}, last_access=i) + self._finish_spill(store) + # the write-thread GC keeps usage within the watermark band [stop, start] (reclaim-to-stop on trigger) + t0 = time.time() + while pool.num_allocated > start and time.time() - t0 < 5.0: + time.sleep(0.01) + allocated = pool.num_allocated + self.assertLessEqual(allocated, start) + self.assertGreaterEqual(allocated, stop) + self.assertEqual(pool.num_free + allocated, n) # no slot leak + n_evicted = n - allocated + self.assertGreater(n_evicted, 0) + for i in range(n_evicted): # the coldest n_evicted pages were reclaimed (index entries deleted) + self.assertEqual(store.exists_prefix([_h(3000 + i)], ["target_kv"]), 0) + for i in range(n_evicted, n): # the warmest survive + self.assertEqual(store.exists_prefix([_h(3000 + i)], ["target_kv"]), 1) + + def test_gc_touch_keeps_a_page_warm(self): + # submit_touch bumps a page's last_access so it survives GC while colder untouched pages are reclaimed. + td = tempfile.TemporaryDirectory() + self.addCleanup(td.cleanup) + cfg = cfg_mod.CpL3Config.from_dict({ + "backend": "posix", "require_plp": False, "index_map_gb": 0.05, + "disks": [{"path": os.path.join(td.name, "d"), "budget_gb": 0.0001}], + }) + store = store_mod.CpL3Store.from_config( + cfg, cp_rank=0, cp_size=1, accessors={"target_kv": self.acc}) + store.connect(cfg) + self.addCleanup(store.close) + pool = store.pools["target_kv"] + n = pool.num_slots + start = int(store._gc_start_frac * n) + stop = int(store._gc_stop_frac * n) + self.assertGreaterEqual(start, 4) + # fill below the start watermark (no GC yet), page i last_access=i + for i in range(start): + store.submit_spill(f"o{i}", {"target_kv": [(i % PAGE_NUM, _h(4000 + i))]}, last_access=i) + self._finish_spill(store) + # touch the coldest page (i=0) to the warmest -> it must survive the upcoming GC + store.submit_touch([(_h(4000 + 0), "target_kv", 10_000)]) + # now push over the start watermark -> GC reclaims coldest down to stop + for i in range(start, n + 1): + store.submit_spill(f"o{i}", {"target_kv": [(i % PAGE_NUM, _h(4000 + i))]}, last_access=i) + self._finish_spill(store) + t0 = time.time() + while pool.num_allocated > start and time.time() - t0 < 5.0: + time.sleep(0.01) + self.assertEqual(store.exists_prefix([_h(4000 + 0)], ["target_kv"]), 1) # touched -> warm -> survived + self.assertEqual(store.exists_prefix([_h(4000 + 1)], ["target_kv"]), 0) # next-coldest -> reclaimed + def test_clear_resets(self): self.store.submit_spill("o", {"target_kv": [(0, _h(0))]}, last_access=1) self._finish_spill()