diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index 60dcd7b7e..537dac798 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/python/sglang/srt/mem_cache/hiradix_cache.py @@ -2943,6 +2943,95 @@ class HiRadixCache(RadixCache): ) return True + def _reserve_write_cp_shared_l2_evict_to_fit( + self, device_indices: torch.Tensor, node_id: int, last_result + ): + """B1 (2b.2) deterministic collective-free evict-to-fit for the shared L2 pool. + + Snapshot the plannable shared-L2 host leaves (excluding the node being backed + up) in the replicated SLRU victim order (``_cp_host_evict_key``), then release + them one at a time -- each release is a replicated free-list mutation -- and + retry the pooled reserve after each. Fragmentation is handled by construction: + the retry succeeds as soon as a contiguous run reopens (free coalesces in + ``_return_range``). Identical victim set + order + deterministic reserve => + every rank evicts the same victims and the same retry succeeds at the same + point (design Thm 1), so NO collective is needed -- this is the replicated + replacement for the #5 ``_evict_host_for_physical_slots(synchronize_across_ranks=True)`` + all_reduce. If the evictable set is exhausted and the reserve still cannot + fit, the backup is skipped (transient: live/pinned/referenced objects hold the + pool, or the single backup exceeds pool capacity) -- a loud, rate-limited + warn, never a hang. + + The victim snapshot is one-level by design: a parent promoted to a new host + leaf when its child is removed is NOT re-pushed mid-loop (unlike the + heap-based _evict_host_for_physical_slots), so deep cascades skip-and-retry + on the next tick's snapshot rather than this round -- missed identically on + every rank, so determinism (Thm 1) is unaffected. + """ + victims = sorted( + ( + node + for node in getattr(self, "evictable_host_leaves", set()) + if getattr(node, "id", None) != node_id + and self._cp_host_leaf_is_plannable_victim(node) + and self._is_cp_shared_l2_metadata(getattr(node, "cp_hicache", None)) + ), + key=self._cp_host_evict_key, + ) + evicted_ids = [] + for victim in victims: + # Re-check plannability immediately before mutating (fail-loud on any + # divergence from the snapshot -- the determinism guard mirrored from + # _evict_cp_host_for_write_admission). + if not self._cp_host_leaf_is_plannable_victim(victim): + raise RuntimeError( + "CP shared-L2 evict-to-fit plan diverged before reservation: " + f"node_id={node_id} victim_id={getattr(victim, 'id', None)}" + ) + # Teardown identical to _evict_cp_host_for_write_admission: evict_cp_host + # releases the shared pool range (-> allocator.release) and drops the + # replicated host leaf. All inputs are in R => all ranks evict in lockstep. + self._record_remove_event(victim) + self._cp_account_leave_node_evict(victim, victim.cp_hicache) + self.cache_controller.evict_cp_host(victim.cp_hicache) + evicted_ids.append(getattr(victim, "id", None)) + victim.host_len = 0 + victim.cp_hicache = None + victim.host_value = None + self._remove_host_leaf(victim) + + result = self.cache_controller.reserve_write_cp( + device_indices=device_indices, + node_id=node_id, + ) + if not isinstance(result, HiCacheWriteFailure): + logger.debug( + "[HiCache-evict] CP shared-L2 evict-to-fit succeeded: " + "node_id=%d evicted=%s", + node_id, + evicted_ids, + ) + return result + if getattr(result, "reason", "host_capacity") != "shared_l2_capacity": + # A non-capacity failure (e.g. undrained_ack) cannot be helped by + # more eviction. + return result + last_result = result + + # Exhausted every evictable shared-L2 leaf and the reserve still does not fit: + # the pool is held by non-evictable (pinned / in-flight / referenced) objects, + # or this single backup exceeds pool capacity. Skip this round (the backup + # retries next tick); loud + rate-limited so a real shortfall stays visible. + self._warn_cp_hicache_fallback( + "cp_shared_l2_evict_exhausted", + "shared_l2_capacity_after_evict", + rate_limit_s=10.0, + node_id=node_id, + required_host_slots=int(getattr(last_result, "required_host_slots", 0)), + evicted=len(evicted_ids), + ) + return last_result + def _reserve_write_cp_indices_no_collective( self, device_indices: torch.Tensor, @@ -2962,15 +3051,16 @@ class HiRadixCache(RadixCache): if getattr(result, "reason", "host_capacity") != "shared_l2_capacity": return result - # B1 (2b.1b): on a shared-L2 capacity miss, SKIP this backup (the - # existing reactive behavior; reserve-at-admission + background rotation - # [2a] keep this off the steady-state path). Do NOT call - # _evict_host_for_physical_slots(synchronize_across_ranks=True) here -- - # its `while len(heap) and not all_ranks_done()` loop is the #5 collective - # deadlock shape (per-rank len(heap)/wall-clock). The collective-free - # deterministic shared-pool eviction-and-retry is 2b.2; until it lands, - # capacity miss degrades to skip, never a collective. - return result + # B1 (2b.2): a shared-L2 capacity miss raises identically on every CP + # rank (replicated free list), so run the deterministic collective-free + # evict-to-fit -- release the coldest plannable shared-L2 host leaves in + # the replicated SLRU order and retry the pooled reserve. Every rank + # frees the same victims in the same order and the same retry succeeds at + # the same point (placement determinism, design Thm 1) -- NO collective, + # NO _evict_host_for_physical_slots(synchronize_across_ranks=True) (#5). + return self._reserve_write_cp_shared_l2_evict_to_fit( + device_indices, node_id, result + ) if not admission_checked: admission = self._cp_build_write_admission( @@ -4309,22 +4399,20 @@ class HiRadixCache(RadixCache): self._update_host_leaf_status(parent) return parent - def _evict_host_for_physical_slots( - self, required_host_slots: int, synchronize_across_ranks: bool = False - ) -> int: - synchronize_across_ranks = ( - synchronize_across_ranks - and getattr(self, "tp_group", None) is not None - and getattr(self, "tp_world_size", 1) > 1 - ) - if required_host_slots <= 0 and not synchronize_across_ranks: + def _evict_host_for_physical_slots(self, required_host_slots: int) -> int: + # B1 (2b.2): the #5 synchronize_across_ranks all_reduce (`while len(heap) and + # not all_ranks_done()` gated on a ReduceOp.MIN host_evict_done_min) is + # DELETED. It was never reached (no caller passed synchronize_across_ranks=True) + # and was the headline CP-deadlock shape (a collective gated on per-rank + # len(heap)). The shared-L2 capacity path now uses the deterministic + # replicated _reserve_write_cp_shared_l2_evict_to_fit instead. + if required_host_slots <= 0: return 0 leaves = list(self.evictable_host_leaves) logger.debug( - "[HiCache-evict] _evict_host_for_physical_slots: required_slots=%d sync=%s leaves=%d", + "[HiCache-evict] _evict_host_for_physical_slots: required_slots=%d leaves=%d", required_host_slots, - synchronize_across_ranks, len(leaves), ) eviction_heap = [ @@ -4334,20 +4422,7 @@ class HiRadixCache(RadixCache): num_evicted = 0 - def all_ranks_done() -> bool: - local_done = int(num_evicted >= required_host_slots) - if not synchronize_across_ranks: - return bool(local_done) - done = torch.tensor(local_done, dtype=torch.int, device="cpu") - self._cp_hicache_all_reduce( - done, - op=torch.distributed.ReduceOp.MIN, - group=self.tp_group, - tag="host_evict_done_min", - ) - return bool(done.item()) - - while len(eviction_heap) and not all_ranks_done(): + while len(eviction_heap) and num_evicted < required_host_slots: _priority, x = heapq.heappop(eviction_heap) if x == self.root_node: break diff --git a/test/registered/unit/mem_cache/test_cp_shared_l2_pool.py b/test/registered/unit/mem_cache/test_cp_shared_l2_pool.py index 9e17c0cba..f00ef3b5e 100644 --- a/test/registered/unit/mem_cache/test_cp_shared_l2_pool.py +++ b/test/registered/unit/mem_cache/test_cp_shared_l2_pool.py @@ -2779,6 +2779,51 @@ class TestCpSharedL2EightRankReserveDeterminism(unittest.TestCase): allocs[3].reserve("cp_hicache_node:1", PAYLOAD_TARGET_KV, 1) # rank 3 only self.assertGreater(len({a.placement_digest() for a in allocs}), 1) + def test_eight_ranks_evict_to_fit_stays_identical(self): + # B1 (2b.2): on a shared-L2 capacity miss every rank runs the SAME + # deterministic evict-to-fit -- release the coldest victims (here the + # ascending-node-id analog of the replicated SLRU order) one at a time, + # retrying the pooled reserve after each, stopping at the first success. + # Identical free list + identical victim order + deterministic reserve => + # every rank evicts the same victims and stops at the same point (design + # Thm 1), so the hiradix _reserve_write_cp_shared_l2_evict_to_fit needs NO + # collective. Also exercises fragmentation: the win comes only after the + # freed ranges COALESCE into a contiguous run. + allocs = self.make_rank_allocators(pages=8) + for nid, npages in ((0, 3), (1, 3), (2, 2)): # fill 3+3+2 = 8/8 pages + self._reserve_on_all(allocs, nid, (PAYLOAD_TARGET_KV,), npages) + for a in allocs: + a.mark_object_committed(f"cp_hicache_node:{nid}") + self.assertEqual(len({a.placement_digest() for a in allocs}), 1) + + # New backup needs 4 contiguous pages -> capacity miss on EVERY rank. + for a in allocs: + with self.assertRaises(ValueError): + a.reserve("cp_hicache_node:9", PAYLOAD_TARGET_KV, 4) + a.abort("cp_hicache_node:9") # SF3: drop any partial reserve + + victim_order = [f"cp_hicache_node:{nid}" for nid in (0, 1, 2)] + evicted_counts = [] + for a in allocs: + evicted = 0 + for victim in victim_order: + a.release(victim) + evicted += 1 + try: + a.reserve("cp_hicache_node:9", PAYLOAD_TARGET_KV, 4) + break # contiguous run reopened -> fit + except ValueError: + a.abort("cp_hicache_node:9") + evicted_counts.append(evicted) + + # Every rank evicted the SAME victims (release 0 -> free=3<4 no fit; release + # 1 -> [0,3)+[3,6) coalesce to [0,6) -> 4 fits) and reached identical placement. + self.assertEqual(set(evicted_counts), {2}) + self.assertEqual(len({a.placement_digest() for a in allocs}), 1) + for a in allocs: + self.assertTrue(a.is_committed("cp_hicache_node:2")) # untouched survivor + self.assertIsNotNone(a.get_range("cp_hicache_node:9", PAYLOAD_TARGET_KV)) + if __name__ == "__main__": unittest.main()