Stabilize CP HiCache residency under L1/L2 pressure
CP shared KV now keeps explicit L1 and host free-room targets so pressure is handled by planned eviction instead of repeated capacity-edge retries. The host allocator gains contiguous-preferred page reservation, L1 owner-lane allocation prefers contiguous physical pages, and CP HiCache metadata preserves pending backup safety for page-granular radix updates. Mooncake transfer stats and allocator microbenchmarks are included to make the remaining transfer bottlenecks measurable rather than inferred. Constraint: CP shared KV uses decode CP size 1 with all prefill CP ranks participating in transfer, so L1/L2 cache residency must remain page-granular and avoid extra collectives.\nConstraint: Production HiCache can be hundreds of GB, so allocator metadata overhead must be visible before enabling aggressive contiguous allocation broadly.\nRejected: Evict only the exact deficit | this keeps the cache at the cliff and causes repeated evict/allocate pressure.\nRejected: Rely on allocator scans alone for contiguity | remote microbenchmarks show fragmented 220GB-equivalent host metadata can make contiguous-preferred scans multi-ms.\nConfidence: medium\nScope-risk: moderate\nDirective: Do not increase L1/L2 free-room defaults or add new CP collectives without ETE evidence and transfer/allocator measurements.\nTested: python -m py_compile on touched runtime/test/benchmark files.\nTested: PYTHONPATH=. python -m pytest -q test/registered/unit/benchmark/test_cp_hicache_allocator_bench.py => 4 passed, 1 warning.\nTested: Remote g0034 log /mnt/beegfs/cjy/log/sglang_cp_hicache_20260601_233723.log shows active prefill process with L1/L2 free-room args, 702 HTTP 200 chat completions, 6272 prefill batches, and no fatal scheduler traceback in latest scan.\nTested: User-reported L1/L2 cache ETE validation passed on remote run.\nNot-tested: Full local pytest suite; local environment is missing several runtime dependencies.\nNot-tested: CUDA allocator microbenchmark during active production prefill process.\nNot-tested: Mooncake straggler fix; stats show transfer tail latency remains a separate bottleneck.
This commit is contained in:
laoyao0822
@@ -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())
|
||||
Reference in New Issue
Block a user