chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

CP HiCache: redesign owner-lane eviction planner as a lazy-re-evaluation heap (O(n log n))

Browse Source 
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>
This commit is contained in:
孙远航
2026-06-24 01:35:37 +00:00
parent 2e15cb257a
commit fdc23f7d61
3 changed files with 297 additions and 102 deletions
Download Patch File Download Diff File Expand all files Collapse all files

65
benchmark/hicache/bench_cp_owner_lane_planner.py
View File

@@ -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)}")

183
python/sglang/srt/mem_cache/hiradix_cache.py
View File

@@ -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,

151
test/registered/unit/mem_cache/test_cp_hicache_load_back_owner_lanes.py
View File

@@ -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)