From 2e15cb257aa7b33986ad7a5d63a6506783c51f92 Mon Sep 17 00:00:00 2001 From: leavelet Date: Wed, 24 Jun 2026 01:25:18 +0000 Subject: [PATCH] bench: CP owner-lane eviction planner micro-benchmark (validates the de-sync fix) Faithful replica of _plan_cp_load_back_owner_lane_evictions' inner loop with the real _cp_load_back_node_owner_page_counts bodies, over CUDA value tensors at the b300 hang's scale (candidates~2000, deficit 7877, ~36 owner pages/node). Measured on a V100: A. ORIGINAL (.item() x cp, no memo) 49.307s victims=219 B. FIXED (bincount + per-plan memo) 0.469s victims=219 (105x) C. REDESIGN (counts-once + sorted pass) 0.146s victims=219 (338x) selection identical: A==B, A==C Confirms the committed memo+bincount fix (4e9b4d05c0) is ~105x and byte-identical; C is the idealized single-owner ceiling (the real heap redesign must also handle leaf-up eligibility + ancestor-unlock + multi-owner). Co-Authored-By: Claude Opus 4.8 (1M context) --- .../hicache/bench_cp_owner_lane_planner.py | 143 ++++++++++++++++++ 1 file changed, 143 insertions(+) create mode 100644 benchmark/hicache/bench_cp_owner_lane_planner.py diff --git a/benchmark/hicache/bench_cp_owner_lane_planner.py b/benchmark/hicache/bench_cp_owner_lane_planner.py new file mode 100644 index 000000000..5e5acb875 --- /dev/null +++ b/benchmark/hicache/bench_cp_owner_lane_planner.py @@ -0,0 +1,143 @@ +"""Micro-benchmark for the CP owner-lane eviction planner hot loop. + +Faithfully replicates _plan_cp_load_back_owner_lane_evictions' inner loop with the REAL +_cp_load_back_node_owner_page_counts bodies (pooled shared-L2 path), over real CUDA value +tensors, at the b300 hang's scale (candidates~2000, deficit[owner]~7877, ~36 owner pages/node). + +Variants: + A. ORIGINAL : cp_size per-owner .item() syncs, recomputed every iteration (no memo) -> O(V*C*cp_size) syncs + B. FIXED : one bincount().tolist() sync, memoized per plan -> O(C) syncs + O(V*C) python + C. HEAP-REDES: compute counts once, then a single sorted pass (single-owner deficit) -> O(C) syncs + O(C log C) python + +Run on a GPU box: python bench_owner_lane_planner.py +""" +import time + +import torch + +DEV = "cuda" +PAGE = 64 +CP = 4 +N_CAND = 2000 # candidate evictable leaves (b300: ~2034) +N_PAGES_PER_NODE = 144 # -> ~36 owner pages each (b300: ~36) +DEFICIT_OWNER = 1 +DEFICIT = 7877 # b300 deficit_by_owner = [0, 7877, 0, 0] + +torch.manual_seed(0) + +# Build N_CAND fake nodes, each a CUDA int64 value tensor of distinct page-aligned token locs. +# first_locs = value[::PAGE] -> consecutive logical pages -> owners cycle 0..cp-1 (even ~36/owner). +NODES = [] +for i in range(N_CAND): + base = (i * N_PAGES_PER_NODE + 1) * PAGE + locs = torch.arange(base, base + N_PAGES_PER_NODE * PAGE, dtype=torch.int64, device=DEV) + NODES.append((i, locs)) + + +def counts_original(value, cp_size=CP): + # exact ORIGINAL _cp_load_back_node_owner_page_counts pooled body (cp_size .item() syncs) + first_locs = value[::PAGE] + logical_pages = torch.div(first_locs, PAGE, rounding_mode="floor") + owners = torch.remainder(logical_pages - 1, cp_size) + return tuple(int((owners == o).sum().item()) for o in range(cp_size)) + + +def counts_fixed(value, cp_size=CP): + # exact FIXED body (one bincount().tolist() sync) + first_locs = value[::PAGE] + logical_pages = torch.div(first_locs, PAGE, rounding_mode="floor") + owners = torch.remainder(logical_pages - 1, cp_size) + return tuple(torch.bincount(owners, minlength=cp_size).tolist()) + + +def plan_greedy(counts_fn, use_memo): + """The current planner: while deficit>0, rescan all candidates, pick best contribution.""" + deficits = [0] * CP + deficits[DEFICIT_OWNER] = DEFICIT + planned_evicted = set() + planned_freed = [0] * CP + memo = {} if use_memo else None + victims = 0 + while any(v > 0 for v in deficits): + best_score = None + best = None + best_counts = None + for nid, value in NODES: + if nid in planned_evicted: + continue + if use_memo and nid in memo: + counts = memo[nid] + else: + counts = counts_fn(value) + if use_memo: + memo[nid] = counts + contribution = sum(min(c, d) for c, d in zip(counts, deficits)) + if contribution <= 0: + continue + score = (-contribution, nid) + if best_score is None or score < best_score: + best_score, best, best_counts = score, nid, counts + if best is None: + break + planned_evicted.add(best) + for o, c in enumerate(best_counts): + planned_freed[o] += c + deficits[o] = max(0, deficits[o] - c) + victims += 1 + return victims, tuple(planned_freed) + + +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 + planned_freed = [0] * CP + victims = 0 + for counts, nid in scored: + if all(v <= 0 for v in rem): + break + for o, c in enumerate(counts): + planned_freed[o] += c + rem[o] = max(0, rem[o] - c) + victims += 1 + return victims, tuple(planned_freed) + + +def timed(label, fn): + torch.cuda.synchronize() + t0 = time.perf_counter() + v, freed = fn() + torch.cuda.synchronize() + dt = time.perf_counter() - t0 + print(f"{label:42s} {dt:9.3f}s victims={v:4d} planned_freed={freed}") + return v, freed + + +def main(): + print(f"device={torch.cuda.get_device_name(0)} cp={CP} page={PAGE} " + f"candidates={N_CAND} pages/node={N_PAGES_PER_NODE} deficit={DEFICIT}\n") + # warmup (kernels/caches) + plan_greedy(counts_fixed, True) + torch.cuda.synchronize() + + 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)) + + print("\nselection identical (victims, planned_freed):") + print(f" A==B: {(vO, fO) == (vF, fF)} A==C: {(vO, fO) == (vH, fH)}") + + +if __name__ == "__main__": + main()