sglang/benchmark/hicache/bench_cp_owner_lane_planner.py

"""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()