diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index 9b022ddd1..9754fe77c 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/python/sglang/srt/mem_cache/hiradix_cache.py @@ -1045,6 +1045,9 @@ class HiRadixCache(RadixCache): # pooled shared-L2 (server_args validation already enforced the prerequisites + mutual exclusion). self.enable_cp_l3 = bool(getattr(server_args, "enable_cp_l3", False)) and self._uses_cp_hicache self.cp_l3_store = None # constructed after the host pool + allocator (slab accessors need them) + self.ongoing_l3_spill = {} # op_id -> pinned (protect_host'd) resident TreeNode being copied to L3 + # backpressure: cap concurrent background spills (bounds the bg queue + staging RAM, rank-uniform) + self.cp_l3_spill_max_inflight = 32 ( extra_config, @@ -1276,8 +1279,80 @@ class HiRadixCache(RadixCache): qsizes, op=torch.distributed.ReduceOp.MIN, group=self.tp_group, tag="l3_queue_min" ) n_spill, n_reload = map(int, qsizes.tolist()) - store.drain_spill_acks(n_spill) - store.drain_reload_acks(n_reload) + for op_id, _ok in store.drain_spill_acks(n_spill): + # spill done (ok or not): unpin the resident node. ok=True -> object now L3-durable + L2-resident; + # ok=False -> the copy failed, just unpin (the maintainer retries it next pass). No allocator + # release here (approach D: the object was never freed -- it's a copy, not a move). + node = self.ongoing_l3_spill.pop(op_id, None) + if node is not None: + if _ok: + node.l3_durable = True # now reloadable from L3; maintainer skips it (O(1)) + node.release_host() + store.drain_reload_acks(n_reload) # 3.2 wires reload-ack admission + pin release + + def _cp_l3_spill_maintainer(self) -> None: + """Background spill (approach D, 3.1): COPY the coldest resident committed objects (replicated SLRU + order) that are not yet in L3 to the disk tier, off the critical path. The object stays L2-resident; + protect_host pins the LIVE node during the async copy (so the slab can't be evicted underneath the bg + gather). Rank-synced by construction: selection rides replicated state, so all ranks submit the SAME + objects in lockstep -> object-granular acks + the MIN-drain stay aligned. Bounded per tick + backpressured.""" + store = self.cp_l3_store + if store is None: + return + budget = self.cp_l3_spill_max_inflight - len(self.ongoing_l3_spill) + if budget <= 0: + return + allocator = getattr(self.cache_controller, "cp_shared_l2_page_allocator", None) + if allocator is None: + return + # coldest-first = the eviction order, so we make objects durable just before they'd be evicted. + candidates = sorted(self.evictable_host_leaves, key=self._cp_host_evict_key) + submitted = 0 + for node in candidates: + if submitted >= budget: + break + if getattr(node, "l3_durable", False): + continue # O(1): already copied to L3 (set on ack / memoized below) + meta = getattr(node, "cp_hicache", None) + if meta is None or not self._is_cp_shared_l2_metadata(meta): + continue + if int(getattr(node, "host_ref_counter", 0) or 0) > 0: + continue # pinned (already spilling / protected) -> skip (excludes in-flight from re-select) + object_key = getattr(meta, "object_key", "") + if not object_key or not allocator.is_committed(object_key): + continue # M-1: only committed objects (owned-page D2H complete) + page_keys = getattr(node, "hash_value", None) + if not page_keys: + continue + required = tuple(getattr(meta, "required_payloads", ())) + if store.exists_prefix(page_keys, required) >= len(page_keys): + node.l3_durable = True # memoize (e.g. a reloaded-from-L3 node) -> O(1) skip next tick + continue # already fully durable in L3 (dedup) + owned_pages = self._cp_l3_build_owned_pages(meta, page_keys, required, store) + last_access = int(getattr(node, "last_access_time", 0) or 0) + node.protect_host() + op_id = store.submit_spill(object_key, owned_pages, last_access) + self.ongoing_l3_spill[op_id] = node + submitted += 1 + + def _cp_l3_build_owned_pages(self, meta, page_keys, required, store): + """Per payload: this rank's owned (slab_page, content_hash) pairs. Owner = object-LOCAL page index + i % cp_size (verified vs _shared_l2_current_page_owners = owner_for_logical_pages(arange(1,n+1))). + Logical pages only (len(page_keys)); padding has no content hash. slab_page = payload range base_page+i.""" + cp_size, cp_rank = store.cp_size, store.cp_rank + n = len(page_keys) + owned = {} + for pk in required: + if pk not in store.payloads: + continue + rng = meta.object_ranges.get(pk) + if rng is None: + continue + base = rng.base_page + owned[pk] = [ + (base + i, page_keys[i]) for i in range(n) if i % cp_size == cp_rank + ] + return owned def _cp_hicache_all_reduce(self, tensor, *, op, group, tag: str) -> None: start = time.perf_counter() @@ -5073,6 +5148,7 @@ class HiRadixCache(RadixCache): self.drain_storage_control_queues() if self.enable_cp_l3: self._drain_l3_control_queues() + self._cp_l3_spill_maintainer() if self.enable_storage_metrics: self.storage_metrics_collector.log_storage_metrics( self.cache_controller.storage_backend.get_stats()