from __future__ import annotations """Microbench CP HiCache allocator/control-path overhead. This benchmark targets metadata-only allocator costs that matter for large HiCache deployments. It intentionally does not allocate real 220GB host KV buffers; instead it constructs the same token/page index metadata sizes used by HostKVCache and CPSharedPagedTokenToKVPoolAllocator. Examples: # Host allocator model for 220GB-equivalent metadata. PYTHONPATH=. python benchmark/hicache/bench_cp_hicache_allocator_overhead.py \ --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 \ --request-pages 8,64,512 --l1-impl current,fifo \ --l1-ops stats,free_room_stats,select_only,alloc_pages """ import argparse import json import math import os import random import statistics import sys import threading import time import types from dataclasses import asdict, dataclass from typing import Callable, Iterable, Optional import torch @dataclass(frozen=True) class BenchResult: bench: str impl: str pattern: str device: str total_pages: int request_pages: int page_size: int repeat: int mean_us: float p50_us: float p95_us: float p99_us: float min_us: float max_us: float 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()] return [int(item) for item in value] def _parse_float_list(value: str | Iterable[float]) -> list[float]: if isinstance(value, str): return [float(item.strip()) for item in value.split(",") if item.strip()] return [float(item) for item in value] def _host_pages_from_gb(size_gb: float, *, bytes_per_token: int, page_size: int) -> int: if size_gb <= 0: raise ValueError(f"size_gb must be positive, got {size_gb}") if bytes_per_token <= 0: raise ValueError(f"bytes_per_token must be positive, got {bytes_per_token}") if page_size <= 0: raise ValueError(f"page_size must be positive, got {page_size}") return int(math.floor(size_gb * 1e9 / float(bytes_per_token * page_size))) def _percentile(samples: list[float], percentile: float) -> float: if not samples: return 0.0 ordered = sorted(samples) if len(ordered) == 1: return ordered[0] rank = (len(ordered) - 1) * percentile / 100.0 lo = int(math.floor(rank)) hi = int(math.ceil(rank)) if lo == hi: return ordered[lo] return ordered[lo] * (hi - rank) + ordered[hi] * (rank - lo) def _page_ids_to_token_slots(page_ids: list[int] | torch.Tensor, page_size: int) -> torch.Tensor: pages = torch.as_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 _make_host_page_order( *, total_pages: int, request_pages: int, pattern: str, seed: int ) -> list[int]: if total_pages <= 0: raise ValueError(f"total_pages must be positive, got {total_pages}") if request_pages <= 0: raise ValueError(f"request_pages must be positive, got {request_pages}") if request_pages > total_pages: raise ValueError( f"request_pages must be <= total_pages, got {request_pages=} {total_pages=}" ) if pattern == "contiguous_fifo": return list(range(total_pages)) if pattern == "random_fragmented": generator = torch.Generator(device="cpu") generator.manual_seed(seed) return torch.randperm(total_pages, generator=generator).tolist() if pattern == "fragmented_prefix_later_run": run_start = min(max(1, total_pages // 3), total_pages - request_pages) run = list(range(run_start, run_start + request_pages)) used = set(run) # Put non-contiguous pages before the usable run so FIFO is poor but a # later contiguous run exists. Prefer descending odd/even pages because # consecutive physical pages are unlikely to appear in the prefix. scattered: list[int] = [] for page in range(total_pages - 1, -1, -2): if page not in used: scattered.append(page) used.add(page) if len(scattered) >= request_pages: break for page in range(total_pages - 2, -1, -2): if len(scattered) >= request_pages: break if page not in used: scattered.append(page) used.add(page) rest = [page for page in range(total_pages) if page not in used] return scattered + run + rest raise ValueError(f"unsupported host pattern: {pattern}") def _make_host_free_slots( *, total_pages: int, request_pages: int, page_size: int, pattern: str, seed: int ) -> torch.Tensor: return _page_ids_to_token_slots( _make_host_page_order( total_pages=total_pages, request_pages=request_pages, pattern=pattern, seed=seed, ), page_size, ) class StandaloneHostAllocator: """Small metadata-only copy of HostKVCache allocation behavior.""" def __init__(self, *, page_size: int, free_slots: torch.Tensor): self.page_size = int(page_size) self.free_slots = free_slots self.lock = threading.RLock() def available_size(self) -> int: return int(self.free_slots.numel()) def alloc(self, need_size: int) -> Optional[torch.Tensor]: assert need_size % self.page_size == 0 if need_size > self.available_size(): return None select_index = self.free_slots[:need_size] self.free_slots = self.free_slots[need_size:] return select_index def alloc_contiguous_preferred(self, need_size: int) -> Optional[torch.Tensor]: assert need_size % self.page_size == 0 if need_size > self.available_size(): return None if need_size == 0: return self.alloc(need_size) fifo_prefix = self.free_slots[:need_size] expected_prefix = fifo_prefix[:1] + torch.arange( need_size, dtype=fifo_prefix.dtype, device=fifo_prefix.device ) if torch.equal(fifo_prefix, expected_prefix): return self.alloc(need_size) page_size = int(self.page_size) need_pages = need_size // page_size page_count = int(self.free_slots.numel()) // page_size if page_count < need_pages: return self.alloc(need_size) page_slots = self.free_slots[: page_count * page_size].view(page_count, page_size) page_offsets = torch.arange( page_size, dtype=page_slots.dtype, device=page_slots.device ) page_is_contiguous = torch.all( page_slots == (page_slots[:, :1] + page_offsets), dim=1 ) if not bool(page_is_contiguous.any()): return self.alloc(need_size) chunk_indices = torch.arange(page_count, dtype=torch.int64)[page_is_contiguous.cpu()] page_starts = page_slots[:, 0][page_is_contiguous] sorted_starts, order = torch.sort(page_starts) sorted_chunks = chunk_indices[order.cpu()] run_start = -1 if need_pages == 1: run_start = 0 else: current_run_start = 0 current_run_len = 1 diffs = sorted_starts[1:] - sorted_starts[:-1] for offset, diff in enumerate(diffs.tolist(), start=1): if diff == page_size: current_run_len += 1 if current_run_len >= need_pages: run_start = current_run_start break else: current_run_start = offset current_run_len = 1 if run_start < 0: return self.alloc(need_size) selected_chunks = sorted_chunks[run_start : run_start + need_pages] token_offsets = ( selected_chunks[:, None] * page_size + torch.arange(page_size, dtype=torch.int64)[None, :] ).reshape(-1) select_index = self.free_slots[token_offsets] keep_mask = torch.ones(self.free_slots.numel(), dtype=torch.bool) keep_mask[token_offsets] = False 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( *, required: list[int], available: list[int], capacities: list[int], target_ratio: float, trigger_ratio: float, ) -> list[int]: deficits: list[int] = [] for req, avail, capacity in zip(required, available, capacities): target_room = ( int(math.ceil(float(capacity) * float(target_ratio))) if capacity > 0 and target_ratio > 0 else 0 ) trigger_room = ( int(math.ceil(float(capacity) * float(trigger_ratio))) if capacity > 0 and trigger_ratio > 0 else 0 ) if int(avail) >= int(req) + trigger_room: deficits.append(0) else: deficits.append(max(0, int(req) + target_room - int(avail))) return deficits class StandaloneCPSharedPagedAllocator: """Metadata-only copy of the CP shared-KV page owner allocator. This intentionally mirrors the current Python/Torch control path used by ``CPSharedPagedTokenToKVPoolAllocator`` so CPU-only environments can measure the allocator shape without importing the full SGLang runtime dependency stack. The benchmark still uses the production allocator when imports are available. """ def __init__( self, *, physical_pages: int, page_size: int, cp_size: int, device: torch.device, ): self.physical_size = int(physical_pages) * int(page_size) self.page_size = int(page_size) self.cp_size = int(cp_size) self.device = device logical_pages = int(physical_pages) * int(cp_size) self._owner_free_pages = None self._owner_release_pages = None self._flat_free_pages_cache = None self._flat_release_pages_cache = None self.free_pages = torch.arange( 1, logical_pages + 1, dtype=torch.int64, device=device ) self.release_pages = torch.empty((0,), dtype=torch.int64, device=device) self.debug_mode = False def _empty_pages(self) -> torch.Tensor: return torch.empty((0,), dtype=torch.int64, device=self.device) def _split_owner_buckets( self, pages: Optional[torch.Tensor] ) -> Optional[list[torch.Tensor]]: if pages is None: return None if pages.numel() == 0: return [torch.empty_like(pages) for _ in range(self.cp_size)] owner_ids = torch.remainder(pages - 1, self.cp_size) buckets: list[torch.Tensor] = [] for owner in range(self.cp_size): owner_pages = pages[owner_ids == owner] if owner_pages.numel() > 1: owner_pages, _ = torch.sort(owner_pages) buckets.append(owner_pages) return buckets def _materialize_owner_buckets( self, buckets: Optional[list[torch.Tensor]] ) -> Optional[torch.Tensor]: if buckets is None: return None non_empty = [bucket for bucket in buckets if bucket.numel() > 0] if not non_empty: return self._empty_pages() return torch.cat(non_empty) @property def free_pages(self): if self._owner_free_pages is not None: if self._flat_free_pages_cache is None: self._flat_free_pages_cache = self._materialize_owner_buckets( self._owner_free_pages ) return self._flat_free_pages_cache return self._flat_free_pages_cache @free_pages.setter def free_pages(self, pages): self._flat_free_pages_cache = pages self._owner_free_pages = self._split_owner_buckets(pages) @property def release_pages(self): if self._owner_release_pages is not None: if self._flat_release_pages_cache is None: self._flat_release_pages_cache = self._materialize_owner_buckets( self._owner_release_pages ) return self._flat_release_pages_cache return self._flat_release_pages_cache @release_pages.setter def release_pages(self, pages): self._flat_release_pages_cache = pages self._owner_release_pages = self._split_owner_buckets(pages) def _owner_bucket_counts(self, buckets: Optional[list[torch.Tensor]]) -> list[int]: if buckets is None: return [0 for _ in range(self.cp_size)] return [int(bucket.numel()) for bucket in buckets] def _owner_available_counts(self) -> list[int]: free_counts = self._owner_bucket_counts(self._owner_free_pages) release_counts = self._owner_bucket_counts(self._owner_release_pages) return [ free_count + release_count for free_count, release_count in zip(free_counts, release_counts) ] def _consume_owner_bucket_prefix( self, *, release: bool, counts_by_owner: list[int], ) -> None: target_attr = "_owner_release_pages" if release else "_owner_free_pages" cache_attr = "_flat_release_pages_cache" if release else "_flat_free_pages_cache" buckets = getattr(self, target_attr) mutated = False for owner, count in enumerate(counts_by_owner): if count <= 0: continue buckets[owner] = buckets[owner][count:] mutated = True if mutated: setattr(self, target_attr, buckets) setattr(self, cache_attr, None) def compute_owner_lane_stats( self, page_compute_owners: list[int], ) -> tuple[list[int], list[int], list[int]]: required = [0 for _ in range(self.cp_size)] for owner in page_compute_owners: if owner < 0 or owner >= self.cp_size: raise ValueError( f"compute owner must be in [0, {self.cp_size}), got {owner}" ) required[owner] += 1 available = self._owner_available_counts() deficits = [ max(0, required_count - available_count) for required_count, available_count in zip(required, available) ] return required, available, deficits def compute_owner_lane_capacity_pages(self) -> list[int]: capacity_pages = int(self.physical_size // self.page_size) return [capacity_pages for _ in range(self.cp_size)] def compute_owner_lane_free_room_stats( self, page_compute_owners: list[int], *, target_ratio: float, trigger_ratio: float, ) -> tuple[list[int], list[int], list[int]]: required, available, _exact_deficits = self.compute_owner_lane_stats( page_compute_owners ) deficits = _compute_owner_lane_free_room_deficits( required=required, available=available, capacities=self.compute_owner_lane_capacity_pages(), target_ratio=target_ratio, trigger_ratio=trigger_ratio, ) return required, available, deficits def _select_owner_free_pages_prefer_contiguous( self, owner_pages: torch.Tensor, required_count: int, ) -> torch.Tensor: return owner_pages[: min(required_count, int(owner_pages.numel()))] def _select_compute_owner_pages( self, page_compute_owners: list[int], ) -> Optional[tuple[torch.Tensor, list[int], list[int]]]: if not page_compute_owners: return ( torch.empty((0,), dtype=torch.int64, device=self.device), [0 for _ in range(self.cp_size)], [0 for _ in range(self.cp_size)], ) required_by_owner = [0 for _ in range(self.cp_size)] positions_by_owner: list[list[int]] = [[] for _ in range(self.cp_size)] for position, owner in enumerate(page_compute_owners): if owner < 0 or owner >= self.cp_size: raise ValueError( f"compute owner must be in [0, {self.cp_size}), got {owner}" ) required_by_owner[owner] += 1 positions_by_owner[owner].append(position) lane_pages = [None for _ in range(self.cp_size)] selected_free_counts = [0 for _ in range(self.cp_size)] selected_release_counts = [0 for _ in range(self.cp_size)] for owner, required_count in enumerate(required_by_owner): if required_count == 0: continue selected_owner_free_mask = self._select_owner_free_pages_prefer_contiguous( self._owner_free_pages[owner], required_count ) selected_owner_pages = selected_owner_free_mask free_count = int(selected_owner_pages.numel()) remaining_count = required_count - free_count if remaining_count > 0: release_bucket = self._owner_release_pages[owner] if remaining_count > release_bucket.numel(): return None selected_owner_release_pages = release_bucket[:remaining_count] selected_owner_pages = torch.cat( (selected_owner_pages, selected_owner_release_pages) ) selected_release_counts[owner] = remaining_count selected_free_counts[owner] = free_count lane_pages[owner] = selected_owner_pages selected_pages = torch.empty( (len(page_compute_owners),), dtype=torch.int64, device=self.device ) for owner, positions in enumerate(positions_by_owner): if not positions: continue position_tensor = torch.tensor( positions, dtype=torch.int64, device=self.device ) selected_pages[position_tensor] = lane_pages[owner] return ( selected_pages, selected_free_counts, selected_release_counts, ) def alloc_pages_with_owners( self, page_compute_owners: list[int], ) -> Optional[torch.Tensor]: if not page_compute_owners: return torch.empty((0,), dtype=torch.int64, device=self.device) selected = self._select_compute_owner_pages(page_compute_owners) if selected is None: return None selected_pages, selected_free_counts, selected_release_counts = selected page_size = self.page_size base = selected_pages.to(torch.int64).unsqueeze(1) * page_size offsets = torch.arange( page_size, dtype=torch.int64, device=self.device ).unsqueeze(0) out_indices = (base + offsets).reshape(-1) self._consume_owner_bucket_prefix( release=False, counts_by_owner=selected_free_counts ) self._consume_owner_bucket_prefix( release=True, counts_by_owner=selected_release_counts ) return out_indices def _is_page_contiguous_selection(selected: Optional[torch.Tensor], page_size: int) -> bool: if selected is None or selected.numel() == 0: return False pages = selected.view(-1, page_size)[:, 0] // page_size if pages.numel() <= 1: return True 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) if impl == "production": from sglang.srt.mem_cache.memory_pool_host import HostKVCache class DummyHostKVCacheForBench(HostKVCache): def get_size_per_token(self): return 1 def init_kv_buffer(self): return None def load_to_device_per_layer(self, *args, **kwargs) -> None: pass def backup_from_device_per_layer(self, *args, **kwargs) -> None: pass def backup_from_device_all_layer(self, *args, **kwargs) -> None: pass def get_data_page(self, *args, **kwargs) -> torch.Tensor: return torch.empty((0,), dtype=torch.uint8) def get_dummy_flat_data_page(self) -> torch.Tensor: return torch.empty((0,), dtype=torch.uint8) def set_from_flat_data_page(self, *args, **kwargs) -> None: pass allocator = DummyHostKVCacheForBench.__new__(DummyHostKVCacheForBench) allocator.page_size = int(page_size) allocator.lock = threading.RLock() allocator.free_slots = free_slots return allocator raise ValueError(f"unsupported host impl: {impl}") def _summarize( *, bench: str, impl: str, pattern: str, device: str, total_pages: int, request_pages: int, page_size: int, samples_us: list[float], contiguous_hits: int, ) -> BenchResult: return BenchResult( bench=bench, impl=impl, pattern=pattern, device=device, total_pages=int(total_pages), request_pages=int(request_pages), page_size=int(page_size), repeat=len(samples_us), 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 / len(samples_us)) if samples_us else 0.0, ) 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, method: str, total_pages: int, request_pages: int, page_size: int, pattern: str, repeat: int, warmup: int, seed: int, ) -> BenchResult: base_free_slots = _make_host_free_slots( total_pages=total_pages, request_pages=request_pages, page_size=page_size, pattern=pattern, seed=seed, ) need_size = request_pages * page_size samples_us: list[float] = [] contiguous_hits = 0 for iteration in range(warmup + repeat): allocator = _make_host_allocator( impl, page_size=page_size, free_slots=base_free_slots.clone() ) fn = allocator.alloc if method == "fifo" else allocator.alloc_contiguous_preferred start_ns = time.perf_counter_ns() selected = fn(need_size) elapsed_us = (time.perf_counter_ns() - start_ns) / 1000.0 if iteration >= warmup: samples_us.append(elapsed_us) contiguous_hits += int(_is_page_contiguous_selection(selected, page_size)) return _summarize( bench="host", impl=f"{impl}:{method}", pattern=pattern, device="cpu", total_pages=total_pages, request_pages=request_pages, page_size=page_size, samples_us=samples_us, contiguous_hits=contiguous_hits, ) 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 rem = num_pages % segment_num owners: list[int] = [] for segment_idx in range(segment_num): count = base + (1 if segment_idx < rem else 0) owner = segment_idx if segment_idx < cp_size else segment_num - segment_idx - 1 owners.extend([owner] * count) return owners[:num_pages] def _make_page_compute_owners(request_pages: int, cp_size: int, pattern: str) -> list[int]: if pattern == "round_robin": return [idx % cp_size for idx in range(request_pages)] if pattern == "single_owner": return [0 for _ in range(request_pages)] if pattern == "zigzag": return _zigzag_owners(request_pages, cp_size) raise ValueError(f"unsupported owner pattern: {pattern}") def _make_l1_free_pages( *, physical_pages: int, cp_size: int, request_owners: list[int], pattern: str, device: torch.device, seed: int, ) -> torch.Tensor: logical_pages = physical_pages * cp_size all_pages = list(range(1, logical_pages + 1)) if pattern == "sequential": return torch.tensor(all_pages, dtype=torch.int64, device=device) if pattern == "random": generator = torch.Generator(device="cpu") generator.manual_seed(seed) return torch.randperm(logical_pages, generator=generator, dtype=torch.int64).to(device) + 1 if pattern != "owner_fragmented_later_run": raise ValueError(f"unsupported L1 pattern: {pattern}") required_by_owner = [0 for _ in range(cp_size)] for owner in request_owners: required_by_owner[owner] += 1 used: set[int] = set() front: list[int] = [] runs: list[int] = [] max_ordinal = physical_pages - 1 for owner, count in enumerate(required_by_owner): if count <= 0: continue run_start = min(max(4, count * 3), max(0, physical_pages - count)) owner_run = [owner + 1 + cp_size * (run_start + idx) for idx in range(count)] if owner_run[-1] > logical_pages: owner_run = [] for page in owner_run: used.add(page) runs.extend(owner_run) # Fragmented prefix for the same owner, avoiding the later run. ordinal = max_ordinal added = 0 while ordinal >= 0 and added < count: page = owner + 1 + cp_size * ordinal ordinal -= 2 if page > logical_pages or page in used: continue front.append(page) used.add(page) added += 1 rest = [page for page in all_pages if page not in used] return torch.tensor(front + runs + rest, dtype=torch.int64, device=device) def _install_sgl_kernel_stubs() -> None: if "sgl_kernel" not in sys.modules: mod = types.ModuleType("sgl_kernel") mod.__file__ = "sgl_kernel_stub.py" mod.__path__ = [] def _getattr(name): if name.startswith("__"): raise AttributeError(name) fn = lambda *args, **kwargs: None setattr(mod, name, fn) return fn mod.__getattr__ = _getattr sys.modules["sgl_kernel"] = mod for submodule in ("sgl_kernel.kvcacheio", "sgl_kernel.quantization"): if submodule not in sys.modules: sub = types.ModuleType(submodule) sub.__file__ = submodule.replace(".", "_") + "_stub.py" def _sub_getattr(name, _sub=sub): if name.startswith("__"): raise AttributeError(name) fn = lambda *args, **kwargs: None setattr(_sub, name, fn) return fn sub.__getattr__ = _sub_getattr sys.modules[submodule] = sub def _make_l1_allocator( *, physical_pages: int, page_size: int, cp_size: int, device: torch.device, production: bool, ): if not production: return StandaloneCPSharedPagedAllocator( physical_pages=physical_pages, page_size=page_size, cp_size=cp_size, device=device, ) from sglang.srt.mem_cache.allocator import CPSharedPagedTokenToKVPoolAllocator return CPSharedPagedTokenToKVPoolAllocator( logical_size=physical_pages * cp_size * page_size, physical_size=physical_pages * page_size, page_size=page_size, dtype=torch.bfloat16, device=str(device), kvcache=None, need_sort=False, cp_size=cp_size, cp_rank=0, ) def _patch_l1_fifo_selector(allocator) -> None: def fifo_selector(owner_pages: torch.Tensor, required_count: int) -> torch.Tensor: return owner_pages[: min(required_count, int(owner_pages.numel()))] allocator._select_owner_free_pages_prefer_contiguous = fifo_selector def _is_l1_selection_physically_contiguous( selected: Optional[torch.Tensor], *, page_size: int, cp_size: int ) -> bool: if selected is None or selected.numel() == 0: return False logical_pages = selected.view(-1, page_size)[:, 0] // page_size owners = torch.remainder(logical_pages - 1, cp_size) physical_pages = torch.div(logical_pages - 1, cp_size, rounding_mode="floor") + 1 for owner in torch.unique(owners).tolist(): owner_phys = physical_pages[owners == int(owner)] if owner_phys.numel() <= 1: continue if not bool(torch.all(owner_phys[1:] - owner_phys[:-1] == 1).item()): return False return True def _bench_l1_case( *, impl: str, op: str, physical_pages: int, request_pages: int, page_size: int, cp_size: int, device: torch.device, free_pattern: str, owner_pattern: str, repeat: int, warmup: int, seed: int, free_room_ratio: float, free_room_trigger_ratio: float, production_allocator: bool, ) -> BenchResult: request_owners = _make_page_compute_owners(request_pages, cp_size, owner_pattern) base_free_pages = _make_l1_free_pages( physical_pages=physical_pages, cp_size=cp_size, request_owners=request_owners, pattern=free_pattern, device=device, seed=seed, ) samples_us: list[float] = [] contiguous_hits = 0 use_cuda = device.type == "cuda" if use_cuda: torch.cuda.synchronize(device) for iteration in range(warmup + repeat): allocator = _make_l1_allocator( physical_pages=physical_pages, page_size=page_size, cp_size=cp_size, device=device, production=production_allocator, ) allocator.free_pages = base_free_pages.clone() allocator.release_pages = torch.empty((0,), dtype=torch.int64, device=device) if impl == "fifo": _patch_l1_fifo_selector(allocator) if use_cuda: torch.cuda.synchronize(device) start_ns = time.perf_counter_ns() selected = None if op == "stats": allocator.compute_owner_lane_stats(request_owners) elif op == "free_room_stats": allocator.compute_owner_lane_free_room_stats( request_owners, target_ratio=free_room_ratio, trigger_ratio=free_room_trigger_ratio, ) elif op == "select_only": selected_result = allocator._select_compute_owner_pages(request_owners) if selected_result is not None: selected_pages = selected_result[0] base = selected_pages.to(torch.int64).unsqueeze(1) * page_size offsets = torch.arange( page_size, dtype=torch.int64, device=device ).unsqueeze(0) selected = (base + offsets).reshape(-1) elif op == "alloc_pages": selected = allocator.alloc_pages_with_owners(request_owners) else: raise ValueError(f"unsupported l1 op: {op}") if use_cuda: torch.cuda.synchronize(device) elapsed_us = (time.perf_counter_ns() - start_ns) / 1000.0 if iteration >= warmup: samples_us.append(elapsed_us) contiguous_hits += int( _is_l1_selection_physically_contiguous( selected, page_size=page_size, cp_size=cp_size ) ) return _summarize( bench="l1", impl=f"{impl}:{op}", pattern=f"{free_pattern}:{owner_pattern}", device=device.type, total_pages=physical_pages, request_pages=request_pages, page_size=page_size, samples_us=samples_us, contiguous_hits=contiguous_hits, ) 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} " f"p50={result.p50_us:9.2f}us p95={result.p95_us:9.2f}us " f"p99={result.p99_us:9.2f}us mean={result.mean_us:9.2f}us " f"contig={result.contiguous_ratio:.2f}" ) def _run_host(args) -> list[BenchResult]: 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) patterns = [item.strip() for item in args.patterns.split(",") if item.strip()] 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()] results: list[BenchResult] = [] for total_pages in host_pages: for request_pages in request_pages_list: if request_pages > total_pages: continue for pattern in patterns: for impl in host_impls: for method in methods: results.append( _bench_host_case( impl=impl, method=method, total_pages=total_pages, request_pages=request_pages, page_size=args.page_size, pattern=pattern, repeat=args.repeat, warmup=args.warmup, seed=args.seed, ) ) 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() if args.device == "cuda" and not torch.cuda.is_available(): raise RuntimeError("--device cuda requested but CUDA is not available") device = torch.device(args.device) if device.type == "cuda": torch.cuda.set_device(args.cuda_device) device = torch.device(f"cuda:{args.cuda_device}") physical_pages_list = _parse_int_list(args.physical_pages) request_pages_list = _parse_int_list(args.request_pages) free_patterns = [item.strip() for item in args.l1_free_patterns.split(",") if item.strip()] owner_patterns = [item.strip() for item in args.l1_owner_patterns.split(",") if item.strip()] impls = [item.strip() for item in args.l1_impl.split(",") if item.strip()] ops = [item.strip() for item in args.l1_ops.split(",") if item.strip()] results: list[BenchResult] = [] for physical_pages in physical_pages_list: for request_pages in request_pages_list: if request_pages > physical_pages * args.cp_size: continue for free_pattern in free_patterns: for owner_pattern in owner_patterns: for impl in impls: for op in ops: results.append( _bench_l1_case( impl=impl, op=op, physical_pages=physical_pages, request_pages=request_pages, page_size=args.page_size, cp_size=args.cp_size, device=device, free_pattern=free_pattern, owner_pattern=owner_pattern, repeat=args.repeat, warmup=args.warmup, seed=args.seed, free_room_ratio=args.l1_free_room_ratio, free_room_trigger_ratio=args.l1_free_room_trigger_ratio, production_allocator=args.l1_allocator == "production", ) ) return results def _build_parser() -> argparse.ArgumentParser: parser = argparse.ArgumentParser(description=__doc__) 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) parser.add_argument("--seed", type=int, default=0) parser.add_argument("--json-output", default="") parser.add_argument( "--host-pages", default="", help="comma list of host pages; defaults to 8192,16384,32768 only when no --host-sizes-gb is set", ) parser.add_argument("--host-sizes-gb", default="") parser.add_argument("--bytes-per-token", type=int, default=100_000) parser.add_argument("--request-pages", default="1,8,64,512") parser.add_argument( "--patterns", default="contiguous_fifo,fragmented_prefix_later_run,random_fragmented", ) 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) parser.add_argument("--device", choices=("cpu", "cuda"), default="cpu") parser.add_argument("--cuda-device", type=int, default=0) parser.add_argument("--stub-sgl-kernel", action="store_true") parser.add_argument("--l1-impl", default="fifo,current") parser.add_argument("--l1-ops", default="alloc_pages") parser.add_argument( "--l1-allocator", choices=("production", "standalone"), default="production", help="Use production allocator imports or the dependency-light metadata copy.", ) parser.add_argument("--l1-free-room-ratio", type=float, default=0.15) parser.add_argument("--l1-free-room-trigger-ratio", type=float, default=0.05) parser.add_argument( "--l1-free-patterns", default="sequential,owner_fragmented_later_run,random", ) parser.add_argument("--l1-owner-patterns", default="round_robin,single_owner,zigzag") return parser def main(argv: Optional[list[str]] = None) -> int: args = _build_parser().parse_args(argv) if args.repeat <= 0: raise ValueError("--repeat must be positive") if args.warmup < 0: raise ValueError("--warmup must be non-negative") results: list[BenchResult] = [] 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)) for result in results: print(_format_result(result), flush=True) if args.json_output: output_path = os.path.abspath(args.json_output) with open(output_path, "w", encoding="utf-8") as fout: json.dump([asdict(result) for result in results], fout, indent=2) print(f"wrote_json={output_path}", flush=True) return 0 if __name__ == "__main__": raise SystemExit(main())