From fdc23f7d61cf0b619439f67233dafbb855849e90 Mon Sep 17 00:00:00 2001 From: leavelet Date: Wed, 24 Jun 2026 01:35:37 +0000 Subject: [PATCH] CP HiCache: redesign owner-lane eviction planner as a lazy-re-evaluation heap (O(n log n)) Replaces the O(victims*candidates) per-iteration greedy argmin rescan in _plan_cp_load_back_owner_lane_evictions with a leaf-up min-heap. The score (-contribution, -unlock, slru_priority, node.id) depends on the current deficits, so a static heap is not equivalent; instead use LAZY RE-EVALUATION: a popped entry is consumed only if its score still matches the node's current score (deficits unchanged since push), else it is re-pushed with the fresh score. Stale scores are always optimistic (contribution = sum(min(counts[o], deficits[o])) and the ancestor-unlock contribution only shrink as deficits shrink), so the first up-to-date popped entry is exactly the global argmin the full rescan would have picked -> PROVABLY EQUIVALENT, with the same leaf-up eligibility + ancestor-unlock + parent-push. Determinism/rank-uniformity preserved (selection decided by the total-ordered score; the heap insertion seq only orders structurally-equal tuples). Equivalence proven by test: a verbatim reference greedy + a randomized property test (300 seeds, single- AND multi-owner deficits, non-uniform counts exercising the lazy-re-eval boundary, varied SLRU priorities) asserting byte-identical (victims, planned_freed, remaining), plus an explicit leaf-up + ancestor-unlock case (child-then-parent). 14 planner tests pass (1 pre-existing unrelated EAGLE-tail failure unchanged). Micro-bench (V100, candidates=2000 deficit=7877, benchmark/hicache/bench_cp_owner_lane_planner.py): A. ORIGINAL (.item() x cp, no memo) 50.06s B. memo + bincount (greedy) 0.494s (101x) C. lazy-re-eval heap 0.162s (309x; 3.0x over B) selection A==B==C identical Co-Authored-By: Claude Opus 4.8 (1M context) --- .../hicache/bench_cp_owner_lane_planner.py | 65 +++++-- python/sglang/srt/mem_cache/hiradix_cache.py | 183 ++++++++++-------- .../test_cp_hicache_load_back_owner_lanes.py | 151 +++++++++++++++ 3 files changed, 297 insertions(+), 102 deletions(-) diff --git a/benchmark/hicache/bench_cp_owner_lane_planner.py b/benchmark/hicache/bench_cp_owner_lane_planner.py index 5e5acb875..68430a99c 100644 --- a/benchmark/hicache/bench_cp_owner_lane_planner.py +++ b/benchmark/hicache/bench_cp_owner_lane_planner.py @@ -88,29 +88,54 @@ def plan_greedy(counts_fn, use_memo): def plan_heap(counts_fn): - """Redesign (single-owner deficit): compute counts once, sort by -counts[owner], one pass. - Byte-identical victim SET to greedy for a single-owner deficit (highest owner-count first, - id tiebreak) -- the contribution is min(counts[owner], remaining) == counts[owner] until the - final pick.""" - owner = DEFICIT_OWNER - scored = [] - for nid, value in NODES: - counts = counts_fn(value) - if counts[owner] <= 0: - continue - scored.append((counts, nid)) - scored.sort(key=lambda x: (-x[0][owner], x[1])) - rem = [0] * CP - rem[owner] = DEFICIT + """The SHIPPED heap redesign: lazy-re-evaluation min-heap (a popped entry is used only if + its score is still current, else re-pushed). Provably equivalent to the greedy argmin + rescan; here flat (no tree) so it isolates the algorithmic cost vs B.""" + import heapq + + deficits = [0] * CP + deficits[DEFICIT_OWNER] = DEFICIT + planned_evicted = set() planned_freed = [0] * CP + memo = {} + heap = [] + seq = 0 + + def scored(nid, value): + if nid in memo: + counts = memo[nid] + else: + counts = counts_fn(value) + memo[nid] = counts + contribution = sum(min(c, d) for c, d in zip(counts, deficits)) + if contribution <= 0: + return None + return (-contribution, nid), counts + + for nid, value in NODES: + s = scored(nid, value) + if s is not None: + heapq.heappush(heap, (s[0], seq, nid, s[1])) + seq += 1 + victims = 0 - for counts, nid in scored: - if all(v <= 0 for v in rem): - break + while any(v > 0 for v in deficits) and heap: + score, _s, nid, counts = heapq.heappop(heap) + if nid in planned_evicted: + continue + fresh = scored(nid, None) # counts already memoized + if fresh is None: + continue + if fresh[0] != score: + heapq.heappush(heap, (fresh[0], seq, nid, fresh[1])) + seq += 1 + continue + counts = fresh[1] + planned_evicted.add(nid) + victims += 1 for o, c in enumerate(counts): planned_freed[o] += c - rem[o] = max(0, rem[o] - c) - victims += 1 + deficits[o] = max(0, deficits[o] - c) return victims, tuple(planned_freed) @@ -133,7 +158,7 @@ def main(): vO, fO = timed("A. ORIGINAL (.item() x cp, no memo)", lambda: plan_greedy(counts_original, False)) vF, fF = timed("B. FIXED (bincount + per-plan memo)", lambda: plan_greedy(counts_fixed, True)) - vH, fH = timed("C. REDESIGN (counts-once + sorted pass)", lambda: plan_heap(counts_fixed)) + vH, fH = timed("C. REDESIGN (lazy-re-eval heap)", lambda: plan_heap(counts_fixed)) print("\nselection identical (victims, planned_freed):") print(f" A==B: {(vO, fO) == (vF, fF)} A==C: {(vO, fO) == (vH, fH)}") diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index 2f6efa9f1..9a7f620fc 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/python/sglang/srt/mem_cache/hiradix_cache.py @@ -2583,96 +2583,115 @@ class HiRadixCache(RadixCache): planned_freed = [0 for _ in range(cp_size)] victims: List[TreeNode] = [] planned_evicted_nodes = set() - candidate_nodes = set(getattr(self, "evictable_leaves", set())) - initial_candidate_count = len(candidate_nodes) - iteration = 0 - # Per-plan memo {node.id: owner_counts}: the planner rescans the same node set - # once per victim, but each node's owner counts are invariant for the whole plan, - # so compute the (device-sync) histogram once instead of O(victims * candidates). + initial_candidate_count = len(getattr(self, "evictable_leaves", set())) + # Per-plan memo {node.id: owner_counts}: each node's owner counts are invariant for + # the whole plan, so compute the (device-sync) histogram once. owner_counts_memo: dict = {} - while any(v > 0 for v in deficits): - iteration += 1 - best_node = None - best_counts = None - best_score = None - scanned_candidates = 0 - for node in list(candidate_nodes): - scanned_candidates += 1 - if (scanned_candidates & 1023) == 0: - now = time.perf_counter() - if now - last_progress_time >= 5.0: - logger.warning( - "[HiCache-load] slow CP owner-lane eviction planning scan: " - "iteration=%d scanned=%d candidates=%d victims=%d " - "deficits=%s planned_freed=%s elapsed_ms=%.3f", - iteration, - scanned_candidates, - len(candidate_nodes), - len(victims), - deficits, - planned_freed, - (now - plan_start_time) * 1000.0, - ) - last_progress_time = now - if node in planned_evicted_nodes: - continue - if not self._cp_device_node_is_load_back_victim_after_plan( - node, planned_evicted_nodes - ): - continue - counts = self._cp_load_back_node_owner_page_counts( - node, cp_size, memo=owner_counts_memo + # Heap-based leaf-up eviction, PROVABLY EQUIVALENT to the previous per-iteration + # greedy argmin rescan, via LAZY RE-EVALUATION: a popped entry is consumed only if + # its score still matches the node's CURRENT score (deficits unchanged since it was + # pushed); otherwise it is re-pushed with the fresh score and competes again. Stale + # scores are always OPTIMISTIC -- contribution = sum(min(counts[o], deficits[o])) and + # the ancestor-unlock contribution can only shrink as deficits shrink, so the score + # only worsens -- hence the first up-to-date popped entry is the global argmin the + # full rescan would have picked. Cost O((candidates + repushes) log n) with repushes + # bounded by deficit crossings (~none until a lane's tail) vs O(victims * candidates). + # Determinism/rank-uniformity is preserved: selection is decided entirely by the + # total-ordered score (node.id is the final tiebreak); the heap insertion `seq` only + # orders structurally-equal tuples (and keeps heapq from comparing TreeNode objects) + # and never changes which node is chosen. + heap: list = [] + seq = 0 + + def _scored(node): + counts = self._cp_load_back_node_owner_page_counts( + node, cp_size, memo=owner_counts_memo + ) + contribution = sum( + min(int(count), int(deficit)) + for count, deficit in zip(counts, deficits) + ) + unlock_contribution = 0 + if contribution <= 0: + unlock_contribution = self._cp_load_back_ancestor_unlock_contribution( + node, deficits, planned_evicted_nodes, cp_size, memo=owner_counts_memo ) - contribution = sum( - min(int(count), int(deficit)) - for count, deficit in zip(counts, deficits) - ) - unlock_contribution = 0 - if contribution <= 0: - unlock_contribution = ( - self._cp_load_back_ancestor_unlock_contribution( - node, - deficits, - planned_evicted_nodes, - cp_size, - memo=owner_counts_memo, - ) + if contribution <= 0 and unlock_contribution <= 0: + return None + # get_priority returns the CpReplicatedSLRUStrategy tuple (is_protected, + # last_access_time[logical], node.id) -- rank-replicated; node.id keeps the full + # score a strict total order even if the strategy is ever swapped. + score = ( + -int(contribution), + -int(unlock_contribution), + self.eviction_strategy.get_priority(node), + int(getattr(node, "id", 0) or 0), + ) + return score, counts + + def _push(node): + nonlocal seq + if node in planned_evicted_nodes: + return + if not self._cp_device_node_is_load_back_victim_after_plan( + node, planned_evicted_nodes + ): + return + scored = _scored(node) + if scored is None: + return + heapq.heappush(heap, (scored[0], seq, node, scored[1])) + seq += 1 + + for node in list(getattr(self, "evictable_leaves", set())): + _push(node) + + pops = 0 + while any(v > 0 for v in deficits) and heap: + score, _seq, node, counts = heapq.heappop(heap) + pops += 1 + if (pops & 1023) == 0: + now = time.perf_counter() + if now - last_progress_time >= 5.0: + logger.warning( + "[HiCache-load] slow CP owner-lane eviction planning scan: " + "pops=%d heap=%d victims=%d deficits=%s planned_freed=%s " + "elapsed_ms=%.3f", + pops, + len(heap), + len(victims), + deficits, + planned_freed, + (now - plan_start_time) * 1000.0, ) - if contribution <= 0 and unlock_contribution <= 0: - continue - # get_priority returns the CpReplicatedSLRUStrategy tuple - # (is_protected, last_access_time[logical], node.id) under CP — a - # rank-replicated value; the trailing node.id keeps the full score a - # strict total order even if the strategy is ever swapped. - score = ( - -int(contribution), - -int(unlock_contribution), - self.eviction_strategy.get_priority(node), - int(getattr(node, "id", 0) or 0), - ) - if best_score is None or score < best_score: - best_score = score - best_node = node - best_counts = counts - - if best_node is None or best_counts is None: - break - - victims.append(best_node) - planned_evicted_nodes.add(best_node) - candidate_nodes.discard(best_node) - for owner, count in enumerate(best_counts): + last_progress_time = now + if node in planned_evicted_nodes: + continue # stale entry: node already evicted in this plan + fresh = _scored(node) + if fresh is None: + continue # no longer contributes or unlocks (deficits shrank) -> drop + if fresh[0] != score: + # deficits changed since this entry was pushed -> optimistic stale score; + # re-push at the fresh (>= old) score so it competes again at its true cost. + heapq.heappush(heap, (fresh[0], seq, node, fresh[1])) + seq += 1 + continue + # up-to-date heap minimum == the current global argmin -> evict it. + counts = fresh[1] + victims.append(node) + planned_evicted_nodes.add(node) + for owner, count in enumerate(counts): planned_freed[owner] += int(count) deficits[owner] = max(0, deficits[owner] - int(count)) - ancestor = getattr(best_node, "parent", None) + ancestor = getattr(node, "parent", None) while ancestor is not None and ancestor != getattr(self, "root_node", None): if ancestor in planned_evicted_nodes: break if self._cp_device_node_is_load_back_victim_after_plan( ancestor, planned_evicted_nodes ): - candidate_nodes.add(ancestor) + _push(ancestor) break if getattr(ancestor, "value", None) is not None: break @@ -2682,13 +2701,13 @@ class HiRadixCache(RadixCache): if plan_elapsed_ms >= 1000.0 or any(v > 0 for v in deficits): logger.warning( "[HiCache-load] CP owner-lane eviction planning done: " - "elapsed_ms=%.3f initial_candidates=%d remaining_candidates=%d " - "iterations=%d victims=%d original_deficit_by_owner=%s " + "elapsed_ms=%.3f initial_candidates=%d remaining_heap=%d " + "pops=%d victims=%d original_deficit_by_owner=%s " "planned_freed_by_owner=%s remaining_deficit_by_owner=%s", plan_elapsed_ms, initial_candidate_count, - len(candidate_nodes), - iteration, + len(heap), + pops, len(victims), plan.deficit_by_owner, planned_freed, diff --git a/test/registered/unit/mem_cache/test_cp_hicache_load_back_owner_lanes.py b/test/registered/unit/mem_cache/test_cp_hicache_load_back_owner_lanes.py index fdc648c7a..ef145879f 100644 --- a/test/registered/unit/mem_cache/test_cp_hicache_load_back_owner_lanes.py +++ b/test/registered/unit/mem_cache/test_cp_hicache_load_back_owner_lanes.py @@ -92,6 +92,7 @@ from sglang.srt.mem_cache.allocator import CPSharedPagedTokenToKVPoolAllocator import sglang.srt.mem_cache.common as mem_cache_common from sglang.srt.mem_cache.hiradix_cache import ( CpHiCacheNodeMetadata, + CpLoadBackPlan, HiRadixCache, ) from sglang.srt.mem_cache.radix_cache import ( @@ -273,6 +274,91 @@ def _make_tiny_eagle_req(cache, allocator, *, seq_len, req_pool_idx): return req +def _reference_greedy_plan(cache, plan): + """Verbatim copy of the PRE-HEAP greedy owner-lane eviction planner -- the equivalence + oracle for the heap rewrite (per-iteration full argmin rescan). Returns + (victim_ids, planned_freed, remaining_deficit) to compare against the production planner.""" + deficits = [max(0, int(v)) for v in plan.deficit_by_owner] + cp_size = len(deficits) + planned_freed = [0 for _ in range(cp_size)] + victims = [] + planned_evicted_nodes = set() + candidate_nodes = set(getattr(cache, "evictable_leaves", set())) + memo = {} + while any(v > 0 for v in deficits): + best_node = best_counts = best_score = None + for node in list(candidate_nodes): + if node in planned_evicted_nodes: + continue + if not cache._cp_device_node_is_load_back_victim_after_plan( + node, planned_evicted_nodes + ): + continue + counts = cache._cp_load_back_node_owner_page_counts(node, cp_size, memo=memo) + contribution = sum( + min(int(c), int(d)) for c, d in zip(counts, deficits) + ) + unlock = 0 + if contribution <= 0: + unlock = cache._cp_load_back_ancestor_unlock_contribution( + node, deficits, planned_evicted_nodes, cp_size, memo=memo + ) + if contribution <= 0 and unlock <= 0: + continue + score = ( + -int(contribution), + -int(unlock), + cache.eviction_strategy.get_priority(node), + int(getattr(node, "id", 0) or 0), + ) + if best_score is None or score < best_score: + best_score, best_node, best_counts = score, node, counts + if best_node is None: + break + victims.append(best_node) + planned_evicted_nodes.add(best_node) + candidate_nodes.discard(best_node) + for owner, count in enumerate(best_counts): + planned_freed[owner] += int(count) + deficits[owner] = max(0, deficits[owner] - int(count)) + ancestor = getattr(best_node, "parent", None) + while ancestor is not None and ancestor != getattr(cache, "root_node", None): + if ancestor in planned_evicted_nodes: + break + if cache._cp_device_node_is_load_back_victim_after_plan( + ancestor, planned_evicted_nodes + ): + candidate_nodes.add(ancestor) + break + if getattr(ancestor, "value", None) is not None: + break + ancestor = getattr(ancestor, "parent", None) + return ( + tuple(n.id for n in victims), + tuple(planned_freed), + tuple(deficits), + ) + + +def _plan_for_deficit(deficit): + return CpLoadBackPlan( + page_owners=[], + required_by_owner=[], + available_by_owner=[], + deficit_by_owner=list(deficit), + free_room_deficit_by_owner=list(deficit), + host_hit_len=0, + ) + + +def _planner_result_tuple(result): + return ( + tuple(n.id for n in result.victims), + tuple(result.planned_freed_by_owner), + tuple(result.remaining_deficit_by_owner), + ) + + class TestCpHiCacheLoadBackOwnerLanes(CustomTestCase): def test_load_back_plan_reports_owner_lane_vectors(self): allocator = _make_allocator() @@ -365,6 +451,71 @@ class TestCpHiCacheLoadBackOwnerLanes(CustomTestCase): first, ) + def test_heap_planner_matches_greedy_reference_randomized(self): + # The heap planner must select byte-identically to the previous greedy argmin + # rescan across randomized flat scenarios -- crucially including MULTI-OWNER + # deficits and non-uniform per-node counts (which exercise the lazy-re-evaluation + # boundary, where a node's contribution = min(counts[o], deficits[o]) changes as a + # lane's deficit is consumed) and varied priorities (the SLRU tiebreak). + import random + + cp_size = 4 + multi_owner_seen = 0 + for seed in range(300): + rng = random.Random(seed) + allocator = _make_allocator(page_size=4, cp_size=cp_size) + cache = _make_cache(allocator) + for i in range(rng.randint(1, 25)): + n_pages = rng.randint(1, 6) + page_owners = [rng.randrange(cp_size) for _ in range(n_pages)] + node = _make_node( + 1000 + i, + 10000 + i * 1000, + page_owners, + value=torch.arange(1, n_pages * 4 + 1, dtype=torch.int64), + priority=rng.randint(0, 4), + ) + _attach_child(cache, cache.root_node, node) + cache.evictable_leaves.add(node) + if rng.random() < 0.4: + deficit = [0] * cp_size + deficit[rng.randrange(cp_size)] = rng.randint(1, 40) + else: + deficit = [rng.randint(0, 25) for _ in range(cp_size)] + if sum(1 for v in deficit if v > 0) >= 2: + multi_owner_seen += 1 + ref = _reference_greedy_plan(cache, _plan_for_deficit(deficit)) + got = _planner_result_tuple( + cache._plan_cp_load_back_owner_lane_evictions(_plan_for_deficit(deficit)) + ) + self.assertEqual(got, ref, msg=f"seed={seed} deficit={deficit}") + self.assertGreater(multi_owner_seen, 50) # the multi-owner path was actually exercised + + def test_heap_planner_matches_greedy_with_leaf_up_unlock(self): + # Leaf-up + ancestor-unlock: the parent (owner-0 pages) is evictable only after its + # child (owner-1, zero direct contribution to an owner-0 deficit) is evicted, so the + # child must be picked first to UNLOCK the parent. The heap's parent-push + the + # unlock score must reproduce the greedy's child-then-parent order. + cp_size = 4 + allocator = _make_allocator(page_size=4, cp_size=cp_size) + cache = _make_cache(allocator) + parent = _make_node( + 60, 600, [0, 0], value=torch.arange(1, 9, dtype=torch.int64), priority=0 + ) + child = _make_node( + 61, 700, [1], value=torch.arange(20, 24, dtype=torch.int64), priority=0 + ) + _attach_child(cache, cache.root_node, parent) + _attach_child(cache, parent, child) + cache.evictable_leaves.add(child) # parent has a child -> not yet evictable + + deficit = [2, 0, 0, 0] # only the parent has owner-0 pages + ref = _reference_greedy_plan(cache, _plan_for_deficit(deficit)) + got = cache._plan_cp_load_back_owner_lane_evictions(_plan_for_deficit(deficit)) + self.assertEqual(_planner_result_tuple(got), ref) + self.assertEqual(tuple(n.id for n in got.victims), (61, 60)) # unlock child, then parent + self.assertEqual(got.remaining_deficit_by_owner, (0, 0, 0, 0)) + def test_load_back_plan_fails_closed_without_cp_metadata(self): allocator = _make_allocator() cache = _make_cache(allocator)