sglang/benchmark/hicache/bench_cp_hicache_allocator_overhead.py

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

  # L1 allocator path on CUDA, stubbing sgl_kernel import if needed.
  PYTHONPATH=python:. python benchmark/hicache/bench_cp_hicache_allocator_overhead.py \
    --bench l1 --device cuda --stub-sgl-kernel --physical-pages 8192,32768 \
    --request-pages 8,64,512 --l1-impl current,fifo \
    --l1-ops stats,free_room_stats,select_only,alloc_pages
"""

import argparse
import json
import math
import os
import 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 _compute_owner_lane_free_room_deficits(
    *,
    required: list[int],
    available: list[int],
    capacities: list[int],
    target_ratio: float,
    trigger_ratio: float,
) -> list[int]:
    deficits: list[int] = []
    for req, avail, capacity in zip(required, available, capacities):
        target_room = (
            int(math.ceil(float(capacity) * float(target_ratio)))
            if capacity > 0 and target_ratio > 0
            else 0
        )
        trigger_room = (
            int(math.ceil(float(capacity) * float(trigger_ratio)))
            if capacity > 0 and trigger_ratio > 0
            else 0
        )
        if int(avail) >= int(req) + trigger_room:
            deficits.append(0)
        else:
            deficits.append(max(0, int(req) + target_room - int(avail)))
    return deficits


class StandaloneCPSharedPagedAllocator:
    """Metadata-only copy of the CP shared-KV page owner allocator.

    This intentionally mirrors the current Python/Torch control path used by
    ``CPSharedPagedTokenToKVPoolAllocator`` so CPU-only environments can measure
    the allocator shape without importing the full SGLang runtime dependency
    stack.  The benchmark still uses the production allocator when imports are
    available.
    """

    def __init__(
        self,
        *,
        physical_pages: int,
        page_size: int,
        cp_size: int,
        device: torch.device,
    ):
        self.physical_size = int(physical_pages) * int(page_size)
        self.page_size = int(page_size)
        self.cp_size = int(cp_size)
        self.device = device
        logical_pages = int(physical_pages) * int(cp_size)
        self._owner_free_pages = None
        self._owner_release_pages = None
        self._flat_free_pages_cache = None
        self._flat_release_pages_cache = None
        self.free_pages = torch.arange(
            1, logical_pages + 1, dtype=torch.int64, device=device
        )
        self.release_pages = torch.empty((0,), dtype=torch.int64, device=device)
        self.debug_mode = False

    def _empty_pages(self) -> torch.Tensor:
        return torch.empty((0,), dtype=torch.int64, device=self.device)

    def _split_owner_buckets(
        self, pages: Optional[torch.Tensor]
    ) -> Optional[list[torch.Tensor]]:
        if pages is None:
            return None
        if pages.numel() == 0:
            return [torch.empty_like(pages) for _ in range(self.cp_size)]
        owner_ids = torch.remainder(pages - 1, self.cp_size)
        buckets: list[torch.Tensor] = []
        for owner in range(self.cp_size):
            owner_pages = pages[owner_ids == owner]
            if owner_pages.numel() > 1:
                owner_pages, _ = torch.sort(owner_pages)
            buckets.append(owner_pages)
        return buckets

    def _materialize_owner_buckets(
        self, buckets: Optional[list[torch.Tensor]]
    ) -> Optional[torch.Tensor]:
        if buckets is None:
            return None
        non_empty = [bucket for bucket in buckets if bucket.numel() > 0]
        if not non_empty:
            return self._empty_pages()
        return torch.cat(non_empty)

    @property
    def free_pages(self):
        if self._owner_free_pages is not None:
            if self._flat_free_pages_cache is None:
                self._flat_free_pages_cache = self._materialize_owner_buckets(
                    self._owner_free_pages
                )
            return self._flat_free_pages_cache
        return self._flat_free_pages_cache

    @free_pages.setter
    def free_pages(self, pages):
        self._flat_free_pages_cache = pages
        self._owner_free_pages = self._split_owner_buckets(pages)

    @property
    def release_pages(self):
        if self._owner_release_pages is not None:
            if self._flat_release_pages_cache is None:
                self._flat_release_pages_cache = self._materialize_owner_buckets(
                    self._owner_release_pages
                )
            return self._flat_release_pages_cache
        return self._flat_release_pages_cache

    @release_pages.setter
    def release_pages(self, pages):
        self._flat_release_pages_cache = pages
        self._owner_release_pages = self._split_owner_buckets(pages)

    def _owner_bucket_counts(self, buckets: Optional[list[torch.Tensor]]) -> list[int]:
        if buckets is None:
            return [0 for _ in range(self.cp_size)]
        return [int(bucket.numel()) for bucket in buckets]

    def _owner_available_counts(self) -> list[int]:
        free_counts = self._owner_bucket_counts(self._owner_free_pages)
        release_counts = self._owner_bucket_counts(self._owner_release_pages)
        return [
            free_count + release_count
            for free_count, release_count in zip(free_counts, release_counts)
        ]

    def _consume_owner_bucket_prefix(
        self,
        *,
        release: bool,
        counts_by_owner: list[int],
    ) -> None:
        target_attr = "_owner_release_pages" if release else "_owner_free_pages"
        cache_attr = "_flat_release_pages_cache" if release else "_flat_free_pages_cache"
        buckets = getattr(self, target_attr)
        mutated = False
        for owner, count in enumerate(counts_by_owner):
            if count <= 0:
                continue
            buckets[owner] = buckets[owner][count:]
            mutated = True
        if mutated:
            setattr(self, target_attr, buckets)
            setattr(self, cache_attr, None)

    def compute_owner_lane_stats(
        self,
        page_compute_owners: list[int],
    ) -> tuple[list[int], list[int], list[int]]:
        required = [0 for _ in range(self.cp_size)]
        for owner in page_compute_owners:
            if owner < 0 or owner >= self.cp_size:
                raise ValueError(
                    f"compute owner must be in [0, {self.cp_size}), got {owner}"
                )
            required[owner] += 1

        available = self._owner_available_counts()
        deficits = [
            max(0, required_count - available_count)
            for required_count, available_count in zip(required, available)
        ]
        return required, available, deficits

    def compute_owner_lane_capacity_pages(self) -> list[int]:
        capacity_pages = int(self.physical_size // self.page_size)
        return [capacity_pages for _ in range(self.cp_size)]

    def compute_owner_lane_free_room_stats(
        self,
        page_compute_owners: list[int],
        *,
        target_ratio: float,
        trigger_ratio: float,
    ) -> tuple[list[int], list[int], list[int]]:
        required, available, _exact_deficits = self.compute_owner_lane_stats(
            page_compute_owners
        )
        deficits = _compute_owner_lane_free_room_deficits(
            required=required,
            available=available,
            capacities=self.compute_owner_lane_capacity_pages(),
            target_ratio=target_ratio,
            trigger_ratio=trigger_ratio,
        )
        return required, available, deficits

    def _select_owner_free_pages_prefer_contiguous(
        self,
        owner_pages: torch.Tensor,
        required_count: int,
    ) -> torch.Tensor:
        return owner_pages[: min(required_count, int(owner_pages.numel()))]

    def _select_compute_owner_pages(
        self,
        page_compute_owners: list[int],
    ) -> Optional[tuple[torch.Tensor, list[int], list[int]]]:
        if not page_compute_owners:
            return (
                torch.empty((0,), dtype=torch.int64, device=self.device),
                [0 for _ in range(self.cp_size)],
                [0 for _ in range(self.cp_size)],
            )

        required_by_owner = [0 for _ in range(self.cp_size)]
        positions_by_owner: list[list[int]] = [[] for _ in range(self.cp_size)]
        for position, owner in enumerate(page_compute_owners):
            if owner < 0 or owner >= self.cp_size:
                raise ValueError(
                    f"compute owner must be in [0, {self.cp_size}), got {owner}"
                )
            required_by_owner[owner] += 1
            positions_by_owner[owner].append(position)

        lane_pages = [None for _ in range(self.cp_size)]
        selected_free_counts = [0 for _ in range(self.cp_size)]
        selected_release_counts = [0 for _ in range(self.cp_size)]
        for owner, required_count in enumerate(required_by_owner):
            if required_count == 0:
                continue

            selected_owner_free_mask = self._select_owner_free_pages_prefer_contiguous(
                self._owner_free_pages[owner], required_count
            )
            selected_owner_pages = selected_owner_free_mask

            free_count = int(selected_owner_pages.numel())
            remaining_count = required_count - free_count
            if remaining_count > 0:
                release_bucket = self._owner_release_pages[owner]
                if remaining_count > release_bucket.numel():
                    return None
                selected_owner_release_pages = release_bucket[:remaining_count]
                selected_owner_pages = torch.cat(
                    (selected_owner_pages, selected_owner_release_pages)
                )
                selected_release_counts[owner] = remaining_count

            selected_free_counts[owner] = free_count
            lane_pages[owner] = selected_owner_pages

        selected_pages = torch.empty(
            (len(page_compute_owners),), dtype=torch.int64, device=self.device
        )
        for owner, positions in enumerate(positions_by_owner):
            if not positions:
                continue
            position_tensor = torch.tensor(
                positions, dtype=torch.int64, device=self.device
            )
            selected_pages[position_tensor] = lane_pages[owner]

        return (
            selected_pages,
            selected_free_counts,
            selected_release_counts,
        )

    def alloc_pages_with_owners(
        self,
        page_compute_owners: list[int],
    ) -> Optional[torch.Tensor]:
        if not page_compute_owners:
            return torch.empty((0,), dtype=torch.int64, device=self.device)
        selected = self._select_compute_owner_pages(page_compute_owners)
        if selected is None:
            return None
        selected_pages, selected_free_counts, selected_release_counts = selected
        page_size = self.page_size
        base = selected_pages.to(torch.int64).unsqueeze(1) * page_size
        offsets = torch.arange(
            page_size, dtype=torch.int64, device=self.device
        ).unsqueeze(0)
        out_indices = (base + offsets).reshape(-1)
        self._consume_owner_bucket_prefix(
            release=False, counts_by_owner=selected_free_counts
        )
        self._consume_owner_bucket_prefix(
            release=True, counts_by_owner=selected_release_counts
        )
        return out_indices


def _is_page_contiguous_selection(selected: Optional[torch.Tensor], page_size: int) -> bool:
    if selected is None or selected.numel() == 0:
        return False
    pages = selected.view(-1, page_size)[:, 0] // page_size
    if pages.numel() <= 1:
        return True
    return bool(torch.all(pages[1:] - pages[:-1] == 1).item())


def _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,
    production: bool,
):
    if not production:
        return StandaloneCPSharedPagedAllocator(
            physical_pages=physical_pages,
            page_size=page_size,
            cp_size=cp_size,
            device=device,
        )

    from sglang.srt.mem_cache.allocator import CPSharedPagedTokenToKVPoolAllocator

    return CPSharedPagedTokenToKVPoolAllocator(
        logical_size=physical_pages * cp_size * page_size,
        physical_size=physical_pages * page_size,
        page_size=page_size,
        dtype=torch.bfloat16,
        device=str(device),
        kvcache=None,
        need_sort=False,
        cp_size=cp_size,
        cp_rank=0,
    )


def _patch_l1_fifo_selector(allocator) -> None:
    def fifo_selector(owner_pages: torch.Tensor, required_count: int) -> torch.Tensor:
        return owner_pages[: min(required_count, int(owner_pages.numel()))]

    allocator._select_owner_free_pages_prefer_contiguous = fifo_selector


def _is_l1_selection_physically_contiguous(
    selected: Optional[torch.Tensor], *, page_size: int, cp_size: int
) -> bool:
    if selected is None or selected.numel() == 0:
        return False
    logical_pages = selected.view(-1, page_size)[:, 0] // page_size
    owners = torch.remainder(logical_pages - 1, cp_size)
    physical_pages = torch.div(logical_pages - 1, cp_size, rounding_mode="floor") + 1
    for owner in torch.unique(owners).tolist():
        owner_phys = physical_pages[owners == int(owner)]
        if owner_phys.numel() <= 1:
            continue
        if not bool(torch.all(owner_phys[1:] - owner_phys[:-1] == 1).item()):
            return False
    return True


def _bench_l1_case(
    *,
    impl: str,
    op: str,
    physical_pages: int,
    request_pages: int,
    page_size: int,
    cp_size: int,
    device: torch.device,
    free_pattern: str,
    owner_pattern: str,
    repeat: int,
    warmup: int,
    seed: int,
    free_room_ratio: float,
    free_room_trigger_ratio: float,
    production_allocator: bool,
) -> BenchResult:
    request_owners = _make_page_compute_owners(request_pages, cp_size, owner_pattern)
    base_free_pages = _make_l1_free_pages(
        physical_pages=physical_pages,
        cp_size=cp_size,
        request_owners=request_owners,
        pattern=free_pattern,
        device=device,
        seed=seed,
    )
    samples_us: list[float] = []
    contiguous_hits = 0
    use_cuda = device.type == "cuda"
    if use_cuda:
        torch.cuda.synchronize(device)
    for iteration in range(warmup + repeat):
        allocator = _make_l1_allocator(
            physical_pages=physical_pages,
            page_size=page_size,
            cp_size=cp_size,
            device=device,
            production=production_allocator,
        )
        allocator.free_pages = base_free_pages.clone()
        allocator.release_pages = torch.empty((0,), dtype=torch.int64, device=device)
        if impl == "fifo":
            _patch_l1_fifo_selector(allocator)
        if use_cuda:
            torch.cuda.synchronize(device)
        start_ns = time.perf_counter_ns()
        selected = None
        if op == "stats":
            allocator.compute_owner_lane_stats(request_owners)
        elif op == "free_room_stats":
            allocator.compute_owner_lane_free_room_stats(
                request_owners,
                target_ratio=free_room_ratio,
                trigger_ratio=free_room_trigger_ratio,
            )
        elif op == "select_only":
            selected_result = allocator._select_compute_owner_pages(request_owners)
            if selected_result is not None:
                selected_pages = selected_result[0]
                base = selected_pages.to(torch.int64).unsqueeze(1) * page_size
                offsets = torch.arange(
                    page_size, dtype=torch.int64, device=device
                ).unsqueeze(0)
                selected = (base + offsets).reshape(-1)
        elif op == "alloc_pages":
            selected = allocator.alloc_pages_with_owners(request_owners)
        else:
            raise ValueError(f"unsupported l1 op: {op}")
        if use_cuda:
            torch.cuda.synchronize(device)
        elapsed_us = (time.perf_counter_ns() - start_ns) / 1000.0
        if iteration >= warmup:
            samples_us.append(elapsed_us)
            contiguous_hits += int(
                _is_l1_selection_physically_contiguous(
                    selected, page_size=page_size, cp_size=cp_size
                )
            )
    return _summarize(
        bench="l1",
        impl=f"{impl}:{op}",
        pattern=f"{free_pattern}:{owner_pattern}",
        device=device.type,
        total_pages=physical_pages,
        request_pages=request_pages,
        page_size=page_size,
        samples_us=samples_us,
        contiguous_hits=contiguous_hits,
    )


def _format_result(result: BenchResult) -> 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()]
    ops = [item.strip() for item in args.l1_ops.split(",") if item.strip()]

    results: list[BenchResult] = []
    for physical_pages in physical_pages_list:
        for request_pages in request_pages_list:
            if request_pages > physical_pages * args.cp_size:
                continue
            for free_pattern in free_patterns:
                for owner_pattern in owner_patterns:
                    for impl in impls:
                        for op in ops:
                            results.append(
                                _bench_l1_case(
                                    impl=impl,
                                    op=op,
                                    physical_pages=physical_pages,
                                    request_pages=request_pages,
                                    page_size=args.page_size,
                                    cp_size=args.cp_size,
                                    device=device,
                                    free_pattern=free_pattern,
                                    owner_pattern=owner_pattern,
                                    repeat=args.repeat,
                                    warmup=args.warmup,
                                    seed=args.seed,
                                    free_room_ratio=args.l1_free_room_ratio,
                                    free_room_trigger_ratio=args.l1_free_room_trigger_ratio,
                                    production_allocator=args.l1_allocator == "production",
                                )
                            )
    return results


def _build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--bench", default="host,l1", help="comma list: host,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-ops", default="alloc_pages")
    parser.add_argument(
        "--l1-allocator",
        choices=("production", "standalone"),
        default="production",
        help="Use production allocator imports or the dependency-light metadata copy.",
    )
    parser.add_argument("--l1-free-room-ratio", type=float, default=0.15)
    parser.add_argument("--l1-free-room-trigger-ratio", type=float, default=0.05)
    parser.add_argument(
        "--l1-free-patterns",
        default="sequential,owner_fragmented_later_run,random",
    )
    parser.add_argument("--l1-owner-patterns", default="round_robin,single_owner,zigzag")
    return parser


def main(argv: Optional[list[str]] = None) -> int:
    args = _build_parser().parse_args(argv)
    if args.repeat <= 0:
        raise ValueError("--repeat must be positive")
    if args.warmup < 0:
        raise ValueError("--warmup must be non-negative")

    results: list[BenchResult] = []
    benches = {item.strip() for item in args.bench.split(",") if item.strip()}
    if "host" in benches:
        results.extend(_run_host(args))
    if "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())