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

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,