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) <noreply@anthropic.com>

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)