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>

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.
+    """The SHIPPED heap redesign: lazy-re-evaluation min-heap (a popped entry is used only if
-    Byte-identical victim SET to greedy for a single-owner deficit (highest owner-count first,
+    its score is still current, else re-pushed).  Provably equivalent to the greedy argmin
-    id tiebreak) -- the contribution is min(counts[owner], remaining) == counts[owner] until the
+    rescan; here flat (no tree) so it isolates the algorithmic cost vs B."""
-    final pick."""
+    import heapq
-    owner = DEFICIT_OWNER
+
-    scored = []
+    deficits = [0] * CP
-    for nid, value in NODES:
+    deficits[DEFICIT_OWNER] = DEFICIT
-        counts = counts_fn(value)
+    planned_evicted = set()
-        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
     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:
+    while any(v > 0 for v in deficits) and heap:
-        if all(v <= 0 for v in rem):
+        score, _s, nid, counts = heapq.heappop(heap)
-            break
+        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)
+            deficits[o] = max(0, deficits[o] - c)
-        victims += 1
     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)}")

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(getattr(self, "evictable_leaves", set()))
-        initial_candidate_count = len(candidate_nodes)
+        # Per-plan memo {node.id: owner_counts}: each node's owner counts are invariant for
-        iteration = 0
+        # the whole plan, so compute the (device-sync) histogram once.
-        # 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).
         owner_counts_memo: dict = {}
 
-        while any(v > 0 for v in deficits):
+        # Heap-based leaf-up eviction, PROVABLY EQUIVALENT to the previous per-iteration
-            iteration += 1
+        # greedy argmin rescan, via LAZY RE-EVALUATION: a popped entry is consumed only if
-            best_node = None
+        # its score still matches the node's CURRENT score (deficits unchanged since it was
-            best_counts = None
+        # pushed); otherwise it is re-pushed with the fresh score and competes again.  Stale
-            best_score = None
+        # scores are always OPTIMISTIC -- contribution = sum(min(counts[o], deficits[o])) and
-            scanned_candidates = 0
+        # the ancestor-unlock contribution can only shrink as deficits shrink, so the score
-            for node in list(candidate_nodes):
+        # only worsens -- hence the first up-to-date popped entry is the global argmin the
-                scanned_candidates += 1
+        # full rescan would have picked.  Cost O((candidates + repushes) log n) with repushes
-                if (scanned_candidates & 1023) == 0:
+        # bounded by deficit crossings (~none until a lane's tail) vs O(victims * candidates).
-                    now = time.perf_counter()
+        # Determinism/rank-uniformity is preserved: selection is decided entirely by the
-                    if now - last_progress_time >= 5.0:
+        # total-ordered score (node.id is the final tiebreak); the heap insertion `seq` only
-                        logger.warning(
+        # orders structurally-equal tuples (and keeps heapq from comparing TreeNode objects)
-                            "[HiCache-load] slow CP owner-lane eviction planning scan: "
+        # and never changes which node is chosen.
-                            "iteration=%d scanned=%d candidates=%d victims=%d "
+        heap: list = []
-                            "deficits=%s planned_freed=%s elapsed_ms=%.3f",
+        seq = 0
-                            iteration,
+
-                            scanned_candidates,
+        def _scored(node):
-                            len(candidate_nodes),
+            counts = self._cp_load_back_node_owner_page_counts(
-                            len(victims),
+                node, cp_size, memo=owner_counts_memo
-                            deficits,
+            )
-                            planned_freed,
+            contribution = sum(
-                            (now - plan_start_time) * 1000.0,
+                min(int(count), int(deficit))
-                        )
+                for count, deficit in zip(counts, deficits)
-                        last_progress_time = now
+            )
-                if node in planned_evicted_nodes:
+            unlock_contribution = 0
-                    continue
+            if contribution <= 0:
-                if not self._cp_device_node_is_load_back_victim_after_plan(
+                unlock_contribution = self._cp_load_back_ancestor_unlock_contribution(
-                    node, planned_evicted_nodes
+                    node, deficits, planned_evicted_nodes, cp_size, memo=owner_counts_memo
-                ):
-                    continue
-                counts = self._cp_load_back_node_owner_page_counts(
-                    node, cp_size, memo=owner_counts_memo
                 )
-                contribution = sum(
+            if contribution <= 0 and unlock_contribution <= 0:
-                    min(int(count), int(deficit))
+                return None
-                    for count, deficit in zip(counts, deficits)
+            # get_priority returns the CpReplicatedSLRUStrategy tuple (is_protected,
-                )
+            # last_access_time[logical], node.id) -- rank-replicated; node.id keeps the full
-                unlock_contribution = 0
+            # score a strict total order even if the strategy is ever swapped.
-                if contribution <= 0:
+            score = (
-                    unlock_contribution = (
+                -int(contribution),
-                        self._cp_load_back_ancestor_unlock_contribution(
+                -int(unlock_contribution),
-                            node,
+                self.eviction_strategy.get_priority(node),
-                            deficits,
+                int(getattr(node, "id", 0) or 0),
-                            planned_evicted_nodes,
+            )
-                            cp_size,
+            return score, counts
-                            memo=owner_counts_memo,
+
-                        )
+        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:
+                    last_progress_time = now
-                    continue
+            if node in planned_evicted_nodes:
-                # get_priority returns the CpReplicatedSLRUStrategy tuple
+                continue  # stale entry: node already evicted in this plan
-                # (is_protected, last_access_time[logical], node.id) under CP — a
+            fresh = _scored(node)
-                # rank-replicated value; the trailing node.id keeps the full score a
+            if fresh is None:
-                # strict total order even if the strategy is ever swapped.
+                continue  # no longer contributes or unlocks (deficits shrank) -> drop
-                score = (
+            if fresh[0] != score:
-                    -int(contribution),
+                # deficits changed since this entry was pushed -> optimistic stale score;
-                    -int(unlock_contribution),
+                # re-push at the fresh (>= old) score so it competes again at its true cost.
-                    self.eviction_strategy.get_priority(node),
+                heapq.heappush(heap, (fresh[0], seq, node, fresh[1]))
-                    int(getattr(node, "id", 0) or 0),
+                seq += 1
-                )
+                continue
-                if best_score is None or score < best_score:
+            # up-to-date heap minimum == the current global argmin -> evict it.
-                    best_score = score
+            counts = fresh[1]
-                    best_node = node
+            victims.append(node)
-                    best_counts = counts
+            planned_evicted_nodes.add(node)
-
+            for owner, count in enumerate(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):
                 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 "
+                "elapsed_ms=%.3f initial_candidates=%d remaining_heap=%d "
-                "iterations=%d victims=%d original_deficit_by_owner=%s "
+                "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),
+                len(heap),
-                iteration,
+                pops,
                 len(victims),
                 plan.deficit_by_owner,
                 planned_freed,

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)