Remove full-cache scans from CP owner-lane allocation

The CP shared-KV allocator was still doing total-cache-sized CPU work in the scheduler hot path.  That cannot be hidden by GPU overlap, so owner-lane allocation now maintains per-owner free/release buckets and consumes request-sized prefixes instead of rebuilding masks over the full free-page tensor on each request.\n\nThe benchmark was extended to isolate L1 stats, selection, and allocation costs, and the CPU layout tests now install a complete sgl_kernel stub before importing SGLang helpers so remote unit collection does not abort in native extension loading.\n\nConstraint: Allocator CPU work blocks scheduler progress and cannot overlap with GPU forward execution.\nConstraint: CPU unit tests must not load native sgl_kernel on remote images where the loader can SIGABRT.\nRejected: Keep contiguous-run search over full free_pages | still scales with cache capacity and measured multi-ms overhead.\nRejected: Treat remote collection abort as an environment-only issue | it prevented allocator regression coverage and was fixable with a test-local stub.\nConfidence: high\nScope-risk: moderate\nDirective: CP owner-lane allocation is bucket-based; do not reintroduce full free_pages scans on the hot path without benchmark evidence.\nTested: Local py_compile for touched files\nTested: Local benchmark unit test, 6 passed\nTested: Remote benchmark unit test, 6 passed\nTested: Remote test_alloc_pages_with_owners.py, 10 passed\nTested: Remote test_cp_shared_kv_layout.py, 27 passed\nTested: Remote production allocator microbench shows select/alloc p50 reduced from ms-scale to sub-ms scale\nNot-tested: Full ETE traffic run after allocator bucket change
This commit is contained in:
laoyao0822
2026-06-02 08:41:00 +08:00
parent ce3a20d11b
commit 7c8fa2f71c
5 changed files with 906 additions and 115 deletions
@@ -1,9 +1,11 @@
import torch
from benchmark.hicache.bench_cp_hicache_allocator_overhead import (
StandaloneCPSharedPagedAllocator,
StandaloneHostAllocator,
_host_pages_from_gb,
_make_host_free_slots,
_make_page_compute_owners,
_parse_int_list,
)
@@ -45,3 +47,35 @@ def test_host_random_fragmented_has_requested_size():
)
assert free_slots.numel() == 64 * 16
assert torch.unique(free_slots).numel() == free_slots.numel()
def test_standalone_l1_allocator_reports_owner_lane_stats():
allocator = StandaloneCPSharedPagedAllocator(
physical_pages=4,
page_size=2,
cp_size=4,
device=torch.device("cpu"),
)
owners = _make_page_compute_owners(5, cp_size=4, pattern="round_robin")
required, available, deficits = allocator.compute_owner_lane_stats(owners)
assert required == [2, 1, 1, 1]
assert available == [4, 4, 4, 4]
assert deficits == [0, 0, 0, 0]
def test_standalone_l1_allocator_allocates_owner_matching_pages():
allocator = StandaloneCPSharedPagedAllocator(
physical_pages=8,
page_size=4,
cp_size=4,
device=torch.device("cpu"),
)
owners = [0, 1, 2, 3, 0]
selected = allocator.alloc_pages_with_owners(owners)
logical_pages = (selected.view(-1, allocator.page_size)[:, 0] // allocator.page_size)
selected_owners = torch.remainder(logical_pages - 1, allocator.cp_size).tolist()
assert selected_owners == owners