From 401de0f8ced2d7d6a9c341dcdd3f95afc92fd7f2 Mon Sep 17 00:00:00 2001 From: laoyao0822 Date: Tue, 2 Jun 2026 23:35:12 +0800 Subject: [PATCH] Reduce CP HiCache L2 allocator scan cost Host HiCache reservations were paying token-level free-slot scans when trying to preserve page contiguity. The allocator now keeps a lazy page-extent index so availability checks and contiguous-preferred allocations avoid materializing the full 220GB-equivalent free-slot metadata path. The companion benchmark models steady-state L2 churn near full occupancy, including burn-in and historical node-size effects, so LPF/RDMA descriptor quality can be separated from ETE noise. Constraint: CP HiCache host allocations are page-shaped, but existing callers may still read free_slots directly. Rejected: Sort and scan free_slots on each alloc_contiguous_preferred call | measured ms-level CPU overhead on 220GB-equivalent metadata. Rejected: Remove free_slots compatibility | storage/tests still rely on the public tensor surface. Confidence: medium Scope-risk: moderate Directive: Do not reintroduce per-allocation full free_slots scans on HostKVCache; preserve page-extent metadata or benchmark before changing allocator shape. Tested: Local py_compile for memory_pool_host.py, allocator benchmark, and related tests. Tested: Local test_cp_hicache_allocator_bench.py 10 passed. Tested: Remote g0034 test_hicache_controller_cp.py 67 passed; test_cp_hicache_allocator_bench.py 10 passed. Tested: Remote 220GB-equivalent host_churn benchmark: contiguous path reduced from ms-level to ~30-292us p50 depending on fragmentation. Not-tested: Full CUDA ETE run after allocator change. Not-tested: Production long-run fragmentation behavior under live traffic. --- .../bench_cp_hicache_allocator_overhead.py | 395 +++++++++++++++++- ..._prefill_cp_page_aligned_cache_contract.md | 151 +++++++ .../sglang/srt/mem_cache/memory_pool_host.py | 274 ++++++++++++ .../test_cp_hicache_allocator_bench.py | 84 ++++ .../managers/test_hicache_controller_cp.py | 23 + 5 files changed, 925 insertions(+), 2 deletions(-) diff --git a/benchmark/hicache/bench_cp_hicache_allocator_overhead.py b/benchmark/hicache/bench_cp_hicache_allocator_overhead.py index 3d4e1a7b1..e1b944eef 100644 --- a/benchmark/hicache/bench_cp_hicache_allocator_overhead.py +++ b/benchmark/hicache/bench_cp_hicache_allocator_overhead.py @@ -14,6 +14,13 @@ Examples: --bench host --host-sizes-gb 220 --request-pages 1,8,64,512 \ --patterns contiguous_fifo,fragmented_prefix_later_run,random_fragmented + # Steady-state L2 host churn model near full HiCache occupancy. + PYTHONPATH=. python benchmark/hicache/bench_cp_hicache_allocator_overhead.py \ + --bench host_churn --host-sizes-gb 220 --request-pages 16,64,512 \ + --host-churn-occupancies 0.90,0.97,0.99 \ + --host-churn-evict-pages 64,512,2048 \ + --host-churn-eviction-patterns oldest,random + # L1 allocator path on CUDA, stubbing sgl_kernel import if needed. PYTHONPATH=python:. python benchmark/hicache/bench_cp_hicache_allocator_overhead.py \ --bench l1 --device cuda --stub-sgl-kernel --physical-pages 8192,32768 \ @@ -25,6 +32,7 @@ import argparse import json import math import os +import random import statistics import sys import threading @@ -55,6 +63,36 @@ class BenchResult: contiguous_ratio: float +@dataclass(frozen=True) +class HostChurnBenchResult: + bench: str + impl: str + pattern: str + device: str + total_pages: int + request_pages: int + page_size: int + repeat: int + target_occupancy: float + evict_pages: int + prefill_node_pages: int + burnin: int + mean_us: float + p50_us: float + p95_us: float + p99_us: float + min_us: float + max_us: float + contiguous_ratio: float + page_first_descriptors_per_op: int + lpf_descriptors_mean: float + lpf_descriptor_ratio_mean: float + run_count_p50: float + run_count_p95: float + max_run_pages_mean: float + max_run_pages_p50: float + + def _parse_int_list(value: str | Iterable[int]) -> list[int]: if isinstance(value, str): return [int(item.strip()) for item in value.split(",") if item.strip()] @@ -244,6 +282,10 @@ class StandaloneHostAllocator: self.free_slots = self.free_slots[keep_mask] return select_index + def free(self, indices: torch.Tensor) -> int: + self.free_slots = torch.cat([self.free_slots, indices.cpu()]) + return int(indices.numel()) + def _compute_owner_lane_free_room_deficits( *, @@ -544,6 +586,43 @@ def _is_page_contiguous_selection(selected: Optional[torch.Tensor], page_size: i return bool(torch.all(pages[1:] - pages[:-1] == 1).item()) +def _page_run_lengths_from_token_slots( + selected: Optional[torch.Tensor], page_size: int +) -> list[int]: + """Return consecutive physical-page run lengths for one host selection. + + This is the layout-independent descriptor proxy used by the L2 benchmark: + current ``page_first_direct`` needs one fixed-layer copy descriptor per page, + while ``layer_page_first`` can collapse each consecutive page run to one + descriptor per KV tensor. + """ + + if selected is None or selected.numel() == 0: + return [] + if selected.numel() % page_size != 0: + raise ValueError( + f"selected token slots must be page-shaped, got {selected.numel()=} " + f"{page_size=}" + ) + pages = (selected.view(-1, page_size)[:, 0] // page_size).tolist() + if not pages: + return [] + + run_lengths: list[int] = [] + current_len = 1 + prev_page = int(pages[0]) + for page in pages[1:]: + page = int(page) + if page == prev_page + 1: + current_len += 1 + else: + run_lengths.append(current_len) + current_len = 1 + prev_page = page + run_lengths.append(current_len) + return run_lengths + + def _make_host_allocator(impl: str, *, page_size: int, free_slots: torch.Tensor): if impl == "standalone": return StandaloneHostAllocator(page_size=page_size, free_slots=free_slots) @@ -614,6 +693,63 @@ def _summarize( ) +def _summarize_host_churn( + *, + impl: str, + method: str, + eviction_pattern: str, + total_pages: int, + request_pages: int, + page_size: int, + target_occupancy: float, + evict_pages: int, + prefill_node_pages: int, + burnin: int, + samples_us: list[float], + run_counts: list[int], + max_run_lengths: list[int], + contiguous_hits: int, +) -> HostChurnBenchResult: + repeat = len(samples_us) + lpf_descriptors_mean = ( + float(statistics.mean(run_counts)) if run_counts else 0.0 + ) + return HostChurnBenchResult( + bench="host_churn", + impl=f"{impl}:{method}", + pattern=f"occ={target_occupancy:.2f}:{eviction_pattern}", + device="cpu", + total_pages=int(total_pages), + request_pages=int(request_pages), + page_size=int(page_size), + repeat=repeat, + target_occupancy=float(target_occupancy), + evict_pages=int(evict_pages), + prefill_node_pages=int(prefill_node_pages), + burnin=int(burnin), + mean_us=float(statistics.mean(samples_us)) if samples_us else 0.0, + p50_us=float(_percentile(samples_us, 50)), + p95_us=float(_percentile(samples_us, 95)), + p99_us=float(_percentile(samples_us, 99)), + min_us=float(min(samples_us)) if samples_us else 0.0, + max_us=float(max(samples_us)) if samples_us else 0.0, + contiguous_ratio=float(contiguous_hits / repeat) if repeat else 0.0, + page_first_descriptors_per_op=int(request_pages), + lpf_descriptors_mean=lpf_descriptors_mean, + lpf_descriptor_ratio_mean=( + lpf_descriptors_mean / float(request_pages) if request_pages else 0.0 + ), + run_count_p50=float(_percentile([float(x) for x in run_counts], 50)), + run_count_p95=float(_percentile([float(x) for x in run_counts], 95)), + max_run_pages_mean=( + float(statistics.mean(max_run_lengths)) if max_run_lengths else 0.0 + ), + max_run_pages_p50=float( + _percentile([float(x) for x in max_run_lengths], 50) + ), + ) + + def _bench_host_case( *, impl: str, @@ -660,6 +796,166 @@ def _bench_host_case( ) +def _evict_host_churn_nodes( + *, + allocator, + active_nodes: list[torch.Tensor], + target_pages_to_free: int, + page_size: int, + eviction_pattern: str, + rng: random.Random, +) -> int: + freed_pages = 0 + while active_nodes and freed_pages < target_pages_to_free: + if eviction_pattern == "oldest": + node_index = 0 + elif eviction_pattern == "youngest": + node_index = len(active_nodes) - 1 + elif eviction_pattern == "random": + node_index = rng.randrange(len(active_nodes)) + else: + raise ValueError(f"unsupported host churn eviction pattern: {eviction_pattern}") + node = active_nodes.pop(node_index) + allocator.free(node) + freed_pages += int(node.numel()) // page_size + return freed_pages + + +def _bench_host_churn_case( + *, + impl: str, + method: str, + total_pages: int, + request_pages: int, + page_size: int, + target_occupancy: float, + evict_pages: int, + eviction_pattern: str, + repeat: int, + warmup: int, + seed: int, + prefill_node_pages: Optional[int] = None, + burnin: int = 0, +) -> HostChurnBenchResult: + if not 0 < target_occupancy < 1: + raise ValueError( + f"target_occupancy must be in (0, 1), got {target_occupancy}" + ) + if request_pages <= 0: + raise ValueError(f"request_pages must be positive, got {request_pages}") + if evict_pages <= 0: + raise ValueError(f"evict_pages must be positive, got {evict_pages}") + if burnin < 0: + raise ValueError(f"burnin must be non-negative, got {burnin}") + if request_pages > total_pages: + raise ValueError( + f"request_pages must be <= total_pages, got {request_pages=} {total_pages=}" + ) + if prefill_node_pages is None: + prefill_node_pages = request_pages + if prefill_node_pages <= 0: + raise ValueError( + f"prefill_node_pages must be positive, got {prefill_node_pages}" + ) + if prefill_node_pages > total_pages: + raise ValueError( + "prefill_node_pages must be <= total_pages, got " + f"{prefill_node_pages=} {total_pages=}" + ) + + base_free_slots = _make_host_free_slots( + total_pages=total_pages, + request_pages=request_pages, + page_size=page_size, + pattern="contiguous_fifo", + seed=seed, + ) + allocator = _make_host_allocator(impl, page_size=page_size, free_slots=base_free_slots) + rng = random.Random(seed) + need_size = request_pages * page_size + prefill_need_size = prefill_node_pages * page_size + + # Fill with configurable node sizes so the benchmark can model fragmented + # steady-state HiCache: many old small nodes can be evicted to satisfy one + # larger new request, which is the path LPF allocation policy cares about. + target_used_pages = min( + total_pages - request_pages, + int(math.floor(float(total_pages) * float(target_occupancy))), + ) + target_used_pages = (target_used_pages // prefill_node_pages) * prefill_node_pages + active_nodes: list[torch.Tensor] = [] + used_pages = 0 + while used_pages + prefill_node_pages <= target_used_pages: + selected = allocator.alloc(prefill_need_size) + if selected is None: + break + active_nodes.append(selected) + used_pages += prefill_node_pages + + samples_us: list[float] = [] + run_counts: list[int] = [] + max_run_lengths: list[int] = [] + contiguous_hits = 0 + fn = allocator.alloc if method == "fifo" else allocator.alloc_contiguous_preferred + min_evict_pages = max(evict_pages, request_pages) + + first_sample_iteration = int(burnin) + int(warmup) + for iteration in range(first_sample_iteration + repeat): + _evict_host_churn_nodes( + allocator=allocator, + active_nodes=active_nodes, + target_pages_to_free=min_evict_pages, + page_size=page_size, + eviction_pattern=eviction_pattern, + rng=rng, + ) + while allocator.available_size() < need_size and active_nodes: + _evict_host_churn_nodes( + allocator=allocator, + active_nodes=active_nodes, + target_pages_to_free=request_pages, + page_size=page_size, + eviction_pattern=eviction_pattern, + rng=rng, + ) + + start_ns = time.perf_counter_ns() + selected = fn(need_size) + elapsed_us = (time.perf_counter_ns() - start_ns) / 1000.0 + if selected is None: + raise RuntimeError( + "host churn allocation failed after eviction: " + f"{total_pages=} {request_pages=} {target_occupancy=} " + f"{evict_pages=} {eviction_pattern=}" + ) + active_nodes.append(selected) + + if iteration >= first_sample_iteration: + samples_us.append(elapsed_us) + run_lengths = _page_run_lengths_from_token_slots(selected, page_size) + run_count = len(run_lengths) + run_counts.append(run_count) + max_run_lengths.append(max(run_lengths) if run_lengths else 0) + contiguous_hits += int(run_count <= 1) + + return _summarize_host_churn( + impl=impl, + method=method, + eviction_pattern=eviction_pattern, + total_pages=total_pages, + request_pages=request_pages, + page_size=page_size, + target_occupancy=target_occupancy, + evict_pages=evict_pages, + prefill_node_pages=prefill_node_pages, + burnin=burnin, + samples_us=samples_us, + run_counts=run_counts, + max_run_lengths=max_run_lengths, + contiguous_hits=contiguous_hits, + ) + + def _zigzag_owners(num_pages: int, cp_size: int) -> list[int]: segment_num = cp_size * 2 base = num_pages // segment_num @@ -915,7 +1211,20 @@ def _bench_l1_case( ) -def _format_result(result: BenchResult) -> str: +def _format_result(result: BenchResult | HostChurnBenchResult) -> str: + if isinstance(result, HostChurnBenchResult): + return ( + f"{result.bench:10s} impl={result.impl:20s} pattern={result.pattern:34s} " + f"dev={result.device:4s} pages={result.total_pages:7d} req={result.request_pages:5d} " + f"evict={result.evict_pages:5d} prefill_node={result.prefill_node_pages:5d} " + f"burnin={result.burnin:4d} p50={result.p50_us:9.2f}us " + f"p95={result.p95_us:9.2f}us p99={result.p99_us:9.2f}us " + f"mean={result.mean_us:9.2f}us contig={result.contiguous_ratio:.2f} " + f"pf_desc={result.page_first_descriptors_per_op:d} " + f"lpf_desc_mean={result.lpf_descriptors_mean:.2f} " + f"lpf_ratio={result.lpf_descriptor_ratio_mean:.3f} " + f"run_p50={result.run_count_p50:.1f} max_run_mean={result.max_run_pages_mean:.1f}" + ) return ( f"{result.bench:4s} impl={result.impl:20s} pattern={result.pattern:34s} " f"dev={result.device:4s} pages={result.total_pages:7d} req={result.request_pages:5d} " @@ -965,6 +1274,64 @@ def _run_host(args) -> list[BenchResult]: return results +def _run_host_churn(args) -> list[HostChurnBenchResult]: + host_pages = _parse_int_list(args.host_pages) if args.host_pages else [] + for size_gb in _parse_float_list(args.host_sizes_gb): + host_pages.append( + _host_pages_from_gb( + size_gb, bytes_per_token=args.bytes_per_token, page_size=args.page_size + ) + ) + if not host_pages: + host_pages = [8192, 16384, 32768] + host_pages = sorted(set(page for page in host_pages if page > 0)) + request_pages_list = _parse_int_list(args.request_pages) + host_impls = [item.strip() for item in args.host_impl.split(",") if item.strip()] + methods = [item.strip() for item in args.host_methods.split(",") if item.strip()] + occupancies = _parse_float_list(args.host_churn_occupancies) + evict_pages_list = _parse_int_list(args.host_churn_evict_pages) + prefill_node_pages_list = ( + _parse_int_list(args.host_churn_prefill_node_pages) + if args.host_churn_prefill_node_pages + else [0] + ) + eviction_patterns = [ + item.strip() for item in args.host_churn_eviction_patterns.split(",") if item.strip() + ] + + results: list[HostChurnBenchResult] = [] + for total_pages in host_pages: + for request_pages in request_pages_list: + if request_pages > total_pages: + continue + for target_occupancy in occupancies: + for evict_pages in evict_pages_list: + for prefill_node_pages in prefill_node_pages_list: + for eviction_pattern in eviction_patterns: + for impl in host_impls: + for method in methods: + results.append( + _bench_host_churn_case( + impl=impl, + method=method, + total_pages=total_pages, + request_pages=request_pages, + page_size=args.page_size, + target_occupancy=target_occupancy, + evict_pages=evict_pages, + eviction_pattern=eviction_pattern, + repeat=args.repeat, + warmup=args.warmup, + seed=args.seed, + prefill_node_pages=( + prefill_node_pages or None + ), + burnin=args.host_churn_burnin, + ) + ) + return results + + def _run_l1(args) -> list[BenchResult]: if args.stub_sgl_kernel: _install_sgl_kernel_stubs() @@ -1015,7 +1382,9 @@ def _run_l1(args) -> list[BenchResult]: def _build_parser() -> argparse.ArgumentParser: parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--bench", default="host,l1", help="comma list: host,l1") + parser.add_argument( + "--bench", default="host,l1", help="comma list: host,host_churn,l1" + ) parser.add_argument("--page-size", type=int, default=64) parser.add_argument("--repeat", type=int, default=20) parser.add_argument("--warmup", type=int, default=5) @@ -1036,6 +1405,26 @@ def _build_parser() -> argparse.ArgumentParser: ) parser.add_argument("--host-impl", default="standalone") parser.add_argument("--host-methods", default="fifo,contiguous") + parser.add_argument("--host-churn-occupancies", default="0.90,0.97,0.99") + parser.add_argument("--host-churn-evict-pages", default="64,512,2048") + parser.add_argument( + "--host-churn-prefill-node-pages", + default="", + help=( + "comma list of node sizes used to prefill steady-state host cache; " + "default uses each request_pages value" + ), + ) + parser.add_argument( + "--host-churn-burnin", + type=int, + default=0, + help=( + "unmeasured steady-state evict+allocate iterations before warmup; " + "useful for exhausting the cold contiguous free tail" + ), + ) + parser.add_argument("--host-churn-eviction-patterns", default="oldest,random") parser.add_argument("--physical-pages", default="8192,32768") parser.add_argument("--cp-size", type=int, default=8) @@ -1071,6 +1460,8 @@ def main(argv: Optional[list[str]] = None) -> int: benches = {item.strip() for item in args.bench.split(",") if item.strip()} if "host" in benches: results.extend(_run_host(args)) + if "host_churn" in benches: + results.extend(_run_host_churn(args)) if "l1" in benches: results.extend(_run_l1(args)) diff --git a/docs/advanced_features/nsa_prefill_cp_page_aligned_cache_contract.md b/docs/advanced_features/nsa_prefill_cp_page_aligned_cache_contract.md index 908413b26..e03d11f15 100644 --- a/docs/advanced_features/nsa_prefill_cp_page_aligned_cache_contract.md +++ b/docs/advanced_features/nsa_prefill_cp_page_aligned_cache_contract.md @@ -5311,3 +5311,154 @@ C123 full-suite update: - The CPU stub now also gives `sgl_kernel.kvcacheio` a module-level `__getattr__`, so named imports resolve to inert functions without importing the native extension. + +### C124 — 2026-06-02 L2 host churn benchmark must model fragmented steady-state eviction + +Finding: + +- The earlier host allocator benchmark measured one cold allocation against a + synthetic `free_slots` layout. That is not enough for production HiCache: L2 + host cache can be ~220 GB, runs near full occupancy, and repeatedly frees old + nodes before reserving a new node. +- If prefill nodes and measured allocation requests have the same page size, a + random eviction workload can still produce unrealistically contiguous freed + chunks. This hides the allocator/pathology we care about for layer-page-first + (LPF) transfer planning. + +Correction: + +- Added `benchmark/hicache/bench_cp_hicache_allocator_overhead.py --bench + host_churn` for steady-state L2 metadata churn. +- The benchmark now records both CPU allocation latency and transfer-descriptor + quality proxies: + - `pf_desc`: page-first descriptor count, one descriptor per selected page. + - `lpf_desc_mean` / `lpf_ratio`: layer-page-first descriptor count after + coalescing consecutive physical page runs. + - `run_p50` / `max_run_mean`: selected physical-page run quality. +- Added `--host-churn-prefill-node-pages` so the filled resident set can use a + different historical node size than the new allocation request. This models + cache built from many small chunks, then later serving larger extensions. + +Scope: + +- This is CPU metadata-only. It does not allocate real 220 GB KV buffers and it + does not test CUDA kernel bandwidth. +- It is intended to decide whether L2 `HostKVCache.alloc_contiguous_preferred()` + / future L2 bucket allocator work is on the actual hot path, and whether LPF + layout can get enough physical-page coalescing to help H2D/D2H/RDMA. + +C124 validation update: + +- Local CPU unit coverage: `test_cp_hicache_allocator_bench.py` passes with the + fragmented prefill-node churn case. +- Remote `g0034` container validation: py_compile passed and the benchmark unit + file passed (`9 passed`). +- Remote 220GB-equivalent production `HostKVCache` churn sample + (`pages=34375`, `occ=0.97`, `request_pages=512`, `evict_pages=512`): + - `prefill_node_pages=1`: FIFO p50 ~11 us but LPF run quality is poor + (`lpf_ratio≈0.663`, `run_p50≈505`). Contiguous-preferred p50 ~700 us and + still cannot find a good run under this fragmentation. + - `prefill_node_pages=64`: FIFO p50 ~9 us and LPF run quality is much better + (`lpf_ratio≈0.012`, `run_p50≈8`). Contiguous-preferred p50 ~499 us. +- Interpretation: current production `alloc_contiguous_preferred()` can add + hundreds of microseconds on fragmented 220GB-equivalent metadata. The + benchmark now makes this measurable separately from ETE noise; it also shows + that node-size history strongly affects LPF transfer coalescing potential. + +### C125 — 2026-06-02 Host churn benchmark needs explicit burn-in to avoid cold free-tail bias + +Finding: + +- A high-occupancy prefill still leaves an initial contiguous free tail. With + small `warmup` and small requests, measured allocations can consume this cold + tail before they ever allocate from evicted fragmented nodes. +- That can overstate LPF run quality and understate the allocation/search cost we + expect after the service has churned for a while. + +Correction plan: + +- Add a separate `--host-churn-burnin` iteration count. Burn-in iterations run + the same evict+allocate cycle but are not measured and are independent of + benchmark warmup. +- Use burn-in to exhaust the initial free tail before collecting steady-state + latency and run-quality samples. + +C125 validation update: + +- Local RED/GREEN: `test_host_churn_burnin_exposes_fragmented_evicted_nodes_after_cold_tail` + first failed on missing `burnin`, then passed after adding the burn-in path. +- Local full benchmark unit file: `10 passed` plus py_compile. +- Remote `g0034` container: py_compile passed and + `test_cp_hicache_allocator_bench.py` passed (`10 passed`). +- Remote 220GB-equivalent production `HostKVCache` sample with `burnin=20`, + `occ=0.97`, `evict_pages=512`: + - `request_pages=64`, `prefill_node_pages=1`: FIFO p50 ~11 us but LPF quality + is worst case (`lpf_ratio=1.000`, all single-page runs). The contiguous + search can cost ~2.6 ms p50. + - `request_pages=512`, `prefill_node_pages=1`: FIFO p50 ~10 us, LPF quality + remains worst case (`lpf_ratio=1.000`). Contiguous search p50 ~576 us and + cannot improve run quality because no 512-page run exists. + - `prefill_node_pages=64`: LPF quality is much better (`lpf_ratio≈0.012–0.031`) + but contiguous search can still cost ~0.5–1.5 ms p50. +- Interpretation: this benchmark now exposes the relevant tradeoff: FIFO is cheap + but can create very high descriptor count for LPF/RDMA when historical nodes + are tiny; naive contiguous search can be milliseconds on 220GB-equivalent host + metadata. A future L2 allocator should avoid full free-list scans and preserve + larger free extents/buckets rather than searching linearly per allocation. + +### C126 — 2026-06-02 L2 HostKVCache allocator uses lazy page extents instead of full free-slot scans + +Finding: + +- `HostKVCache.alloc_contiguous_preferred()` previously scanned/materialized the + full token-level `free_slots` tensor to discover contiguous physical page runs. + On 220GB-equivalent host metadata this could cost hundreds of microseconds to + milliseconds per reservation. +- The full scan is the wrong shape for CP HiCache: allocations and releases are + page-shaped, and the fast path only needs page-run metadata plus a token-index + tensor for the chosen pages. + +Correction: + +- `HostKVCache` now maintains a lazy page-extent index: + - `free_slots` remains a compatibility property and is materialized lazily only + when external code reads it. + - `available_size()` reads an integer token count and does not materialize. + - `free()` converts page-shaped token indices into page runs and merges them + into sorted extents. + - `alloc_contiguous_preferred()` first checks the largest free extent for a + single-run allocation; if no run is large enough, it falls back to a batched + fragmented run allocation without scanning/sorting the whole `free_slots` + tensor. +- The fallback is still page-shaped and returns normal token indices. It avoids + silent corruption: overlapping/double-free page extents raise, and non-page + shaped frees disable the extent index with a warning before falling back to the + legacy tensor path. + +C126 validation update: + +- Remote RED/GREEN target: the new lazy extent-index unit first failed on missing + `_free_slots_dirty`, then passed after implementation. +- Remote `test_hicache_controller_cp.py`: `67 passed`. +- Remote `test_cp_hicache_allocator_bench.py`: `10 passed`. +- Remote 220GB-equivalent production `HostKVCache` churn sample with `burnin=20`, + `occ=0.97`, `evict_pages=512` after optimization: + - `request_pages=64`, `prefill_node_pages=64`: contiguous p50 ~31 us. + - `request_pages=512`, `prefill_node_pages=64`: contiguous p50 ~86 us. + - `request_pages=512`, `prefill_node_pages=1`: contiguous p50 ~292 us; this + is no longer a full free-list scan, but mostly fragmented selection plus + constructing a 32768-token host-index tensor. +- Compared with the previous C125 sample, the contiguous path drops from roughly + 0.5–2.6 ms p50 to roughly 30–292 us p50 on this benchmark shape. + +Remaining risk: + +- `alloc()` now uses the extent-backed fragmented allocation when the extent + index is active, so exact FIFO order is no longer the internal contract. This + should be acceptable for host KV slots because callers require unique free + page-shaped slots, not FIFO identity, but future code must not depend on + token-level FIFO order. +- Very fragmented tiny historical nodes still produce poor LPF coalescing. That + requires higher-level eviction/allocation policy to preserve larger free runs; + the allocator now avoids the worst CPU scan cost but cannot create physical + contiguity that the free set does not contain. diff --git a/python/sglang/srt/mem_cache/memory_pool_host.py b/python/sglang/srt/mem_cache/memory_pool_host.py index b201ed9c6..16af23e05 100644 --- a/python/sglang/srt/mem_cache/memory_pool_host.py +++ b/python/sglang/srt/mem_cache/memory_pool_host.py @@ -1,4 +1,6 @@ import abc +import bisect +import heapq import logging import threading from collections import defaultdict @@ -314,6 +316,249 @@ class HostKVCache(abc.ABC): """ raise NotImplementedError() + @property + def free_slots(self) -> torch.Tensor: + if getattr(self, "_free_slots_dirty", False): + self._free_slots = self._materialize_free_slots_from_extents() + self._free_slots_dirty = False + return self._free_slots + + @free_slots.setter + def free_slots(self, value: torch.Tensor) -> None: + self._free_slots = value.cpu().to(dtype=torch.int64).contiguous() + self._rebuild_free_extent_index_from_slots(self._free_slots) + + def _reset_free_extent_index(self, *, enabled: bool, token_count: int) -> None: + self._free_extent_index_enabled = bool(enabled) + self._free_token_count = int(token_count) + self._free_extents_by_start: dict[int, int] = {} + self._free_extent_starts: list[int] = [] + self._free_extent_heap: list[tuple[int, int]] = [] + self._free_slots_dirty = False + + def _rebuild_free_extent_index_from_slots(self, free_slots: torch.Tensor) -> None: + token_count = int(free_slots.numel()) + self._reset_free_extent_index(enabled=False, token_count=token_count) + page_size = int(self.page_size) + if token_count == 0: + self._free_extent_index_enabled = True + return + if token_count % page_size != 0: + logger.warning( + "[HiCache-L2-allocator] disabling extent index for non-page-shaped free_slots: tokens=%d page_size=%d", + token_count, + page_size, + ) + return + + page_slots = free_slots.view(-1, page_size) + offsets = torch.arange(page_size, dtype=page_slots.dtype) + if not bool(torch.all(page_slots == (page_slots[:, :1] + offsets)).item()): + logger.warning( + "[HiCache-L2-allocator] disabling extent index for non-contiguous page slots" + ) + return + + self._free_extent_index_enabled = True + pages = torch.div(page_slots[:, 0], page_size, rounding_mode="floor").tolist() + if not pages: + return + run_start = int(pages[0]) + run_len = 1 + prev = run_start + for page in pages[1:]: + page = int(page) + if page == prev + 1: + run_len += 1 + else: + self._insert_free_extent(run_start, run_len) + run_start = page + run_len = 1 + prev = page + self._insert_free_extent(run_start, run_len) + + def _materialize_free_slots_from_extents(self) -> torch.Tensor: + page_size = int(self.page_size) + chunks: list[torch.Tensor] = [] + offsets = torch.arange(page_size, dtype=torch.int64) + for start in self._free_extent_starts: + length = self._free_extents_by_start.get(start, 0) + if length <= 0: + continue + pages = torch.arange(start, start + length, dtype=torch.int64) + chunks.append((pages[:, None] * page_size + offsets[None, :]).reshape(-1)) + if not chunks: + return torch.empty((0,), dtype=torch.int64) + return torch.cat(chunks).contiguous() + + def _insert_free_extent(self, start_page: int, page_count: int) -> None: + if page_count <= 0: + return + start_page = int(start_page) + page_count = int(page_count) + end_page = start_page + page_count + + pos = bisect.bisect_left(self._free_extent_starts, start_page) + if pos > 0: + prev_start = self._free_extent_starts[pos - 1] + prev_len = self._free_extents_by_start[prev_start] + prev_end = prev_start + prev_len + if prev_end > start_page: + raise RuntimeError( + "[HiCache-L2-allocator] double free or overlapping host extent" + ) + if prev_end == start_page: + start_page = prev_start + page_count += prev_len + del self._free_extents_by_start[prev_start] + del self._free_extent_starts[pos - 1] + pos -= 1 + if pos < len(self._free_extent_starts): + next_start = self._free_extent_starts[pos] + next_len = self._free_extents_by_start[next_start] + if end_page > next_start: + raise RuntimeError( + "[HiCache-L2-allocator] double free or overlapping host extent" + ) + if end_page == next_start: + page_count += next_len + del self._free_extents_by_start[next_start] + del self._free_extent_starts[pos] + + self._free_extents_by_start[start_page] = page_count + bisect.insort(self._free_extent_starts, start_page) + heapq.heappush(self._free_extent_heap, (-page_count, start_page)) + + def _pop_largest_free_extent(self) -> Optional[tuple[int, int]]: + while self._free_extent_heap: + neg_len, start = heapq.heappop(self._free_extent_heap) + length = -int(neg_len) + current = self._free_extents_by_start.get(start) + if current == length: + del self._free_extents_by_start[start] + pos = bisect.bisect_left(self._free_extent_starts, start) + if ( + pos < len(self._free_extent_starts) + and self._free_extent_starts[pos] == start + ): + del self._free_extent_starts[pos] + return start, length + return None + + def _peek_largest_free_extent(self) -> Optional[tuple[int, int]]: + while self._free_extent_heap: + neg_len, start = self._free_extent_heap[0] + length = -int(neg_len) + if self._free_extents_by_start.get(start) == length: + return start, length + heapq.heappop(self._free_extent_heap) + return None + + def _page_ids_to_token_indices(self, page_ids: list[int]) -> torch.Tensor: + if not page_ids: + return torch.empty((0,), dtype=torch.int64) + page_size = int(self.page_size) + pages = torch.tensor(page_ids, dtype=torch.int64) + offsets = torch.arange(page_size, dtype=torch.int64) + return (pages[:, None] * page_size + offsets[None, :]).reshape(-1).contiguous() + + def _alloc_from_free_extents( + self, need_size: int, *, require_single_run: bool + ) -> Optional[torch.Tensor]: + page_size = int(self.page_size) + need_pages = int(need_size) // page_size + if need_size > self.available_size(): + return None + if need_pages == 0: + return torch.empty((0,), dtype=torch.int64) + + if require_single_run: + largest = self._peek_largest_free_extent() + if largest is None or largest[1] < need_pages: + return None + start, length = self._pop_largest_free_extent() + selected_pages: list[int] = [] + selected_pages.extend(range(start, start + need_pages)) + remaining = length - need_pages + if remaining > 0: + self._insert_free_extent(start + need_pages, remaining) + else: + selected_pages = self._alloc_fragmented_from_free_extents(need_pages) + + self._free_token_count -= need_size + self._free_slots_dirty = True + return self._page_ids_to_token_indices(selected_pages) + + def _alloc_fragmented_from_free_extents(self, need_pages: int) -> list[int]: + remaining_pages = int(need_pages) + selected_pages: list[int] = [] + consumed_count = 0 + replacement: Optional[tuple[int, int]] = None + + for start in self._free_extent_starts: + if remaining_pages <= 0: + break + length = self._free_extents_by_start[start] + take = min(length, remaining_pages) + selected_pages.extend(range(start, start + take)) + remaining_pages -= take + if take == length: + consumed_count += 1 + continue + replacement = (start + take, length - take) + break + + if remaining_pages > 0: + raise RuntimeError( + "[HiCache-L2-allocator] extent index underflow during fragmented allocation" + ) + + consumed_starts = self._free_extent_starts[:consumed_count] + for start in consumed_starts: + del self._free_extents_by_start[start] + del self._free_extent_starts[:consumed_count] + + if replacement is not None: + old_start = self._free_extent_starts[0] + del self._free_extents_by_start[old_start] + self._free_extent_starts[0] = replacement[0] + self._free_extents_by_start[replacement[0]] = replacement[1] + heapq.heappush(self._free_extent_heap, (-replacement[1], replacement[0])) + + return selected_pages + + def _page_runs_from_token_indices( + self, indices: torch.Tensor + ) -> Optional[list[tuple[int, int]]]: + indices = indices.cpu().to(dtype=torch.int64).contiguous() + token_count = int(indices.numel()) + page_size = int(self.page_size) + if token_count == 0: + return [] + if token_count % page_size != 0: + return None + page_slots = indices.view(-1, page_size) + offsets = torch.arange(page_size, dtype=page_slots.dtype) + if not bool(torch.all(page_slots == (page_slots[:, :1] + offsets)).item()): + return None + + pages = torch.div(page_slots[:, 0], page_size, rounding_mode="floor").tolist() + runs: list[tuple[int, int]] = [] + run_start = int(pages[0]) + run_len = 1 + prev = run_start + for page in pages[1:]: + page = int(page) + if page == prev + 1: + run_len += 1 + else: + runs.append((run_start, run_len)) + run_start = page + run_len = 1 + prev = page + runs.append((run_start, run_len)) + return runs + @synchronized def clear(self): # Initialize memory states and tracking structures. @@ -323,6 +568,8 @@ class HostKVCache(abc.ABC): self.free_slots = torch.arange(self.size, dtype=torch.int64) def available_size(self): + if hasattr(self, "_free_token_count"): + return int(self._free_token_count) return len(self.free_slots) @synchronized @@ -332,6 +579,10 @@ class HostKVCache(abc.ABC): ), "The requested size should be a multiple of the page size." if need_size > self.available_size(): return None + if getattr(self, "_free_extent_index_enabled", False): + return self._alloc_from_free_extents( + need_size, require_single_run=False + ) select_index = self.free_slots[:need_size] self.free_slots = self.free_slots[need_size:] @@ -356,6 +607,15 @@ class HostKVCache(abc.ABC): return None if need_size == 0: return self.alloc(need_size) + if getattr(self, "_free_extent_index_enabled", False): + select_index = self._alloc_from_free_extents( + need_size, require_single_run=True + ) + if select_index is not None: + return select_index + return self._alloc_from_free_extents( + need_size, require_single_run=False + ) fifo_prefix = self.free_slots[:need_size] expected_prefix = fifo_prefix[:1] + torch.arange( @@ -422,6 +682,20 @@ class HostKVCache(abc.ABC): @synchronized def free(self, indices: torch.Tensor) -> int: + if getattr(self, "_free_extent_index_enabled", False): + runs = self._page_runs_from_token_indices(indices) + if runs is not None: + for start, length in runs: + self._insert_free_extent(start, length) + self._free_token_count += int(indices.numel()) + self._free_slots_dirty = True + return len(indices) + logger.warning( + "[HiCache-L2-allocator] disabling extent index for non-page-shaped free indices: tokens=%d page_size=%d", + int(indices.numel()), + int(self.page_size), + ) + self._free_extent_index_enabled = False self.free_slots = torch.cat([self.free_slots, indices.cpu()]) return len(indices) diff --git a/test/registered/unit/benchmark/test_cp_hicache_allocator_bench.py b/test/registered/unit/benchmark/test_cp_hicache_allocator_bench.py index 1c98b7a50..a02897c5b 100644 --- a/test/registered/unit/benchmark/test_cp_hicache_allocator_bench.py +++ b/test/registered/unit/benchmark/test_cp_hicache_allocator_bench.py @@ -1,11 +1,14 @@ import torch from benchmark.hicache.bench_cp_hicache_allocator_overhead import ( + HostChurnBenchResult, StandaloneCPSharedPagedAllocator, StandaloneHostAllocator, + _bench_host_churn_case, _host_pages_from_gb, _make_host_free_slots, _make_page_compute_owners, + _page_run_lengths_from_token_slots, _parse_int_list, ) @@ -49,6 +52,65 @@ def test_host_random_fragmented_has_requested_size(): assert torch.unique(free_slots).numel() == free_slots.numel() +def test_page_run_lengths_from_token_slots_counts_lpf_descriptors(): + page_size = 4 + selected = torch.tensor( + [ + *range(10 * page_size, 13 * page_size), + *range(20 * page_size, 22 * page_size), + *range(25 * page_size, 26 * page_size), + ], + dtype=torch.int64, + ) + + assert _page_run_lengths_from_token_slots(selected, page_size) == [3, 2, 1] + + +def test_host_churn_case_reports_l2_run_quality(): + result = _bench_host_churn_case( + impl="standalone", + method="contiguous", + total_pages=48, + request_pages=4, + page_size=8, + target_occupancy=0.75, + evict_pages=4, + eviction_pattern="random", + repeat=4, + warmup=1, + seed=7, + ) + + assert isinstance(result, HostChurnBenchResult) + assert result.bench == "host_churn" + assert result.repeat == 4 + assert result.page_first_descriptors_per_op == 4 + assert 0 < result.lpf_descriptor_ratio_mean <= 1 + assert result.run_count_p50 >= 1 + assert result.max_run_pages_mean >= 1 + + +def test_host_churn_prefill_node_pages_can_model_fragmented_free_chunks(): + result = _bench_host_churn_case( + impl="standalone", + method="fifo", + total_pages=64, + request_pages=8, + page_size=4, + target_occupancy=0.75, + evict_pages=8, + eviction_pattern="random", + prefill_node_pages=1, + repeat=8, + warmup=2, + seed=11, + ) + + assert result.prefill_node_pages == 1 + assert result.run_count_p50 > 1 + assert result.lpf_descriptor_ratio_mean > 1 / result.request_pages + + def test_standalone_l1_allocator_reports_owner_lane_stats(): allocator = StandaloneCPSharedPagedAllocator( physical_pages=4, @@ -79,3 +141,25 @@ def test_standalone_l1_allocator_allocates_owner_matching_pages(): logical_pages = (selected.view(-1, allocator.page_size)[:, 0] // allocator.page_size) selected_owners = torch.remainder(logical_pages - 1, allocator.cp_size).tolist() assert selected_owners == owners + + +def test_host_churn_burnin_exposes_fragmented_evicted_nodes_after_cold_tail(): + result = _bench_host_churn_case( + impl="standalone", + method="fifo", + total_pages=32, + request_pages=8, + page_size=4, + target_occupancy=0.75, + evict_pages=8, + eviction_pattern="random", + prefill_node_pages=1, + burnin=1, + repeat=1, + warmup=0, + seed=23, + ) + + assert result.burnin == 1 + assert result.run_count_p50 > 1 + assert result.lpf_descriptor_ratio_mean > 1 / result.request_pages diff --git a/test/registered/unit/managers/test_hicache_controller_cp.py b/test/registered/unit/managers/test_hicache_controller_cp.py index e0c5dd96c..7199ebc52 100644 --- a/test/registered/unit/managers/test_hicache_controller_cp.py +++ b/test/registered/unit/managers/test_hicache_controller_cp.py @@ -1016,6 +1016,29 @@ class TestHiCacheControllerCPWrite(CustomTestCase): self.assertEqual(selected.tolist(), [8, 9, 10, 11, 12, 13, 14, 15]) self.assertEqual(host_pool.free_slots.tolist(), [100, 101, 102, 103]) + def test_host_alloc_contiguous_preferred_uses_lazy_extent_index(self): + host_pool = DummyHostKVCacheForAlloc.__new__(DummyHostKVCacheForAlloc) + host_pool.page_size = 4 + host_pool.lock = __import__("threading").RLock() + pages = [50, 51, 52, 53, 100, 7, 8] + host_pool.free_slots = torch.tensor( + [page * 4 + offset for page in pages for offset in range(4)], + dtype=torch.int64, + ) + + selected = host_pool.alloc_contiguous_preferred(16) + + self.assertEqual( + selected.tolist(), + [page * 4 + offset for page in [50, 51, 52, 53] for offset in range(4)], + ) + self.assertEqual(host_pool.available_size(), 12) + self.assertTrue(host_pool._free_slots_dirty) + self.assertEqual( + host_pool.free_slots.tolist(), + [page * 4 + offset for page in [7, 8, 100] for offset in range(4)], + ) + def test_cp_reserve_zero_owned_queues_no_ack_until_submit(self): host_pool = FakeHostPool(torch.tensor([], dtype=torch.int64)) controller = self.make_controller(host_pool, cp_rank=3)