L3 3.1: background spill maintainer (resident->L3 copy, approach D)
Spill is decoupled from eviction (approach D, user-confirmed): a rank-synced background maintainer at check_hicache_events COPIES the coldest resident committed objects (replicated SLRU order, not-yet-in-L3, unpinned) to L3 off the critical path. The object stays L2-resident; protect_host pins the LIVE node during the async bg gather (no freed-range hazard). owned_pages owner = object-local i%cp_size (verified vs _shared_l2_current_page_owners); per-payload base_page from object_ranges; logical pages only. Bounded + backpressured (cp_l3_spill_max_inflight). Replicated selection -> lockstep object-granular acks. Spill-ack (in _drain_l3_control_queues): release_host + node.l3_durable=True (O(1) skip; memoized from exists_prefix for reloaded nodes). Eviction unchanged (drops L2; L3 copy survives -> reload in 3.2; not-yet-spilled = recompute, fail-soft). Compiles; default-off. Live validation via flag-on smoke (branch sync). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
@@ -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()
|
||||
|
||||
Reference in New Issue
Block a user