chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

Reduce CP HiCache L2 allocator scan cost

Browse Source 
Host HiCache reservations were paying token-level free-slot scans when trying to preserve page contiguity. The allocator now keeps a lazy page-extent index so availability checks and contiguous-preferred allocations avoid materializing the full 220GB-equivalent free-slot metadata path.

The companion benchmark models steady-state L2 churn near full occupancy, including burn-in and historical node-size effects, so LPF/RDMA descriptor quality can be separated from ETE noise.

Constraint: CP HiCache host allocations are page-shaped, but existing callers may still read free_slots directly.
Rejected: Sort and scan free_slots on each alloc_contiguous_preferred call | measured ms-level CPU overhead on 220GB-equivalent metadata.
Rejected: Remove free_slots compatibility | storage/tests still rely on the public tensor surface.
Confidence: medium
Scope-risk: moderate
Directive: Do not reintroduce per-allocation full free_slots scans on HostKVCache; preserve page-extent metadata or benchmark before changing allocator shape.
Tested: Local py_compile for memory_pool_host.py, allocator benchmark, and related tests.
Tested: Local test_cp_hicache_allocator_bench.py 10 passed.
Tested: Remote g0034 test_hicache_controller_cp.py 67 passed; test_cp_hicache_allocator_bench.py 10 passed.
Tested: Remote 220GB-equivalent host_churn benchmark: contiguous path reduced from ms-level to ~30-292us p50 depending on fragmentation.
Not-tested: Full CUDA ETE run after allocator change.
Not-tested: Production long-run fragmentation behavior under live traffic.
This commit is contained in:
laoyao0822
2026-06-02 23:35:45 +08:00
parent 7c8fa2f71c
commit 401de0f8ce
5 changed files with 925 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
benchmark/hicache/bench_cp_hicache_allocator_overhead.py
+393 -2
View File
@@ -14,6 +14,13 @@ Examples:
--bench host --host-sizes-gb 220 --request-pages 1,8,64,512 \
--patterns contiguous_fifo,fragmented_prefix_later_run,random_fragmented
# Steady-state L2 host churn model near full HiCache occupancy.
PYTHONPATH=. python benchmark/hicache/bench_cp_hicache_allocator_overhead.py \
--bench host_churn --host-sizes-gb 220 --request-pages 16,64,512 \
--host-churn-occupancies 0.90,0.97,0.99 \
--host-churn-evict-pages 64,512,2048 \
--host-churn-eviction-patterns oldest,random
# L1 allocator path on CUDA, stubbing sgl_kernel import if needed.
PYTHONPATH=python:. python benchmark/hicache/bench_cp_hicache_allocator_overhead.py \
--bench l1 --device cuda --stub-sgl-kernel --physical-pages 8192,32768 \
@@ -25,6 +32,7 @@ import argparse
import json
import math
import os
import random
import statistics
import sys
import threading
@@ -55,6 +63,36 @@ class BenchResult:
contiguous_ratio: float
@dataclass(frozen=True)
class HostChurnBenchResult:
bench: str
impl: str
pattern: str
device: str
total_pages: int
request_pages: int
page_size: int
repeat: int
target_occupancy: float
evict_pages: int
prefill_node_pages: int
burnin: int
mean_us: float
p50_us: float
p95_us: float
p99_us: float
min_us: float
max_us: float
contiguous_ratio: float
page_first_descriptors_per_op: int
lpf_descriptors_mean: float
lpf_descriptor_ratio_mean: float
run_count_p50: float
run_count_p95: float
max_run_pages_mean: float
max_run_pages_p50: float
def _parse_int_list(value: str | Iterable[int]) -> list[int]:
if isinstance(value, str):
return [int(item.strip()) for item in value.split(",") if item.strip()]
@@ -244,6 +282,10 @@ class StandaloneHostAllocator:
self.free_slots = self.free_slots[keep_mask]
return select_index
def free(self, indices: torch.Tensor) -> int:
self.free_slots = torch.cat([self.free_slots, indices.cpu()])
return int(indices.numel())
def _compute_owner_lane_free_room_deficits(
*,
@@ -544,6 +586,43 @@ def _is_page_contiguous_selection(selected: Optional[torch.Tensor], page_size: i
return bool(torch.all(pages[1:] - pages[:-1] == 1).item())
def _page_run_lengths_from_token_slots(
selected: Optional[torch.Tensor], page_size: int
) -> list[int]:
"""Return consecutive physical-page run lengths for one host selection.
This is the layout-independent descriptor proxy used by the L2 benchmark:
current ``page_first_direct`` needs one fixed-layer copy descriptor per page,
while ``layer_page_first`` can collapse each consecutive page run to one
descriptor per KV tensor.
"""
if selected is None or selected.numel() == 0:
return []
if selected.numel() % page_size != 0:
raise ValueError(
f"selected token slots must be page-shaped, got {selected.numel()=} "
f"{page_size=}"
)
pages = (selected.view(-1, page_size)[:, 0] // page_size).tolist()
if not pages:
return []
run_lengths: list[int] = []
current_len = 1
prev_page = int(pages[0])
for page in pages[1:]:
page = int(page)
if page == prev_page + 1:
current_len += 1
else:
run_lengths.append(current_len)
current_len = 1
prev_page = page
run_lengths.append(current_len)
return run_lengths
def _make_host_allocator(impl: str, *, page_size: int, free_slots: torch.Tensor):
if impl == "standalone":
return StandaloneHostAllocator(page_size=page_size, free_slots=free_slots)
@@ -614,6 +693,63 @@ def _summarize(
)
def _summarize_host_churn(
*,
impl: str,
method: str,
eviction_pattern: str,
total_pages: int,
request_pages: int,
page_size: int,
target_occupancy: float,
evict_pages: int,
prefill_node_pages: int,
burnin: int,
samples_us: list[float],
run_counts: list[int],
max_run_lengths: list[int],
contiguous_hits: int,
) -> HostChurnBenchResult:
repeat = len(samples_us)
lpf_descriptors_mean = (
float(statistics.mean(run_counts)) if run_counts else 0.0
)
return HostChurnBenchResult(
bench="host_churn",
impl=f"{impl}:{method}",
pattern=f"occ={target_occupancy:.2f}:{eviction_pattern}",
device="cpu",
total_pages=int(total_pages),
request_pages=int(request_pages),
page_size=int(page_size),
repeat=repeat,
target_occupancy=float(target_occupancy),
evict_pages=int(evict_pages),
prefill_node_pages=int(prefill_node_pages),
burnin=int(burnin),
mean_us=float(statistics.mean(samples_us)) if samples_us else 0.0,
p50_us=float(_percentile(samples_us, 50)),
p95_us=float(_percentile(samples_us, 95)),
p99_us=float(_percentile(samples_us, 99)),
min_us=float(min(samples_us)) if samples_us else 0.0,
max_us=float(max(samples_us)) if samples_us else 0.0,
contiguous_ratio=float(contiguous_hits / repeat) if repeat else 0.0,
page_first_descriptors_per_op=int(request_pages),
lpf_descriptors_mean=lpf_descriptors_mean,
lpf_descriptor_ratio_mean=(
lpf_descriptors_mean / float(request_pages) if request_pages else 0.0
),
run_count_p50=float(_percentile([float(x) for x in run_counts], 50)),
run_count_p95=float(_percentile([float(x) for x in run_counts], 95)),
max_run_pages_mean=(
float(statistics.mean(max_run_lengths)) if max_run_lengths else 0.0
),
max_run_pages_p50=float(
_percentile([float(x) for x in max_run_lengths], 50)
),
)
def _bench_host_case(
*,
impl: str,
@@ -660,6 +796,166 @@ def _bench_host_case(
)
def _evict_host_churn_nodes(
*,
allocator,
active_nodes: list[torch.Tensor],
target_pages_to_free: int,
page_size: int,
eviction_pattern: str,
rng: random.Random,
) -> int:
freed_pages = 0
while active_nodes and freed_pages < target_pages_to_free:
if eviction_pattern == "oldest":
node_index = 0
elif eviction_pattern == "youngest":
node_index = len(active_nodes) - 1
elif eviction_pattern == "random":
node_index = rng.randrange(len(active_nodes))
else:
raise ValueError(f"unsupported host churn eviction pattern: {eviction_pattern}")
node = active_nodes.pop(node_index)
allocator.free(node)
freed_pages += int(node.numel()) // page_size
return freed_pages
def _bench_host_churn_case(
*,
impl: str,
method: str,
total_pages: int,
request_pages: int,
page_size: int,
target_occupancy: float,
evict_pages: int,
eviction_pattern: str,
repeat: int,
warmup: int,
seed: int,
prefill_node_pages: Optional[int] = None,
burnin: int = 0,
) -> HostChurnBenchResult:
if not 0 < target_occupancy < 1:
raise ValueError(
f"target_occupancy must be in (0, 1), got {target_occupancy}"
)
if request_pages <= 0:
raise ValueError(f"request_pages must be positive, got {request_pages}")
if evict_pages <= 0:
raise ValueError(f"evict_pages must be positive, got {evict_pages}")
if burnin < 0:
raise ValueError(f"burnin must be non-negative, got {burnin}")
if request_pages > total_pages:
raise ValueError(
f"request_pages must be <= total_pages, got {request_pages=} {total_pages=}"
)
if prefill_node_pages is None:
prefill_node_pages = request_pages
if prefill_node_pages <= 0:
raise ValueError(
f"prefill_node_pages must be positive, got {prefill_node_pages}"
)
if prefill_node_pages > total_pages:
raise ValueError(
"prefill_node_pages must be <= total_pages, got "
f"{prefill_node_pages=} {total_pages=}"
)
base_free_slots = _make_host_free_slots(
total_pages=total_pages,
request_pages=request_pages,
page_size=page_size,
pattern="contiguous_fifo",
seed=seed,
)
allocator = _make_host_allocator(impl, page_size=page_size, free_slots=base_free_slots)
rng = random.Random(seed)
need_size = request_pages * page_size
prefill_need_size = prefill_node_pages * page_size
# Fill with configurable node sizes so the benchmark can model fragmented
# steady-state HiCache: many old small nodes can be evicted to satisfy one
# larger new request, which is the path LPF allocation policy cares about.
target_used_pages = min(
total_pages - request_pages,
int(math.floor(float(total_pages) * float(target_occupancy))),
)
target_used_pages = (target_used_pages // prefill_node_pages) * prefill_node_pages
active_nodes: list[torch.Tensor] = []
used_pages = 0
while used_pages + prefill_node_pages <= target_used_pages:
selected = allocator.alloc(prefill_need_size)
if selected is None:
break
active_nodes.append(selected)
used_pages += prefill_node_pages
samples_us: list[float] = []
run_counts: list[int] = []
max_run_lengths: list[int] = []
contiguous_hits = 0
fn = allocator.alloc if method == "fifo" else allocator.alloc_contiguous_preferred
min_evict_pages = max(evict_pages, request_pages)
first_sample_iteration = int(burnin) + int(warmup)
for iteration in range(first_sample_iteration + repeat):
_evict_host_churn_nodes(
allocator=allocator,
active_nodes=active_nodes,
target_pages_to_free=min_evict_pages,
page_size=page_size,
eviction_pattern=eviction_pattern,
rng=rng,
)
while allocator.available_size() < need_size and active_nodes:
_evict_host_churn_nodes(
allocator=allocator,
active_nodes=active_nodes,
target_pages_to_free=request_pages,
page_size=page_size,
eviction_pattern=eviction_pattern,
rng=rng,
)
start_ns = time.perf_counter_ns()
selected = fn(need_size)
elapsed_us = (time.perf_counter_ns() - start_ns) / 1000.0
if selected is None:
raise RuntimeError(
"host churn allocation failed after eviction: "
f"{total_pages=} {request_pages=} {target_occupancy=} "
f"{evict_pages=} {eviction_pattern=}"
)
active_nodes.append(selected)
if iteration >= first_sample_iteration:
samples_us.append(elapsed_us)
run_lengths = _page_run_lengths_from_token_slots(selected, page_size)
run_count = len(run_lengths)
run_counts.append(run_count)
max_run_lengths.append(max(run_lengths) if run_lengths else 0)
contiguous_hits += int(run_count <= 1)
return _summarize_host_churn(
impl=impl,
method=method,
eviction_pattern=eviction_pattern,
total_pages=total_pages,
request_pages=request_pages,
page_size=page_size,
target_occupancy=target_occupancy,
evict_pages=evict_pages,
prefill_node_pages=prefill_node_pages,
burnin=burnin,
samples_us=samples_us,
run_counts=run_counts,
max_run_lengths=max_run_lengths,
contiguous_hits=contiguous_hits,
)
def _zigzag_owners(num_pages: int, cp_size: int) -> list[int]:
segment_num = cp_size * 2
base = num_pages // segment_num
@@ -915,7 +1211,20 @@ def _bench_l1_case(
)
def _format_result(result: BenchResult) -> str:
def _format_result(result: BenchResult | HostChurnBenchResult) -> str:
if isinstance(result, HostChurnBenchResult):
return (
f"{result.bench:10s} impl={result.impl:20s} pattern={result.pattern:34s} "
f"dev={result.device:4s} pages={result.total_pages:7d} req={result.request_pages:5d} "
f"evict={result.evict_pages:5d} prefill_node={result.prefill_node_pages:5d} "
f"burnin={result.burnin:4d} p50={result.p50_us:9.2f}us "
f"p95={result.p95_us:9.2f}us p99={result.p99_us:9.2f}us "
f"mean={result.mean_us:9.2f}us contig={result.contiguous_ratio:.2f} "
f"pf_desc={result.page_first_descriptors_per_op:d} "
f"lpf_desc_mean={result.lpf_descriptors_mean:.2f} "
f"lpf_ratio={result.lpf_descriptor_ratio_mean:.3f} "
f"run_p50={result.run_count_p50:.1f} max_run_mean={result.max_run_pages_mean:.1f}"
)
return (
f"{result.bench:4s} impl={result.impl:20s} pattern={result.pattern:34s} "
f"dev={result.device:4s} pages={result.total_pages:7d} req={result.request_pages:5d} "
@@ -965,6 +1274,64 @@ def _run_host(args) -> list[BenchResult]:
return results
def _run_host_churn(args) -> list[HostChurnBenchResult]:
host_pages = _parse_int_list(args.host_pages) if args.host_pages else []
for size_gb in _parse_float_list(args.host_sizes_gb):
host_pages.append(
_host_pages_from_gb(
size_gb, bytes_per_token=args.bytes_per_token, page_size=args.page_size
)
)
if not host_pages:
host_pages = [8192, 16384, 32768]
host_pages = sorted(set(page for page in host_pages if page > 0))
request_pages_list = _parse_int_list(args.request_pages)
host_impls = [item.strip() for item in args.host_impl.split(",") if item.strip()]
methods = [item.strip() for item in args.host_methods.split(",") if item.strip()]
occupancies = _parse_float_list(args.host_churn_occupancies)
evict_pages_list = _parse_int_list(args.host_churn_evict_pages)
prefill_node_pages_list = (
_parse_int_list(args.host_churn_prefill_node_pages)
if args.host_churn_prefill_node_pages
else [0]
)
eviction_patterns = [
item.strip() for item in args.host_churn_eviction_patterns.split(",") if item.strip()
]
results: list[HostChurnBenchResult] = []
for total_pages in host_pages:
for request_pages in request_pages_list:
if request_pages > total_pages:
continue
for target_occupancy in occupancies:
for evict_pages in evict_pages_list:
for prefill_node_pages in prefill_node_pages_list:
for eviction_pattern in eviction_patterns:
for impl in host_impls:
for method in methods:
results.append(
_bench_host_churn_case(
impl=impl,
method=method,
total_pages=total_pages,
request_pages=request_pages,
page_size=args.page_size,
target_occupancy=target_occupancy,
evict_pages=evict_pages,
eviction_pattern=eviction_pattern,
repeat=args.repeat,
warmup=args.warmup,
seed=args.seed,
prefill_node_pages=(
prefill_node_pages or None
),
burnin=args.host_churn_burnin,
)
)
return results
def _run_l1(args) -> list[BenchResult]:
if args.stub_sgl_kernel:
_install_sgl_kernel_stubs()
@@ -1015,7 +1382,9 @@ def _run_l1(args) -> list[BenchResult]:
def _build_parser() -> argparse.ArgumentParser:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--bench", default="host,l1", help="comma list: host,l1")
parser.add_argument(
"--bench", default="host,l1", help="comma list: host,host_churn,l1"
)
parser.add_argument("--page-size", type=int, default=64)
parser.add_argument("--repeat", type=int, default=20)
parser.add_argument("--warmup", type=int, default=5)
@@ -1036,6 +1405,26 @@ def _build_parser() -> argparse.ArgumentParser:
)
parser.add_argument("--host-impl", default="standalone")
parser.add_argument("--host-methods", default="fifo,contiguous")
parser.add_argument("--host-churn-occupancies", default="0.90,0.97,0.99")
parser.add_argument("--host-churn-evict-pages", default="64,512,2048")
parser.add_argument(
"--host-churn-prefill-node-pages",
default="",
help=(
"comma list of node sizes used to prefill steady-state host cache; "
"default uses each request_pages value"
),
)
parser.add_argument(
"--host-churn-burnin",
type=int,
default=0,
help=(
"unmeasured steady-state evict+allocate iterations before warmup; "
"useful for exhausting the cold contiguous free tail"
),
)
parser.add_argument("--host-churn-eviction-patterns", default="oldest,random")
parser.add_argument("--physical-pages", default="8192,32768")
parser.add_argument("--cp-size", type=int, default=8)
@@ -1071,6 +1460,8 @@ def main(argv: Optional[list[str]] = None) -> int:
benches = {item.strip() for item in args.bench.split(",") if item.strip()}
if "host" in benches:
results.extend(_run_host(args))
if "host_churn" in benches:
results.extend(_run_host_churn(args))
if "l1" in benches:
results.extend(_run_l1(args))