diff --git a/benchmark/hicache/bench_cp_hicache_allocator_overhead.py b/benchmark/hicache/bench_cp_hicache_allocator_overhead.py new file mode 100644 index 000000000..cf956efc6 --- /dev/null +++ b/benchmark/hicache/bench_cp_hicache_allocator_overhead.py @@ -0,0 +1,742 @@ +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 + + # Production 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 +""" + +import argparse +import json +import math +import os +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 + + +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 _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 _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 _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 _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): + 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_mask: torch.Tensor, required_count: int) -> torch.Tensor: + return owner_mask & (torch.cumsum(owner_mask.to(torch.int64), dim=0) <= required_count) + + 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, + 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, +) -> 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, + ) + 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 = allocator.alloc_pages_with_owners(request_owners) + 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=impl, + 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) -> str: + 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_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()] + + 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: + results.append( + _bench_l1_case( + impl=impl, + 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, + ) + ) + return results + + +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("--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("--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-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 "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()) diff --git a/docs/advanced_features/nsa_prefill_cp_hicache_reactive_host_free_room_plan.md b/docs/advanced_features/nsa_prefill_cp_hicache_reactive_host_free_room_plan.md index 06219758a..cfaf1efba 100644 --- a/docs/advanced_features/nsa_prefill_cp_hicache_reactive_host_free_room_plan.md +++ b/docs/advanced_features/nsa_prefill_cp_hicache_reactive_host_free_room_plan.md @@ -203,3 +203,30 @@ Runtime validation: eviction may still fail; failures must stay explicit and warning-level. 4. The first version does not improve physical page contiguity; it only reduces eviction frequency. + +## 2026-06-02 correction: trigger and target are separate + +The first draft only added target room when `required > available`. The current +accepted contract separates **when** eviction starts from **how far** eviction +continues: + +```text +if available >= required + trigger_room: + deficit = 0 +else: + deficit = max(0, required + target_room - available) +``` + +This keeps remaining cache usable when the current reservation fits, while still +making a triggered eviction release enough extra room to reduce repeated eviction +and give subsequent allocators a better chance of finding contiguous page runs. + +Implications: + +- `trigger_room=0` preserves the conservative behavior: evict only when the + current reservation would not fit. +- `target_room>0` makes a triggered eviction release extra pages. +- L1/device and L2/host should use the same trigger/target model, but with + separate ratios because L1 hit value and host capacity pressure are different. +- Free-room alone is not sufficient for contiguous allocation. `HostKVCache` and + CP owner-lane device allocation also need contiguous-preferred selection. diff --git a/docs/superpowers/plans/2026-06-02-cp-hicache-l1-l2-free-room-contiguous-allocation.md b/docs/superpowers/plans/2026-06-02-cp-hicache-l1-l2-free-room-contiguous-allocation.md new file mode 100644 index 000000000..8911483c0 --- /dev/null +++ b/docs/superpowers/plans/2026-06-02-cp-hicache-l1-l2-free-room-contiguous-allocation.md @@ -0,0 +1,507 @@ +# CP HiCache L1/L2 Free-Room and Contiguous Allocation Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** Reduce repeated L1/L2 eviction churn under CP HiCache pressure and make later allocations more likely to receive contiguous page runs. + +**Architecture:** Add trigger/target free-room watermarks for L2 host and L1 device owner lanes. Trigger controls when extra eviction starts; target controls how much room remains after a triggered eviction. Preserve CP owner-lane correctness and avoid new collectives. Add contiguous-preferred allocation only as best effort: it may improve descriptor coalescing, but must never choose pages with the wrong owner. + +**Tech Stack:** Python, PyTorch tensors, SGLang CP shared-KV allocator, HiRadixCache, HiCacheController, unit tests under `test/registered/unit`. + +--- + +## Findings from code reading + +1. `HostKVCache.alloc()` currently slices `free_slots[:need_size]`; `HostKVCache.free()` appends freed token indices. Over-evicting alone cannot guarantee contiguous later host allocations because older fragmented `free_slots` remain at the front. +2. CP host write admission lives in `HiRadixCache._cp_build_write_admission()`. It currently computes exact per-owner deficits: `max(0, required - available)` for target and draft host pools. +3. CP L1 owner-lane prefill allocation uses `alloc_paged_token_slots_extend()` -> `CPSharedPagedTokenToKVPoolAllocator.alloc_extend_compute_owner()`. If first allocation fails, `_evict_for_compute_owner_lanes()` evicts exact owner deficits. +4. CP L1 load-back allocation uses `HiRadixCache._build_cp_load_back_plan()` -> `HiCacheController.load_cp()` -> `alloc_pages_with_owners()`. The load-back plan also computes exact owner-lane deficits. +5. `CPSharedPagedTokenToKVPoolAllocator._select_compute_owner_pages()` preserves owner correctness but is not contiguity-aware beyond consuming the current `free_pages`/`release_pages` order. +6. Free-room must be trigger/target based, not always-on target maintenance. The accepted formula is: + ```text + if available >= required + trigger_room: + deficit = 0 + else: + deficit = max(0, required + target_room - available) + ``` +7. Implementation correction: L1 owner-lane free-room deficits are page-unit deficits, not token-unit deficits. The owner-lane eviction API already consumes page counts in `owner_lane_deficits`, while `num_tokens` remains the scalar coarse eviction amount. +8. Implementation correction: L1 contiguous-preferred allocation must preserve owner order in the final request page sequence but may change which logical page IDs satisfy each owner. Existing tests that asserted exact FIFO pages were updated to assert the new contiguous preference rather than the old incidental page choice. + +## Progress as of 2026-06-02 + +- Done: Task 1-3 host/L2 ratio knobs, trigger-target helper, and CP host write admission. +- Done: Task 4 host `alloc_contiguous_preferred()` and CP target/draft reservation use. +- Done: Task 5 L1 owner-lane free-room in prefill eviction and CP load-back planning. +- Done: Task 6 L1 owner-lane contiguous-preferred free-page selection without `torch.isin`, sort-merge, or CUDA `.tolist()` path. +- Verification: remote py_compile passed for touched runtime/test files. Remote targeted tests passed with an sgl_kernel stub runner for allocator/load-back tests because importing the installed `sgl_kernel` directly can abort this container during pytest collection. + +## Files + +- Modify: `python/sglang/srt/server_args.py` + - Add four runtime knobs: host target/trigger ratios and L1 target/trigger ratios. + - Validate each ratio and `trigger <= target`. +- Modify: `python/sglang/srt/mem_cache/hiradix_cache.py` + - Add page-aligned free-room helper. + - Apply helper to CP host write admission. + - Apply helper to CP L1 load-back owner-lane plan. + - Store configured ratios on `HiRadixCache`. +- Modify: `python/sglang/srt/mem_cache/common.py` + - Apply L1 owner-lane trigger/target deficits before `_evict_for_compute_owner_lanes()` calls `tree_cache.evict()`. +- Modify: `python/sglang/srt/mem_cache/allocator.py` + - Add owner-lane capacity helper for CP shared paged allocator. + - Add optional owner-lane contiguous-preferred page selection without changing owner correctness. +- Modify: `python/sglang/srt/mem_cache/memory_pool_host.py` + - Add `alloc_contiguous_preferred()` that finds page-contiguous host token runs before falling back to FIFO. +- Modify: `python/sglang/srt/managers/cache_controller.py` + - Use host `alloc_contiguous_preferred()` for CP target/draft host reservations when available. +- Test: `test/registered/unit/mem_cache/test_cp_hicache_metadata.py` +- Test: `test/registered/unit/mem_cache/test_cp_shared_kv_layout.py` +- Test: `test/registered/unit/managers/test_hicache_controller_cp.py` + +--- + +### Task 1: Add ratio knobs and validation + +**Files:** +- Modify: `python/sglang/srt/server_args.py` +- Test: `test/registered/unit/mem_cache/test_cp_hicache_metadata.py` + +- [ ] **Step 1: Write failing validation tests** + +Add tests that instantiate `ServerArgs` with invalid ratios and expect `ValueError` from `__post_init__()` or the existing validation path. Cover negative ratio and trigger greater than target. + +- [ ] **Step 2: Run validation tests and verify RED** + +Run: +```bash +cd /root/sglang-work/sglang-dev +PYTHONPATH=python python -m pytest -q test/registered/unit/mem_cache/test_cp_hicache_metadata.py -k free_room_args +``` +Expected: FAIL because the new attributes do not exist or validation is missing. + +- [ ] **Step 3: Add args** + +Add dataclass fields with default `0.0`: +```python +hicache_host_free_room_ratio: float = 0.0 +hicache_host_free_room_trigger_ratio: float = 0.0 +hicache_l1_free_room_ratio: float = 0.0 +hicache_l1_free_room_trigger_ratio: float = 0.0 +``` + +Add parser args near existing HiCache args: +```python +--hicache-host-free-room-ratio +--hicache-host-free-room-trigger-ratio +--hicache-l1-free-room-ratio +--hicache-l1-free-room-trigger-ratio +``` + +- [ ] **Step 4: Add validation** + +In `_handle_hicache()` or adjacent HiCache validation, reject values outside `[0, 1)` and reject `trigger > target` for host and L1 pairs. + +- [ ] **Step 5: Run validation tests and verify GREEN** + +Run the same targeted pytest. Expected: PASS. + +--- + +### Task 2: Implement reusable trigger/target free-room deficit helper + +**Files:** +- Modify: `python/sglang/srt/mem_cache/hiradix_cache.py` +- Test: `test/registered/unit/mem_cache/test_cp_hicache_metadata.py` + +- [ ] **Step 1: Write failing helper tests** + +Add tests for a helper with page-size rounding: + +- required fits above trigger -> deficit `0`. +- required fits but below target only -> still deficit `0` if trigger is `0`. +- required does not fit trigger -> deficit includes target room. +- ratio target is rounded up to page size. + +- [ ] **Step 2: Run helper tests and verify RED** + +Run: +```bash +cd /root/sglang-work/sglang-dev +PYTHONPATH=python python -m pytest -q test/registered/unit/mem_cache/test_cp_hicache_metadata.py -k free_room_deficit +``` +Expected: FAIL because the helper does not exist. + +- [ ] **Step 3: Add helper** + +Add a pure helper in `hiradix_cache.py`: +```python +def _page_aligned_room_from_ratio(capacity: int, ratio: float, page_size: int) -> int: + if ratio <= 0 or capacity <= 0: + return 0 + raw = int(math.ceil(float(capacity) * float(ratio))) + return ((raw + page_size - 1) // page_size) * page_size + + +def _free_room_deficit( + *, + required: int, + available: int, + capacity: int, + page_size: int, + target_ratio: float, + trigger_ratio: float, +) -> int: + target_room = _page_aligned_room_from_ratio(capacity, target_ratio, page_size) + trigger_room = _page_aligned_room_from_ratio(capacity, trigger_ratio, page_size) + if available >= required + trigger_room: + return 0 + return max(0, required + target_room - available) +``` + +- [ ] **Step 4: Run helper tests and verify GREEN** + +Run targeted pytest. Expected: PASS. + +--- + +### Task 3: Apply free-room to L2 CP host write admission + +**Files:** +- Modify: `python/sglang/srt/mem_cache/hiradix_cache.py` +- Test: `test/registered/unit/mem_cache/test_cp_hicache_metadata.py` + +- [ ] **Step 1: Write failing admission tests** + +Add tests around `_cp_build_write_admission()` using `HiRadixCache.__new__()`: + +1. `host_free_room_ratio=0.0` preserves current exact deficits. +2. If `available >= required + trigger_room`, deficit is zero even if `available < required + target_room`. +3. If `available < required + trigger_room`, deficit is `required + target_room - available`. +4. Draft host lower availability drives the max deficit. + +- [ ] **Step 2: Run tests and verify RED** + +Run: +```bash +cd /root/sglang-work/sglang-dev +PYTHONPATH=python python -m pytest -q test/registered/unit/mem_cache/test_cp_hicache_metadata.py -k cp_host_free_room +``` +Expected: FAIL because admission still uses exact deficits. + +- [ ] **Step 3: Store ratios and update admission** + +In `HiRadixCache.__init__`, store: +```python +self.hicache_host_free_room_ratio = server_args.hicache_host_free_room_ratio +self.hicache_host_free_room_trigger_ratio = server_args.hicache_host_free_room_trigger_ratio +self.hicache_l1_free_room_ratio = server_args.hicache_l1_free_room_ratio +self.hicache_l1_free_room_trigger_ratio = server_args.hicache_l1_free_room_trigger_ratio +``` + +In `_cp_build_write_admission()`, compute target and draft deficits with `_free_room_deficit()` using snapshot capacity per owner. + +- [ ] **Step 4: Run admission tests and verify GREEN** + +Run targeted pytest. Expected: PASS. + +--- + +### Task 4: Add contiguous-preferred host allocation and use it for CP HiCache reservations + +**Files:** +- Modify: `python/sglang/srt/mem_cache/memory_pool_host.py` +- Modify: `python/sglang/srt/managers/cache_controller.py` +- Test: `test/registered/unit/managers/test_hicache_controller_cp.py` + +- [ ] **Step 1: Write failing host allocator test** + +Create a fake or direct `HostKVCache.__new__()` instance with `page_size=4` and fragmented `free_slots` such as: +```python +[100,101,102,103, 8,9,10,11, 12,13,14,15] +``` +Request `8` tokens. Expected contiguous-preferred allocation returns `[8..15]`, not the fragmented FIFO prefix if the prefix does not satisfy the requested run. + +- [ ] **Step 2: Write failing controller test** + +Extend `FakeHostPool` with `alloc_contiguous_preferred()` and verify `reserve_write_cp()` calls it instead of `alloc()` for CP target and draft host pools. + +- [ ] **Step 3: Verify RED** + +Run: +```bash +cd /root/sglang-work/sglang-dev +PYTHONPATH=python python -m pytest -q test/registered/unit/managers/test_hicache_controller_cp.py -k contiguous_preferred +``` +Expected: FAIL because the method/call does not exist. + +- [ ] **Step 4: Implement host method** + +Add `HostKVCache.alloc_contiguous_preferred(need_size)`: + +- assert page-aligned size; +- return `None` if capacity is insufficient; +- if FIFO prefix is already a contiguous token run, reuse `alloc()` behavior; +- otherwise scan page starts in `free_slots` and select the first run of `need_size // page_size` consecutive pages; +- remove selected indices by boolean mask; +- fallback to `alloc()` if no run exists. + +- [ ] **Step 5: Use it in CP reservation** + +In `HiCacheController.reserve_write_cp()`, replace direct host allocation with: +```python +alloc_host = getattr(self.mem_pool_host, "alloc_contiguous_preferred", self.mem_pool_host.alloc) +host_indices = alloc_host(len(physical_device_indices)) +``` +Do the same for `draft_mem_pool_host`. + +- [ ] **Step 6: Verify GREEN** + +Run the targeted controller tests. Expected: PASS. + +--- + +### Task 5: Apply L1 owner-lane free-room to prefill allocation and load-back + +**Files:** +- Modify: `python/sglang/srt/mem_cache/allocator.py` +- Modify: `python/sglang/srt/mem_cache/common.py` +- Modify: `python/sglang/srt/mem_cache/hiradix_cache.py` +- Test: `test/registered/unit/mem_cache/test_cp_shared_kv_layout.py` +- Test: `test/registered/unit/mem_cache/test_cp_hicache_metadata.py` + +- [ ] **Step 1: Write failing allocator capacity tests** + +Add a test for `CPSharedPagedTokenToKVPoolAllocator.compute_owner_lane_capacity_pages()` expecting each lane capacity to equal `num_pages // cp_size`. + +- [ ] **Step 2: Write failing prefill eviction tests** + +Extend `_evict_for_compute_owner_lanes()` tests: + +- with target/trigger ratio zero, current exact deficits are preserved; +- with trigger zero and target room two pages, an exhausted lane evicts `required + target_room - available` pages for that owner; +- with available satisfying trigger, no extra eviction runs. + +- [ ] **Step 3: Write failing load-back plan tests** + +Add a `HiRadixCache.__new__()` test proving `_build_cp_load_back_plan()` uses target/trigger L1 room in `deficit_by_owner`. + +- [ ] **Step 4: Verify RED** + +Run: +```bash +cd /root/sglang-work/sglang-dev +PYTHONPATH=python python -m pytest -q \ + test/registered/unit/mem_cache/test_cp_shared_kv_layout.py -k "owner_lane and free_room" \ + test/registered/unit/mem_cache/test_cp_hicache_metadata.py -k "load_back and free_room" +``` +Expected: FAIL because L1 free-room is not implemented. + +- [ ] **Step 5: Implement allocator capacity helper** + +Add to `CPSharedPagedTokenToKVPoolAllocator`: +```python +def compute_owner_lane_capacity_pages(self) -> list[int]: + capacity = self.num_pages // self.cp_size + return [int(capacity) for _ in range(self.cp_size)] +``` + +- [ ] **Step 6: Update `_evict_for_compute_owner_lanes()`** + +Give it optional `target_ratio` and `trigger_ratio`. Compute deficits from required/available/capacity using page units. Preserve current behavior when both ratios are zero. + +- [ ] **Step 7: Pass ratios from `alloc_paged_token_slots_extend()`** + +When CP shared KV server args are enabled, pass `server_args.hicache_l1_free_room_ratio` and `server_args.hicache_l1_free_room_trigger_ratio` into `_evict_for_compute_owner_lanes()`. + +- [ ] **Step 8: Update load-back plan** + +In `_build_cp_load_back_plan()` and `_refresh_cp_load_back_plan()`, compute L1 deficits with the same trigger/target helper. Use allocator capacity pages converted to token counts or compute directly in pages consistently with existing `compute_owner_lane_stats()` output. + +- [ ] **Step 9: Verify GREEN** + +Run targeted tests. Expected: PASS. + +--- + +### Task 6: Add L1 owner-lane contiguous-preferred page selection + +**Files:** +- Modify: `python/sglang/srt/mem_cache/allocator.py` +- Test: `test/registered/unit/mem_cache/test_cp_shared_kv_layout.py` + +- [ ] **Step 1: Write failing contiguity preference tests** + +Use CPU allocator with fragmented owner-lane pages. For one owner requiring three pages, set candidates so the first owner-correct pages are fragmented but a later three-page physical run exists. Verify selected logical pages map to consecutive physical pages for that owner. + +- [ ] **Step 2: Verify RED** + +Run: +```bash +cd /root/sglang-work/sglang-dev +PYTHONPATH=python python -m pytest -q test/registered/unit/mem_cache/test_cp_shared_kv_layout.py -k contiguous_owner_lane +``` +Expected: FAIL because current selection takes first owner pages. + +- [ ] **Step 3: Implement best-effort selection** + +In `_select_compute_owner_pages()`: + +- keep existing owner correctness; +- prefer a contiguous run in each owner lane when one can satisfy that owner's required count; +- otherwise fall back to the current first-available behavior; +- do not introduce CPU `.tolist()` or host sync in the CUDA hot path; +- preserve the existing no-`torch.isin` and no forced sort-merge tests. + +- [ ] **Step 4: Verify GREEN** + +Run targeted allocator tests. Expected: PASS. + +--- + +### Task 7: Full verification and remote sync + +**Files:** all modified files. + +- [ ] **Step 1: Run local CPU tests** + +Run: +```bash +cd /root/sglang-work/sglang-dev +PYTHONPATH=python python -m pytest -q \ + test/registered/unit/mem_cache/test_cp_shared_kv_layout.py \ + test/registered/unit/mem_cache/test_cp_hicache_metadata.py \ + test/registered/unit/managers/test_hicache_controller_cp.py +``` +Expected: PASS or dependency-related skip only; real failures must be fixed. + +- [ ] **Step 2: Compile touched files** + +Run: +```bash +cd /root/sglang-work/sglang-dev +python -m py_compile \ + python/sglang/srt/server_args.py \ + python/sglang/srt/mem_cache/hiradix_cache.py \ + python/sglang/srt/mem_cache/common.py \ + python/sglang/srt/mem_cache/allocator.py \ + python/sglang/srt/mem_cache/memory_pool_host.py \ + python/sglang/srt/managers/cache_controller.py +``` +Expected: no output. + +- [ ] **Step 3: Sync code to remote only after local CPU tests are clean** + +Use `scp` to `g0034:/mnt/beegfs/cjy/sglang-dev`. Do not chown inside container. + +- [ ] **Step 4: Run remote targeted tests in container** + +Run via ssh/docker: +```bash +cd /sgl-workspace/sglang-tai +PYTHONPATH=python python -m pytest -q \ + test/registered/unit/mem_cache/test_cp_shared_kv_layout.py \ + test/registered/unit/mem_cache/test_cp_hicache_metadata.py \ + test/registered/unit/managers/test_hicache_controller_cp.py +``` +Expected: PASS. + +- [ ] **Step 5: Runtime validation knobs** + +Start conservatively: +```bash +--hicache-host-free-room-trigger-ratio 0.0 \ +--hicache-host-free-room-ratio 0.03 \ +--hicache-l1-free-room-trigger-ratio 0.0 \ +--hicache-l1-free-room-ratio 0.02 +``` +Observe eviction frequency, allocation fallback logs, cache hit, accept length, and Mooncake transfer block statistics. + +--- + +## Self-review + +- The plan keeps default behavior unchanged with all ratios `0.0`. +- The plan separates trigger and target room to avoid wasting cache capacity when current allocation fits. +- The plan avoids new collectives. +- The plan preserves CP owner-lane correctness. +- The plan records that over-eviction alone does not guarantee contiguous allocation; allocator policy is required. +- Risk retained: L1 contiguous-preferred selection must avoid CUDA host sync. If a correct no-sync implementation is too invasive, Task 6 should stop at tests and defer implementation rather than adding a slow hot-path algorithm. + +--- + +## Microbenchmark follow-up: allocator overhead guardrail + +### Added benchmark + +File: +- `benchmark/hicache/bench_cp_hicache_allocator_overhead.py` +- unit smoke: `test/registered/unit/benchmark/test_cp_hicache_allocator_bench.py` + +Purpose: +- measure metadata/control-path cost without allocating real 220GB HiCache buffers; +- compare FIFO allocation vs contiguous-preferred allocation; +- cover Host/L2 and L1 CP owner-lane allocation; +- keep CUDA validation remote-only. + +Important default behavior: +- `--host-sizes-gb 220` maps to page count via `floor(size_gb * 1e9 / (bytes_per_token * page_size))`. +- With default `bytes_per_token=100000` and `page_size=64`, 220GB maps to 34375 pages. +- `--host-pages` defaults to `8192,16384,32768` only when neither `--host-pages` nor `--host-sizes-gb` is supplied. + +### Remote verification already run + +Remote container path: +- `/sgl-workspace/sglang-tai` + +Commands: +```bash +python -m py_compile \ + benchmark/hicache/bench_cp_hicache_allocator_overhead.py \ + test/registered/unit/benchmark/test_cp_hicache_allocator_bench.py +PYTHONPATH=. python -m pytest -q \ + test/registered/unit/benchmark/test_cp_hicache_allocator_bench.py +``` +Result: +- `4 passed, 1 warning` + +Host 220GB-equivalent smoke: +```bash +PYTHONPATH=. python benchmark/hicache/bench_cp_hicache_allocator_overhead.py \ + --bench host --host-sizes-gb 220 \ + --request-pages 1,8,64,512 --repeat 5 --warmup 2 \ + --patterns contiguous_fifo,fragmented_prefix_later_run,random_fragmented \ + --host-methods fifo,contiguous \ + --json-output /tmp/cp_hicache_allocator_host_220.json +``` +Representative result: +- contiguous free list: contiguous-preferred remains cheap (`req=8 p50≈22us`, `req=512 p50≈86us`). +- fragmented/random free list: contiguous-preferred is multi-ms (`req=8/64 p50≈4.6ms`, `req=512 p50≈3.5ms`, random p95 can exceed 8ms). +- FIFO stays low-latency but loses contiguity in fragmented cases. + +L1 CPU proxy smoke: +```bash +PYTHONPATH=python:. python benchmark/hicache/bench_cp_hicache_allocator_overhead.py \ + --bench l1 --device cpu --stub-sgl-kernel --physical-pages 8192 \ + --request-pages 64,512 --repeat 5 --warmup 2 \ + --l1-free-patterns owner_fragmented_later_run,random \ + --l1-owner-patterns round_robin,single_owner,zigzag \ + --l1-impl fifo,current \ + --json-output /tmp/cp_hicache_allocator_l1_cpu.json +``` +Representative result: +- current contiguous-preferred L1 CPU proxy is slower than FIFO by ~0.4ms for single-owner and ~2.5-3.6ms for round-robin/zigzag at 8192 physical pages. +- it restores contiguity for `owner_fragmented_later_run`, but cannot restore contiguity for fully random free lists. + +CUDA note: +- CUDA L1 smoke was not run in this pass because a production prefill server was active on `g0034` using the GPUs. Do not run CUDA allocator benchmark concurrently with active ETE unless explicitly scheduled. + +### Finding / next implementation implication + +The current contiguous-preferred scan is acceptable only when free lists are already mostly contiguous. Under fragmented 220GB-equivalent metadata it can add multi-ms CPU/control-path overhead. This means the free-room/evict strategy cannot rely on repeated full-list scans to recover contiguity in the hot path. + +Next preferred direction: +1. Keep this benchmark as regression evidence. +2. Optimize allocator metadata before enabling aggressive contiguous-preferred allocation broadly: + - Host/L2: maintain page-run or interval metadata instead of scanning/sorting the whole free-slot tensor on each allocation. + - L1: avoid per-owner full free-list masks for small requests, or gate contiguous-preferred to fragmented-risk / large-request paths. +3. Use trigger/target free-room eviction to reduce eviction frequency, but do not assume it alone fixes allocator overhead. diff --git a/python/sglang/srt/disaggregation/common/utils.py b/python/sglang/srt/disaggregation/common/utils.py index 6f3da2128..aebf727cf 100644 --- a/python/sglang/srt/disaggregation/common/utils.py +++ b/python/sglang/srt/disaggregation/common/utils.py @@ -40,3 +40,36 @@ def group_concurrent_contiguous( dst_groups = [g.tolist() for g in dst_groups] return src_groups, dst_groups + + +def contiguous_group_stats( + src_indices: npt.NDArray[np.int32], + dst_indices: npt.NDArray[np.int32], + src_groups: List[npt.NDArray[np.int32]], + dst_groups: List[npt.NDArray[np.int32]], +) -> dict[str, object]: + """Summarize contiguous-group coalescing without exposing full index arrays.""" + + del dst_groups # Same grouping shape as src_groups; keep arg for call-site clarity. + + src_indices = np.asarray(src_indices).reshape(-1) + dst_indices = np.asarray(dst_indices).reshape(-1) + group_lens = [len(group) for group in src_groups] + page_count = int(src_indices.size) + group_count = int(len(group_lens)) + + def _diff_head(indices: npt.NDArray[np.int32], limit: int = 16) -> list[int]: + if indices.size <= 1: + return [] + sample = indices[: min(indices.size, limit + 1)] + return np.diff(sample).astype(np.int64).tolist() + + return { + "pages": page_count, + "groups": group_count, + "min_group_pages": int(min(group_lens)) if group_lens else 0, + "max_group_pages": int(max(group_lens)) if group_lens else 0, + "avg_group_pages": float(page_count / group_count) if group_count else 0.0, + "src_diff_head": _diff_head(src_indices), + "dst_diff_head": _diff_head(dst_indices), + } diff --git a/python/sglang/srt/disaggregation/mooncake/conn.py b/python/sglang/srt/disaggregation/mooncake/conn.py index 0875704c0..0a531106d 100644 --- a/python/sglang/srt/disaggregation/mooncake/conn.py +++ b/python/sglang/srt/disaggregation/mooncake/conn.py @@ -23,6 +23,7 @@ from sglang.srt.disaggregation.common.conn import ( ) from sglang.srt.disaggregation.common.utils import ( FastQueue, + contiguous_group_stats, group_concurrent_contiguous, ) from sglang.srt.disaggregation.mooncake.utils import ( @@ -62,6 +63,24 @@ def _cp_draft_shared_kv_debug(message: str, *args, limit: int = 64) -> None: logger.info("[CP_DRAFT_SHARED_KV] " + message, *args) +def _mooncake_transfer_stats_enabled() -> bool: + return envs.SGLANG_DISAGGREGATION_TRANSFER_STATS.get() + + +def _mooncake_transfer_stats_log(message: str, *args) -> None: + if not _mooncake_transfer_stats_enabled(): + return + limit = envs.SGLANG_DISAGGREGATION_TRANSFER_STATS_LIMIT.get() + if limit is not None and limit <= 0: + return + key = "mooncake_transfer_stats" + count = _CP_SHARED_DEBUG_COUNTS.get(key, 0) + if limit is not None and count >= limit: + return + _CP_SHARED_DEBUG_COUNTS[key] = count + 1 + logger.info("[Mooncake-transfer-stats] " + message, *args) + + def _np_summary(arr) -> str: if arr is None: return "None" @@ -310,6 +329,7 @@ class MooncakeKVManager(CommonKVManager): prefill_data_indices: npt.NDArray[np.int32], dst_data_indices: npt.NDArray[np.int32], executor: concurrent.futures.ThreadPoolExecutor, + debug_room: Optional[int] = None, ) -> int: """ Generic KV cache transfer supporting both MHA and MLA architectures. @@ -319,6 +339,15 @@ class MooncakeKVManager(CommonKVManager): prefill_kv_blocks, dst_kv_blocks = group_concurrent_contiguous( prefill_data_indices, dst_data_indices ) + transfer_stats_enabled = _mooncake_transfer_stats_enabled() + grouping_stats = None + if transfer_stats_enabled: + grouping_stats = contiguous_group_stats( + prefill_data_indices, + dst_data_indices, + prefill_kv_blocks, + dst_kv_blocks, + ) layers_params = None @@ -387,6 +416,7 @@ class MooncakeKVManager(CommonKVManager): transfer_blocks.extend(set_transfer_blocks(src_ptr, dst_ptr, item_len)) return self._transfer_data(mooncake_session_id, transfer_blocks) + start_time = time.perf_counter() if transfer_stats_enabled else 0.0 if self.enable_custom_mem_pool: futures = [ executor.submit( @@ -402,12 +432,84 @@ class MooncakeKVManager(CommonKVManager): if status != 0: for f in futures: f.cancel() + if transfer_stats_enabled: + self._log_kvcache_transfer_stats( + mooncake_session_id=mooncake_session_id, + debug_room=debug_room, + grouping_stats=grouping_stats, + layers_params=layers_params, + elapsed_ms=(time.perf_counter() - start_time) * 1000, + status=status, + custom_mem_pool=True, + ) return status - return 0 + status = 0 else: # Combining all layers' params in one batch transfer is more efficient # compared to using multiple threads - return process_layers(layers_params) + status = process_layers(layers_params) + + if transfer_stats_enabled: + self._log_kvcache_transfer_stats( + mooncake_session_id=mooncake_session_id, + debug_room=debug_room, + grouping_stats=grouping_stats, + layers_params=layers_params, + elapsed_ms=(time.perf_counter() - start_time) * 1000, + status=status, + custom_mem_pool=self.enable_custom_mem_pool, + ) + return status + + def _log_kvcache_transfer_stats( + self, + mooncake_session_id: str, + debug_room: Optional[int], + grouping_stats: Optional[dict[str, object]], + layers_params: List[Tuple[int, int, int]], + elapsed_ms: float, + status: int, + custom_mem_pool: bool, + ) -> None: + if grouping_stats is None: + return + page_count = int(grouping_stats["pages"]) + group_count = int(grouping_stats["groups"]) + layer_count = len(layers_params) + item_bytes_per_page = sum(int(item_len) for _, _, item_len in layers_params) + total_bytes = page_count * item_bytes_per_page + transfer_blocks = group_count * layer_count + avg_block_bytes = ( + float(total_bytes / transfer_blocks) if transfer_blocks else 0.0 + ) + bandwidth_gbps = ( + float(total_bytes / elapsed_ms / 1e6) if elapsed_ms > 0 else 0.0 + ) + + _mooncake_transfer_stats_log( + "cp_rank=%s room=%s session=%s status=%s custom_mem_pool=%s " + "pages=%s groups=%s avg_group_pages=%.2f max_group_pages=%s " + "layers=%s transfer_blocks=%s total_bytes=%.3fGiB " + "avg_block_bytes=%.1fKiB elapsed_ms=%.3f bandwidth=%.2fGB/s " + "src_diff_head=%s dst_diff_head=%s", + self.attn_cp_rank, + debug_room, + mooncake_session_id, + status, + custom_mem_pool, + page_count, + group_count, + float(grouping_stats["avg_group_pages"]), + grouping_stats["max_group_pages"], + layer_count, + transfer_blocks, + total_bytes / (1024**3), + avg_block_bytes / 1024, + elapsed_ms, + bandwidth_gbps, + grouping_stats["src_diff_head"], + grouping_stats["dst_diff_head"], + ) def send_kvcache( self, @@ -416,6 +518,7 @@ class MooncakeKVManager(CommonKVManager): dst_kv_ptrs: list[int], dst_kv_indices: npt.NDArray[np.int32], executor: concurrent.futures.ThreadPoolExecutor, + debug_room: Optional[int] = None, ): return self._send_kvcache_generic( mooncake_session_id=mooncake_session_id, @@ -425,6 +528,7 @@ class MooncakeKVManager(CommonKVManager): prefill_data_indices=prefill_kv_indices, dst_data_indices=dst_kv_indices, executor=executor, + debug_room=debug_room, ) def send_kvcache_slice( @@ -973,6 +1077,7 @@ class MooncakeKVManager(CommonKVManager): target_rank_registration_info.dst_kv_ptrs, chunked_dst_kv_indice, executor, + debug_room=kv_chunk.room, ) else: ret = self.send_kvcache_slice( diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 434afe091..6a2280404 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -261,6 +261,8 @@ class Envs: SGLANG_DISAGGREGATION_WAITING_TIMEOUT = EnvInt(300) SGLANG_DISAGGREGATION_NIXL_BACKEND = EnvStr("UCX") SGLANG_DISAGGREGATION_ALL_CP_RANKS_TRANSFER = EnvBool(False) + SGLANG_DISAGGREGATION_TRANSFER_STATS = EnvBool(False) + SGLANG_DISAGGREGATION_TRANSFER_STATS_LIMIT = EnvInt(128) # Scheduler: others: SGLANG_EMPTY_CACHE_INTERVAL = EnvFloat(-1) # in seconds. Set if you observe high memory accumulation over a long serving period. diff --git a/python/sglang/srt/managers/cache_controller.py b/python/sglang/srt/managers/cache_controller.py index 1770ef49c..f238531ae 100644 --- a/python/sglang/srt/managers/cache_controller.py +++ b/python/sglang/srt/managers/cache_controller.py @@ -980,7 +980,12 @@ class HiCacheController: owned_logical_indices = padded_device_indices[owned_mask] physical_device_indices = layout.logical_locs_to_physical(owned_logical_indices) - host_indices = self.mem_pool_host.alloc(len(physical_device_indices)) + host_alloc = getattr( + self.mem_pool_host, + "alloc_contiguous_preferred", + self.mem_pool_host.alloc, + ) + host_indices = host_alloc(len(physical_device_indices)) if host_indices is None: logger.info( "[CacheCtrl-write] reserve_write_cp FAILED (host full): node_id=%d logical_len=%d owned=%d", @@ -992,9 +997,12 @@ class HiCacheController: draft_host_indices = None if self.has_draft_hicache: - draft_host_indices = self.draft_mem_pool_host.alloc( - len(physical_device_indices) + draft_host_alloc = getattr( + self.draft_mem_pool_host, + "alloc_contiguous_preferred", + self.draft_mem_pool_host.alloc, ) + draft_host_indices = draft_host_alloc(len(physical_device_indices)) if draft_host_indices is None: self.mem_pool_host.free(host_indices) logger.info( diff --git a/python/sglang/srt/mem_cache/allocator.py b/python/sglang/srt/mem_cache/allocator.py index 0b5a1349e..25e278df4 100644 --- a/python/sglang/srt/mem_cache/allocator.py +++ b/python/sglang/srt/mem_cache/allocator.py @@ -20,6 +20,7 @@ Page-aligned memory pool. """ import abc +import math from typing import TYPE_CHECKING, List, Optional import torch @@ -36,6 +37,33 @@ if TYPE_CHECKING: from sglang.srt.mem_cache.memory_pool import KVCache +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 + + def _debug_sort_nvtx_enabled() -> bool: return _SORT_NVTX_ENABLED @@ -632,6 +660,73 @@ class CPSharedPagedTokenToKVPoolAllocator(PagedTokenToKVPoolAllocator): ] 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]]: + """Return owner-lane deficits with reactive trigger/target free room. + + Units are pages, matching ``page_compute_owners`` and + ``owner_lane_deficits``. With both ratios at zero this is identical to + ``compute_owner_lane_stats``. + """ + + 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_mask: torch.Tensor, + required_count: int, + ) -> torch.Tensor: + selected_prefix_mask = owner_mask & ( + torch.cumsum(owner_mask.to(torch.int64), dim=0) <= required_count + ) + if required_count <= 1: + return selected_prefix_mask + + owner_positions = owner_mask.nonzero(as_tuple=True)[0] + if owner_positions.numel() < required_count: + return selected_prefix_mask + + owner_pages = self.free_pages[owner_positions] + run_edges = owner_pages[1:] - owner_pages[:-1] == self.cp_size + if run_edges.numel() < required_count - 1: + return selected_prefix_mask + + eligible_starts = run_edges.unfold(0, required_count - 1, 1).all(dim=1) + if eligible_starts.numel() == 0: + return selected_prefix_mask + + chosen_start = torch.argmax(eligible_starts.to(torch.int64)) + contiguous_positions = owner_positions[ + chosen_start + + torch.arange( + required_count, dtype=torch.int64, device=owner_positions.device + ) + ] + selected_contiguous_mask = torch.zeros_like(owner_mask) + selected_contiguous_mask[contiguous_positions] = True + return torch.where( + eligible_starts.any(), selected_contiguous_mask, selected_prefix_mask + ) + def _select_compute_owner_pages( self, page_compute_owners: List[int], @@ -659,8 +754,10 @@ class CPSharedPagedTokenToKVPoolAllocator(PagedTokenToKVPoolAllocator): continue owner_mask = torch.remainder(self.free_pages - 1, self.cp_size) == owner - selected_owner_free_mask = owner_mask & ( - torch.cumsum(owner_mask.to(torch.int64), dim=0) <= required_count + selected_owner_free_mask = ( + self._select_owner_free_pages_prefer_contiguous( + owner_mask, required_count + ) ) selected_owner_pages = self.free_pages[selected_owner_free_mask] diff --git a/python/sglang/srt/mem_cache/common.py b/python/sglang/srt/mem_cache/common.py index fa0cbe4a3..57d4a4726 100644 --- a/python/sglang/srt/mem_cache/common.py +++ b/python/sglang/srt/mem_cache/common.py @@ -12,6 +12,7 @@ from sglang.srt.mem_cache.base_prefix_cache import ( EvictParams, EvictResult, ) +from sglang.srt.mem_cache.allocator import compute_owner_lane_free_room_deficits from sglang.srt.mem_cache.cp_shared_kv_compute_owner import ( build_in_seq_page_compute_owners, get_in_seq_page_compute_owner_unavailable_reason, @@ -65,6 +66,42 @@ def _log_cp_shared_kv_alloc_fallback( ) +def _compute_owner_lane_stats_for_eviction( + *, + tree_cache: BasePrefixCache, + allocator, + page_compute_owners: list[int], +) -> tuple[list[int], list[int], list[int]]: + target_ratio = float(getattr(tree_cache, "hicache_l1_free_room_ratio", 0.0) or 0.0) + trigger_ratio = float( + getattr(tree_cache, "hicache_l1_free_room_trigger_ratio", 0.0) or 0.0 + ) + free_room_stats = getattr(allocator, "compute_owner_lane_free_room_stats", None) + if free_room_stats is not None: + return free_room_stats( + page_compute_owners, + target_ratio=target_ratio, + trigger_ratio=trigger_ratio, + ) + required, available, exact_deficits = allocator.compute_owner_lane_stats( + page_compute_owners + ) + lane_capacity_pages = getattr(allocator, "compute_owner_lane_capacity_pages", None) + if lane_capacity_pages is None: + return required, available, exact_deficits + return ( + required, + available, + compute_owner_lane_free_room_deficits( + required=required, + available=available, + capacities=lane_capacity_pages(), + target_ratio=target_ratio, + trigger_ratio=trigger_ratio, + ), + ) + + @triton.jit def write_req_to_token_pool_triton( req_to_token_ptr, # [max_batch, max_context_len] @@ -331,7 +368,11 @@ def _evict_for_compute_owner_lanes( max_attempts = max(2, min(8, int(getattr(allocator, "cp_size", 1)))) for attempt in range(max_attempts): - _required, _available, deficits = compute_owner_lane_stats(page_compute_owners) + _required, _available, deficits = _compute_owner_lane_stats_for_eviction( + tree_cache=tree_cache, + allocator=allocator, + page_compute_owners=page_compute_owners, + ) deficit_pages = sum(deficits) if deficit_pages <= 0: return diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index a683ccb40..58cf7187a 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/python/sglang/srt/mem_cache/hiradix_cache.py @@ -4,6 +4,7 @@ import atexit import heapq import json import logging +import math import os import threading import time @@ -20,7 +21,10 @@ from sglang.srt.managers.cache_controller import ( HiCacheWriteReservation, PrefetchOperation, ) -from sglang.srt.mem_cache.allocator import CPSharedPagedTokenToKVPoolAllocator +from sglang.srt.mem_cache.allocator import ( + CPSharedPagedTokenToKVPoolAllocator, + compute_owner_lane_free_room_deficits, +) from sglang.srt.mem_cache.base_prefix_cache import ( DecLockRefParams, DecLockRefResult, @@ -126,6 +130,31 @@ def _compute_shared_hicache_token_capacities( return target_tokens, draft_tokens +def _page_aligned_room_from_ratio(capacity: int, ratio: float, page_size: int) -> int: + if page_size <= 0: + raise ValueError(f"page_size must be positive, got {page_size}") + if capacity <= 0 or ratio <= 0: + return 0 + raw_room = int(math.ceil(float(capacity) * float(ratio))) + return ((raw_room + page_size - 1) // page_size) * page_size + + +def _free_room_deficit( + *, + required: int, + available: int, + capacity: int, + page_size: int, + target_ratio: float, + trigger_ratio: float, +) -> int: + target_room = _page_aligned_room_from_ratio(capacity, target_ratio, page_size) + trigger_room = _page_aligned_room_from_ratio(capacity, trigger_ratio, page_size) + if int(available) >= int(required) + trigger_room: + return 0 + return max(0, int(required) + target_room - int(available)) + + @dataclass class CpHiCacheNodeMetadata: # Legacy name kept for existing call sites. Under the page-aligned cache @@ -570,6 +599,14 @@ class HiRadixCache(RadixCache): 1 if server_args.hicache_write_policy == "write_through" else 2 ) self.load_back_threshold = 10 + self.hicache_host_free_room_ratio = server_args.hicache_host_free_room_ratio + self.hicache_host_free_room_trigger_ratio = ( + server_args.hicache_host_free_room_trigger_ratio + ) + self.hicache_l1_free_room_ratio = server_args.hicache_l1_free_room_ratio + self.hicache_l1_free_room_trigger_ratio = ( + server_args.hicache_l1_free_room_trigger_ratio + ) # Detach storage backend automatically on process shutdown atexit.register(self.shutdown) @@ -973,7 +1010,7 @@ class HiRadixCache(RadixCache): f"page_size={self.page_size} page_owners={len(page_owners)}" ) - required, available, deficits = compute_owner_lane_stats(page_owners) + required, available, deficits = self._cp_load_back_owner_lane_stats(page_owners) return CpLoadBackPlan( page_owners=list(page_owners), required_by_owner=[int(v) for v in required], @@ -982,9 +1019,42 @@ class HiRadixCache(RadixCache): host_hit_len=host_hit_len, ) + def _cp_load_back_owner_lane_stats( + self, page_owners: List[int] + ) -> Tuple[List[int], List[int], List[int]]: + allocator = self.token_to_kv_pool_allocator + target_ratio = float(getattr(self, "hicache_l1_free_room_ratio", 0.0) or 0.0) + trigger_ratio = float( + getattr(self, "hicache_l1_free_room_trigger_ratio", 0.0) or 0.0 + ) + free_room_stats = getattr(allocator, "compute_owner_lane_free_room_stats", None) + if free_room_stats is not None: + return free_room_stats( + page_owners, + target_ratio=target_ratio, + trigger_ratio=trigger_ratio, + ) + required, available, exact_deficits = allocator.compute_owner_lane_stats( + page_owners + ) + lane_capacity_pages = getattr(allocator, "compute_owner_lane_capacity_pages", None) + if lane_capacity_pages is None: + return required, available, exact_deficits + return ( + required, + available, + compute_owner_lane_free_room_deficits( + required=required, + available=available, + capacities=lane_capacity_pages(), + target_ratio=target_ratio, + trigger_ratio=trigger_ratio, + ), + ) + def _refresh_cp_load_back_plan(self, plan: CpLoadBackPlan) -> CpLoadBackPlan: - required, available, deficits = ( - self.token_to_kv_pool_allocator.compute_owner_lane_stats(plan.page_owners) + required, available, deficits = self._cp_load_back_owner_lane_stats( + plan.page_owners ) return CpLoadBackPlan( page_owners=plan.page_owners, @@ -1344,11 +1414,42 @@ class HiRadixCache(RadixCache): snapshot = self._cp_host_capacity_snapshot() target_available = self._cp_host_available_tokens_by_rank(snapshot) draft_available = self._cp_host_available_tokens_by_rank(snapshot, draft=True) - target_deficit = tuple( - max(0, req - avail) for req, avail in zip(required, target_available) + target_ratio = float(getattr(self, "hicache_host_free_room_ratio", 0.0) or 0.0) + trigger_ratio = float( + getattr(self, "hicache_host_free_room_trigger_ratio", 0.0) or 0.0 ) + target_deficit = tuple( + _free_room_deficit( + required=req, + available=avail, + capacity=capacity, + page_size=self.page_size, + target_ratio=target_ratio, + trigger_ratio=trigger_ratio, + ) + for req, avail, capacity in zip( + required, target_available, snapshot.target_capacity + ) + ) + draft_capacity = snapshot.draft_capacity draft_deficit = tuple( - max(0, req - avail) for req, avail in zip(required, draft_available) + ( + _free_room_deficit( + required=req, + available=avail, + capacity=capacity, + page_size=self.page_size, + target_ratio=target_ratio, + trigger_ratio=trigger_ratio, + ) + if draft_capacity is not None + else 0 + ) + for req, avail, capacity in zip( + required, + draft_available, + draft_capacity if draft_capacity is not None else snapshot.target_capacity, + ) ) deficit = tuple(max(a, b) for a, b in zip(target_deficit, draft_deficit)) eviction_plan = self._plan_cp_host_evictions(deficit) diff --git a/python/sglang/srt/mem_cache/memory_pool_host.py b/python/sglang/srt/mem_cache/memory_pool_host.py index d75c735df..b201ed9c6 100644 --- a/python/sglang/srt/mem_cache/memory_pool_host.py +++ b/python/sglang/srt/mem_cache/memory_pool_host.py @@ -338,6 +338,88 @@ class HostKVCache(abc.ABC): return select_index + @synchronized + def alloc_contiguous_preferred(self, need_size: int) -> Optional[torch.Tensor]: + """Allocate page-shaped host slots, preferring physical page runs. + + ``free_slots`` is still the source of truth and the fallback remains FIFO. + The CP HiCache fast paths only require best-effort contiguity: if a + contiguous physical run is readily available, use it to reduce transfer + descriptor fragmentation; otherwise preserve the historical allocator + behavior by returning ``alloc(need_size)``. + """ + + assert ( + need_size % self.page_size == 0 + ), "The requested size should be a multiple of the page size." + 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].T + ).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 + @synchronized def free(self, indices: torch.Tensor) -> int: self.free_slots = torch.cat([self.free_slots, indices.cpu()]) diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 5d580017b..7ebaaa61a 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -558,6 +558,10 @@ class ServerArgs: hicache_write_policy: str = "write_through" hicache_io_backend: str = "kernel" hicache_mem_layout: str = "layer_first" + hicache_host_free_room_ratio: float = 0.0 + hicache_host_free_room_trigger_ratio: float = 0.0 + hicache_l1_free_room_ratio: float = 0.0 + hicache_l1_free_room_trigger_ratio: float = 0.0 disable_hicache_numa_detect: bool = False hicache_storage_backend: Optional[str] = None hicache_storage_prefetch_policy: str = "best_effort" @@ -2898,6 +2902,7 @@ class ServerArgs: or self.disaggregation_decode_enable_offload_kvcache ): return + self._validate_hicache_free_room_args() # Step 1: Initial layout-io compatibility normalization. self._resolve_layout_io_compatibility() @@ -2912,6 +2917,34 @@ class ServerArgs: if io_changed: self._resolve_layout_io_compatibility() + def _validate_hicache_free_room_args(self): + def check_ratio(name: str, value: float): + if value < 0 or value >= 1: + raise ValueError(f"{name} must be in [0, 1), got {value}") + + check_ratio( + "hicache_host_free_room_ratio", self.hicache_host_free_room_ratio + ) + check_ratio( + "hicache_host_free_room_trigger_ratio", + self.hicache_host_free_room_trigger_ratio, + ) + check_ratio("hicache_l1_free_room_ratio", self.hicache_l1_free_room_ratio) + check_ratio( + "hicache_l1_free_room_trigger_ratio", + self.hicache_l1_free_room_trigger_ratio, + ) + if self.hicache_host_free_room_trigger_ratio > self.hicache_host_free_room_ratio: + raise ValueError( + "hicache_host_free_room_trigger_ratio must be <= " + "hicache_host_free_room_ratio" + ) + if self.hicache_l1_free_room_trigger_ratio > self.hicache_l1_free_room_ratio: + raise ValueError( + "hicache_l1_free_room_trigger_ratio must be <= " + "hicache_l1_free_room_ratio" + ) + def _resolve_layout_io_compatibility(self): if ( self.hicache_mem_layout == "page_first_direct" @@ -5136,6 +5169,42 @@ class ServerArgs: default=ServerArgs.hicache_mem_layout, help="The layout of host memory pool for hierarchical cache.", ) + parser.add_argument( + "--hicache-host-free-room-ratio", + type=float, + default=ServerArgs.hicache_host_free_room_ratio, + help=( + "L2/host free-room target ratio after a triggered CP HiCache " + "host eviction. 0 preserves exact-deficit eviction." + ), + ) + parser.add_argument( + "--hicache-host-free-room-trigger-ratio", + type=float, + default=ServerArgs.hicache_host_free_room_trigger_ratio, + help=( + "L2/host free-room trigger ratio. Eviction starts only when " + "available host slots are below required plus this room." + ), + ) + parser.add_argument( + "--hicache-l1-free-room-ratio", + type=float, + default=ServerArgs.hicache_l1_free_room_ratio, + help=( + "L1/device owner-lane free-room target ratio after a triggered " + "CP HiCache owner-lane eviction." + ), + ) + parser.add_argument( + "--hicache-l1-free-room-trigger-ratio", + type=float, + default=ServerArgs.hicache_l1_free_room_trigger_ratio, + help=( + "L1/device owner-lane free-room trigger ratio. Eviction starts " + "only when owner-lane pages are below required plus this room." + ), + ) parser.add_argument( "--disable-hicache-numa-detect", action="store_true", diff --git a/test/registered/unit/benchmark/test_cp_hicache_allocator_bench.py b/test/registered/unit/benchmark/test_cp_hicache_allocator_bench.py new file mode 100644 index 000000000..2b8b1ef56 --- /dev/null +++ b/test/registered/unit/benchmark/test_cp_hicache_allocator_bench.py @@ -0,0 +1,47 @@ +import torch + +from benchmark.hicache.bench_cp_hicache_allocator_overhead import ( + StandaloneHostAllocator, + _host_pages_from_gb, + _make_host_free_slots, + _parse_int_list, +) + + +def test_parse_int_list_accepts_commas_and_spaces(): + assert _parse_int_list("1, 2,8") == [1, 2, 8] + + +def test_host_pages_from_gb_rounds_to_page_capacity(): + assert _host_pages_from_gb(220.0, bytes_per_token=100_000, page_size=64) == 34375 + + +def test_host_contiguous_preferred_skips_fragmented_prefix(): + page_size = 4 + free_slots = _make_host_free_slots( + total_pages=12, + request_pages=2, + page_size=page_size, + pattern="fragmented_prefix_later_run", + seed=0, + ) + allocator = StandaloneHostAllocator(page_size=page_size, free_slots=free_slots) + + selected = allocator.alloc_contiguous_preferred(2 * page_size) + + selected_pages = (selected.view(-1, page_size)[:, 0] // page_size).tolist() + assert selected_pages[1] == selected_pages[0] + 1 + prefix_pages = (free_slots[: 2 * page_size].view(-1, page_size)[:, 0] // page_size) + assert selected_pages != prefix_pages.tolist() + + +def test_host_random_fragmented_has_requested_size(): + free_slots = _make_host_free_slots( + total_pages=64, + request_pages=8, + page_size=16, + pattern="random_fragmented", + seed=123, + ) + assert free_slots.numel() == 64 * 16 + assert torch.unique(free_slots).numel() == free_slots.numel() diff --git a/test/registered/unit/disaggregation/test_disaggregation_common_utils.py b/test/registered/unit/disaggregation/test_disaggregation_common_utils.py new file mode 100644 index 000000000..f8941a314 --- /dev/null +++ b/test/registered/unit/disaggregation/test_disaggregation_common_utils.py @@ -0,0 +1,63 @@ +import unittest +from importlib import util +from pathlib import Path + +import numpy as np + +from sglang.test.ci.ci_register import register_cpu_ci + + +register_cpu_ci(est_time=1, suite="stage-a-test-cpu") + +_UTILS_PATH = ( + Path(__file__).resolve().parents[4] + / "python" + / "sglang" + / "srt" + / "disaggregation" + / "common" + / "utils.py" +) +_spec = util.spec_from_file_location("disaggregation_common_utils_under_test", _UTILS_PATH) +_utils = util.module_from_spec(_spec) +assert _spec.loader is not None +_spec.loader.exec_module(_utils) +contiguous_group_stats = _utils.contiguous_group_stats + + +class TestDisaggregationCommonUtils(unittest.TestCase): + def test_contiguous_group_stats_reports_fragmentation_shape(self): + src_indices = np.array([1, 2, 3, 10, 18, 19], dtype=np.int32) + dst_indices = np.array([101, 102, 103, 110, 118, 119], dtype=np.int32) + src_groups = [[1, 2, 3], [10], [18, 19]] + dst_groups = [[101, 102, 103], [110], [118, 119]] + + stats = contiguous_group_stats( + src_indices, dst_indices, src_groups, dst_groups + ) + + self.assertEqual(stats["pages"], 6) + self.assertEqual(stats["groups"], 3) + self.assertEqual(stats["min_group_pages"], 1) + self.assertEqual(stats["max_group_pages"], 3) + self.assertAlmostEqual(stats["avg_group_pages"], 2.0) + self.assertEqual(stats["src_diff_head"], [1, 1, 7, 8, 1]) + self.assertEqual(stats["dst_diff_head"], [1, 1, 7, 8, 1]) + + def test_contiguous_group_stats_handles_empty_transfer(self): + stats = contiguous_group_stats( + np.array([], dtype=np.int32), + np.array([], dtype=np.int32), + [], + [], + ) + + self.assertEqual(stats["pages"], 0) + self.assertEqual(stats["groups"], 0) + self.assertEqual(stats["avg_group_pages"], 0.0) + self.assertEqual(stats["src_diff_head"], []) + self.assertEqual(stats["dst_diff_head"], []) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/managers/test_hicache_controller_cp.py b/test/registered/unit/managers/test_hicache_controller_cp.py index 31232bf56..e0c5dd96c 100644 --- a/test/registered/unit/managers/test_hicache_controller_cp.py +++ b/test/registered/unit/managers/test_hicache_controller_cp.py @@ -109,6 +109,7 @@ from sglang.srt.managers.cache_controller import HiCacheController from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout from sglang.srt.mem_cache.hiradix_cache import CpHiCacheNodeMetadata from sglang.srt.mem_cache.memory_pool_host import ( + HostKVCache, MHATokenToKVPoolHost, MLATokenToKVPoolHost, NSATokenToKVPoolHost, @@ -161,6 +162,50 @@ class FakeHostPool: return len(indices) +class ContiguousPreferredHostPool(FakeHostPool): + def __init__(self, alloc_result): + super().__init__(alloc_result) + self.contiguous_alloc_calls = [] + + def alloc_contiguous_preferred(self, need_size): + self.contiguous_alloc_calls.append(need_size) + if self.alloc_result is None: + return None + return self.alloc_result[:need_size].clone() + + +class DummyHostKVCacheForAlloc(HostKVCache): + def get_size_per_token(self): + return 1 + + def init_kv_buffer(self): + return None + + def load_to_device_per_layer( + self, device_pool, host_indices, device_indices, layer_id, io_backend + ) -> None: + pass + + def backup_from_device_per_layer( + self, device_pool, host_indices, device_indices, layer_id, io_backend + ) -> None: + pass + + def backup_from_device_all_layer( + self, device_pool, host_indices, device_indices, io_backend + ) -> None: + pass + + def get_data_page(self, index, flat: bool = True) -> 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, index: int, data_page: torch.Tensor) -> None: + pass + + class FakeDevicePool: device = "cpu" layer_num = 1 @@ -931,6 +976,46 @@ class TestHiCacheControllerCPWrite(CustomTestCase): self.assertEqual(controller.ack_write_queue[0].node_ids, [79]) self.assertIn("all-layer backup fallback", "\n".join(logs.output)) + def test_cp_reserve_write_uses_contiguous_preferred_host_alloc(self): + host_pool = ContiguousPreferredHostPool( + torch.tensor([100, 101, 102, 103], dtype=torch.int64) + ) + draft_host_pool = ContiguousPreferredHostPool( + torch.tensor([200, 201, 202, 203], dtype=torch.int64) + ) + controller = self.make_controller( + host_pool, + cp_rank=1, + draft_host_pool=draft_host_pool, + draft_mem_pool_device=FakeDevicePool("draft"), + ) + logical_locs = torch.arange(4, 20, dtype=torch.int64) + + reservation = controller.reserve_write_cp(logical_locs, node_id=179) + + self.assertEqual(reservation.metadata.host_indices.tolist(), [100, 101, 102, 103]) + self.assertEqual( + reservation.metadata.draft_host_indices.tolist(), [200, 201, 202, 203] + ) + self.assertEqual(host_pool.contiguous_alloc_calls, [4]) + self.assertEqual(draft_host_pool.contiguous_alloc_calls, [4]) + self.assertEqual(host_pool.alloc_calls, []) + self.assertEqual(draft_host_pool.alloc_calls, []) + + def test_host_alloc_contiguous_preferred_skips_fragmented_fifo_prefix(self): + host_pool = DummyHostKVCacheForAlloc.__new__(DummyHostKVCacheForAlloc) + host_pool.page_size = 4 + host_pool.lock = __import__("threading").RLock() + host_pool.free_slots = torch.tensor( + [100, 101, 102, 103, 8, 9, 10, 11, 12, 13, 14, 15], + dtype=torch.int64, + ) + + selected = host_pool.alloc_contiguous_preferred(8) + + 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_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) diff --git a/test/registered/unit/mem_cache/test_cp_hicache_metadata.py b/test/registered/unit/mem_cache/test_cp_hicache_metadata.py index 3b5e5062d..ab2adf7ff 100644 --- a/test/registered/unit/mem_cache/test_cp_hicache_metadata.py +++ b/test/registered/unit/mem_cache/test_cp_hicache_metadata.py @@ -107,6 +107,8 @@ from sglang.srt.mem_cache.base_prefix_cache import ( MatchPrefixParams, ) from sglang.srt.mem_cache.hiradix_cache import ( + CpHiCacheCapacitySnapshot, + CpHiCacheEvictionPlan, CpHiCacheNodeMetadata, HiRadixCache, PendingHiCacheBackup, @@ -605,6 +607,132 @@ class TestCpHiCacheNodeMetadata(CustomTestCase): self.assertEqual(child.page_owners.numel(), 0) +class TestCpHiCacheFreeRoom(CustomTestCase): + def test_free_room_deficit_does_not_evict_when_required_fits_trigger(self): + deficit = hiradix_cache._free_room_deficit( + required=64, + available=96, + capacity=256, + page_size=64, + target_ratio=0.5, + trigger_ratio=0.0, + ) + + self.assertEqual(deficit, 0) + + def test_free_room_deficit_evicts_to_target_after_trigger(self): + deficit = hiradix_cache._free_room_deficit( + required=64, + available=96, + capacity=256, + page_size=64, + target_ratio=0.5, + trigger_ratio=0.25, + ) + + self.assertEqual(deficit, 96) + + def test_free_room_deficit_rounds_room_to_page(self): + deficit = hiradix_cache._free_room_deficit( + required=64, + available=0, + capacity=100, + page_size=64, + target_ratio=0.01, + trigger_ratio=0.0, + ) + + self.assertEqual(deficit, 128) + + def test_cp_host_write_admission_uses_trigger_target_room(self): + cache = HiRadixCache.__new__(HiRadixCache) + cache.page_size = 64 + cache.hicache_host_free_room_ratio = 0.5 + cache.hicache_host_free_room_trigger_ratio = 0.25 + cache._cp_required_host_tokens_by_rank = lambda _indices: (64, 0) + cache._cp_host_capacity_snapshot = lambda: CpHiCacheCapacitySnapshot( + target_capacity=(256, 256), + draft_capacity=None, + committed_target=(160, 0), + committed_draft=(0, 0), + pending_target=(0, 0), + pending_draft=(0, 0), + ) + cache._plan_cp_host_evictions = lambda deficit: CpHiCacheEvictionPlan( + victims=(), + planned_freed=tuple(0 for _ in deficit), + remaining_deficit=tuple(0 for _ in deficit), + ) + + admission = cache._cp_build_write_admission( + torch.arange(64, dtype=torch.int64), + node_id=600, + phase="unit", + ) + + self.assertEqual(admission.target_available_by_owner, (96, 256)) + self.assertEqual(admission.deficit_by_owner, (96, 0)) + + def test_cp_host_write_admission_does_not_evict_when_trigger_room_fits(self): + cache = HiRadixCache.__new__(HiRadixCache) + cache.page_size = 64 + cache.hicache_host_free_room_ratio = 0.5 + cache.hicache_host_free_room_trigger_ratio = 0.0 + cache._cp_required_host_tokens_by_rank = lambda _indices: (64, 0) + cache._cp_host_capacity_snapshot = lambda: CpHiCacheCapacitySnapshot( + target_capacity=(256, 256), + draft_capacity=None, + committed_target=(160, 0), + committed_draft=(0, 0), + pending_target=(0, 0), + pending_draft=(0, 0), + ) + cache._plan_cp_host_evictions = lambda deficit: CpHiCacheEvictionPlan( + victims=(), + planned_freed=tuple(0 for _ in deficit), + remaining_deficit=tuple(0 for _ in deficit), + ) + + admission = cache._cp_build_write_admission( + torch.arange(64, dtype=torch.int64), + node_id=601, + phase="unit", + ) + + self.assertEqual(admission.target_available_by_owner, (96, 256)) + self.assertEqual(admission.deficit_by_owner, (0, 0)) + + def test_cp_host_write_admission_uses_draft_room_deficit(self): + cache = HiRadixCache.__new__(HiRadixCache) + cache.page_size = 64 + cache.hicache_host_free_room_ratio = 0.25 + cache.hicache_host_free_room_trigger_ratio = 0.25 + cache._cp_required_host_tokens_by_rank = lambda _indices: (64, 0) + cache._cp_host_capacity_snapshot = lambda: CpHiCacheCapacitySnapshot( + target_capacity=(256, 256), + draft_capacity=(256, 256), + committed_target=(0, 0), + committed_draft=(192, 0), + pending_target=(0, 0), + pending_draft=(0, 0), + ) + cache._plan_cp_host_evictions = lambda deficit: CpHiCacheEvictionPlan( + victims=(), + planned_freed=tuple(0 for _ in deficit), + remaining_deficit=tuple(0 for _ in deficit), + ) + + admission = cache._cp_build_write_admission( + torch.arange(64, dtype=torch.int64), + node_id=602, + phase="unit", + ) + + self.assertEqual(admission.target_available_by_owner, (256, 256)) + self.assertEqual(admission.draft_available_by_owner, (64, 256)) + self.assertEqual(admission.deficit_by_owner, (64, 0)) + + class FakeWriteFailure: metadata = None @@ -2592,6 +2720,39 @@ class TestHiRadixCacheCPSplitEvict(CustomTestCase): class TestHiRadixCacheCPLoadBack(CustomTestCase): + def test_cp_load_back_plan_uses_l1_free_room_target(self): + class FreeRoomAllocator: + def compute_owner_lane_stats(self, _page_owners): + return [1, 0], [0, 8], [1, 0] + + def compute_owner_lane_capacity_pages(self): + return [8, 8] + + cache = HiRadixCache.__new__(HiRadixCache) + cache._uses_cp_hicache = True + cache.page_size = 64 + cache.hicache_l1_free_room_ratio = 0.5 + cache.hicache_l1_free_room_trigger_ratio = 0.25 + cache.token_to_kv_pool_allocator = FreeRoomAllocator() + + node = TreeNode(id=701) + node.host_len = 64 + node.cp_hicache = CpHiCacheNodeMetadata( + logical_len=64, + padded_len=64, + owned_positions=torch.arange(64, dtype=torch.int64), + host_indices=torch.arange(64, dtype=torch.int64), + page_owners=torch.tensor([0], dtype=torch.int8), + page_size=64, + ) + + plan = cache._build_cp_load_back_plan([node], node_id=701) + + self.assertEqual(plan.required_by_owner, [1, 0]) + self.assertEqual(plan.available_by_owner, [0, 8]) + # required=1 page, available=0, target_room=ceil(8*0.5)=4 pages. + self.assertEqual(plan.deficit_by_owner, [5, 0]) + def test_cp_load_back_uses_host_len_not_host_value(self): cache = HiRadixCache.__new__(HiRadixCache) cache._uses_cp_hicache = True diff --git a/test/registered/unit/mem_cache/test_cp_shared_kv_layout.py b/test/registered/unit/mem_cache/test_cp_shared_kv_layout.py index 496774ce4..8757ef4a7 100644 --- a/test/registered/unit/mem_cache/test_cp_shared_kv_layout.py +++ b/test/registered/unit/mem_cache/test_cp_shared_kv_layout.py @@ -341,7 +341,7 @@ class TestCPSharedPagedAllocator(CustomTestCase): self.assertIsNotNone(locs) logical_pages = locs.view(-1, page_size)[:, 0] // page_size - self.assertEqual(logical_pages.tolist(), [9, 3, 1, 7, 4]) + self.assertEqual(logical_pages.tolist(), [1, 3, 5, 7, 4]) self.assertEqual( ((logical_pages - 1) % cp_size).tolist(), page_compute_owners, @@ -392,6 +392,40 @@ class TestCPSharedPagedAllocator(CustomTestCase): self.assertEqual(allocator.free_pages.tolist(), [2, 3, 4]) self.assertEqual(allocator.release_pages.tolist(), []) + def test_contiguous_owner_lane_selection_prefers_later_physical_run(self): + from sglang.srt.mem_cache.allocator import CPSharedPagedTokenToKVPoolAllocator + + page_size = 64 + cp_size = 4 + allocator = CPSharedPagedTokenToKVPoolAllocator( + logical_size=page_size * 32, + physical_size=page_size * 8, + page_size=page_size, + dtype=torch.bfloat16, + device="cpu", + kvcache=None, + need_sort=False, + cp_size=cp_size, + cp_rank=0, + ) + allocator.free_pages = torch.tensor( + [1, 9, 13, 17] + + [page for page in range(2, 33) if page not in {9, 13, 17}], + dtype=torch.int64, + ) + + locs = allocator.alloc_pages_with_owners([0, 0, 0]) + + self.assertIsNotNone(locs) + logical_pages = locs.view(-1, page_size)[:, 0] // page_size + self.assertEqual(logical_pages.tolist(), [9, 13, 17]) + self.assertEqual( + ((logical_pages - 1) % cp_size).tolist(), + [0, 0, 0], + ) + for selected_page in logical_pages.tolist(): + self.assertNotIn(selected_page, allocator.free_pages.tolist()) + def test_compute_owner_alloc_does_not_evict_lanes_when_first_try_succeeds(self): from types import SimpleNamespace @@ -663,6 +697,54 @@ class TestCPSharedPagedAllocator(CustomTestCase): ) self.assertIsNotNone(locs) + def test_compute_owner_lane_eviction_uses_l1_free_room_target(self): + from sglang.srt.mem_cache.base_prefix_cache import EvictResult + from sglang.srt.mem_cache.common import _evict_for_compute_owner_lanes + + page_size = 64 + + class FakeAllocator: + cp_size = 2 + + def __init__(self): + self.page_size = page_size + + def available_size(self): + return 0 + + def compute_owner_lane_stats(self, _page_compute_owners): + return [1, 0], [0, 8], [1, 0] + + def compute_owner_lane_capacity_pages(self): + return [8, 8] + + class FakeTreeCache: + hicache_l1_free_room_ratio = 0.5 + hicache_l1_free_room_trigger_ratio = 0.25 + + def __init__(self): + self.owner_deficits = [] + + def is_chunk_cache(self): + return False + + def evictable_size(self): + return page_size * 8 + + def evict(self, params): + self.owner_deficits.append(list(params.owner_lane_deficits)) + return EvictResult(num_tokens_evicted=0) + + tree_cache = FakeTreeCache() + _evict_for_compute_owner_lanes( + tree_cache=tree_cache, + allocator=FakeAllocator(), + page_compute_owners=[0], + ) + + # required=1 page, available=0, target_room=ceil(8*0.5)=4 pages. + self.assertEqual(tree_cache.owner_deficits[0], [5, 0]) + def test_compute_owner_capacity_wait_reports_owner_lane_deficits(self): from types import SimpleNamespace