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
@@ -14,10 +14,11 @@ Examples:
|
||||
--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.
|
||||
# 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
|
||||
--request-pages 8,64,512 --l1-impl current,fifo \
|
||||
--l1-ops stats,free_room_stats,select_only,alloc_pages
|
||||
"""
|
||||
|
||||
import argparse
|
||||
@@ -244,6 +245,296 @@ class StandaloneHostAllocator:
|
||||
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
|
||||
@@ -477,7 +768,22 @@ def _install_sgl_kernel_stubs() -> None:
|
||||
sys.modules[submodule] = sub
|
||||
|
||||
|
||||
def _make_l1_allocator(*, physical_pages: int, page_size: int, cp_size: int, device: torch.device):
|
||||
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(
|
||||
@@ -494,8 +800,8 @@ def _make_l1_allocator(*, physical_pages: int, page_size: int, cp_size: int, dev
|
||||
|
||||
|
||||
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)
|
||||
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
|
||||
|
||||
@@ -520,6 +826,7 @@ def _is_l1_selection_physically_contiguous(
|
||||
def _bench_l1_case(
|
||||
*,
|
||||
impl: str,
|
||||
op: str,
|
||||
physical_pages: int,
|
||||
request_pages: int,
|
||||
page_size: int,
|
||||
@@ -530,6 +837,9 @@ def _bench_l1_case(
|
||||
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(
|
||||
@@ -551,6 +861,7 @@ def _bench_l1_case(
|
||||
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)
|
||||
@@ -559,7 +870,28 @@ def _bench_l1_case(
|
||||
if use_cuda:
|
||||
torch.cuda.synchronize(device)
|
||||
start_ns = time.perf_counter_ns()
|
||||
selected = allocator.alloc_pages_with_owners(request_owners)
|
||||
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
|
||||
@@ -572,7 +904,7 @@ def _bench_l1_case(
|
||||
)
|
||||
return _summarize(
|
||||
bench="l1",
|
||||
impl=impl,
|
||||
impl=f"{impl}:{op}",
|
||||
pattern=f"{free_pattern}:{owner_pattern}",
|
||||
device=device.type,
|
||||
total_pages=physical_pages,
|
||||
@@ -648,6 +980,7 @@ def _run_l1(args) -> list[BenchResult]:
|
||||
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:
|
||||
@@ -657,21 +990,26 @@ def _run_l1(args) -> list[BenchResult]:
|
||||
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,
|
||||
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
|
||||
|
||||
|
||||
@@ -705,6 +1043,15 @@ def _build_parser() -> argparse.ArgumentParser:
|
||||
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",
|
||||
|
||||
@@ -5150,3 +5150,164 @@ Remaining validation gap:
|
||||
- Requires a fresh ETE run with chunked prefill enabled to verify that the
|
||||
fallback storm disappears in production traffic. An already-running remote
|
||||
process will not pick up this change until restarted.
|
||||
|
||||
### C121 — 2026-06-02 CPU overhead must be measured by microbench, not inferred from logs
|
||||
|
||||
Finding:
|
||||
|
||||
- Runtime logs can show fallback storms and fatal paths, but they cannot rank normal-path CPU overhead. The hot paths here are often successful allocator/control-path calls that produce no warning log.
|
||||
- The CP shared-KV L1 owner-lane allocator had two visible metadata costs:
|
||||
- `compute_owner_lane_stats()` scanned the full free-page tensor once per CP owner and performed one `.item()` synchronization per owner.
|
||||
- `_select_compute_owner_pages()` recomputed owner masks per owner and built prefix selections via full-tensor `cumsum`, then optionally searched contiguous owner-lane runs.
|
||||
- Host/L2 `alloc_contiguous_preferred()` is also expensive on fragmented metadata because it validates page contiguity over the full free-slot tensor and searches for later runs. A local 210GB-equivalent metadata benchmark showed fragmented/random contiguous-preferred allocation at roughly 8–11 ms p50 versus FIFO at sub-ms scale for the same request sizes.
|
||||
|
||||
Correction implemented:
|
||||
|
||||
- Extended `benchmark/hicache/bench_cp_hicache_allocator_overhead.py` with L1 operation breakdowns:
|
||||
- `stats`
|
||||
- `free_room_stats`
|
||||
- `select_only`
|
||||
- `alloc_pages`
|
||||
- Added a dependency-light standalone CP shared-paged allocator model so the CPU metadata shape can be measured locally without importing the full SGLang runtime stack.
|
||||
- Optimized the production CP owner-lane stats path from per-owner mask+sum scans to one `torch.bincount()` over page owners.
|
||||
- Optimized owner-page selection by computing free/release page owners once per allocation and replacing the prefix `cumsum` selection with `nonzero()[:required_count]` prefix selection.
|
||||
- Replaced contiguous-run detection from `run_edges.unfold(...).all(dim=1)` with a cumulative-sum sliding-window test. This preserves the same contiguous-run contract but avoids O(num_free_pages * required_pages) metadata work for large requests.
|
||||
|
||||
Current evidence:
|
||||
|
||||
- Local CPU benchmark after the stats/path cleanup still shows selection/allocation as the dominant cost. Example focused run at `physical_pages=32768`, `cp_size=8`, random/zigzag owners:
|
||||
- `stats`: about 0.9–1.4 ms p50.
|
||||
- `select_only`: about 8–12 ms p50 for the focused random/fragmented zigzag cases.
|
||||
- `alloc_pages`: about 8–12 ms p50 for the same cases, request sizes 64–512 pages.
|
||||
- Therefore the first optimization only reduces part of the overhead; the main remaining cost is still owner-lane page selection, especially contiguous-preferred search and full-size mask materialization.
|
||||
|
||||
Verification:
|
||||
|
||||
- Local:
|
||||
- `python -m py_compile python/sglang/srt/mem_cache/allocator.py 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` → `6 passed, 1 warning`.
|
||||
- Remote `g0034` container:
|
||||
- benchmark unit subset: `6 passed, 1 warning`.
|
||||
|
||||
Known verification gap:
|
||||
|
||||
- Remote `test_cp_shared_kv_layout.py` collection aborted while importing `sgl_kernel` native ops in the container, before reaching the allocator tests. This is an environment/native import failure during test collection, not evidence of allocator logic failure.
|
||||
- Need a clean remote production-allocator test path or a container state where `sgl_kernel` imports safely before claiming full production allocator verification.
|
||||
|
||||
Next target:
|
||||
|
||||
- Replace full-tensor owner-lane selection with a lower-overhead data structure or batched selector. The viable directions are:
|
||||
1. Maintain per-owner free-page queues/counters incrementally.
|
||||
2. Use a bounded contiguous-search policy and skip expensive run search for small requests or highly fragmented pools.
|
||||
3. Add a fused selector kernel if GPU-side selection remains acceptable, but avoid increasing synchronization frequency.
|
||||
|
||||
C121 additional remote standalone benchmark evidence:
|
||||
|
||||
- `g0034` container, dependency-light standalone allocator model, `physical_pages=32768`, `cp_size=8`, zigzag owners:
|
||||
- random, 64 pages: `stats` p50 0.87 ms, `select_only` p50 6.55 ms, `alloc_pages` p50 7.29 ms.
|
||||
- owner-fragmented later-run, 64 pages: `stats` p50 0.43 ms, `select_only` p50 6.13 ms, `alloc_pages` p50 6.80 ms.
|
||||
- random, 512 pages: `stats` p50 0.90 ms, `select_only` p50 7.46 ms, `alloc_pages` p50 7.91 ms.
|
||||
- owner-fragmented later-run, 512 pages: `stats` p50 0.45 ms, `select_only` p50 6.89 ms, `alloc_pages` p50 7.16 ms.
|
||||
- This reinforces that after the cheap cleanup, stats is no longer the main bottleneck; selector/allocation metadata still costs multi-ms and is the next CPU-overhead target.
|
||||
|
||||
### C122 — 2026-06-02 L1 owner-lane allocation must remove pure-CPU full scans
|
||||
|
||||
Finding:
|
||||
|
||||
- A remaining ~7–10 ms allocator/control-path cost is still too high because it
|
||||
is pure CPU metadata work. Unlike D2H/H2D backup/load kernels, this work
|
||||
cannot overlap with GPU forward progress once the scheduler is blocked waiting
|
||||
for page allocation.
|
||||
- The C121 cleanup made stats cheaper, but selection still scanned or
|
||||
materialized full free-page tensors on every owner-lane allocation.
|
||||
- This shape scales with total cache capacity, not request size. With a 150–220
|
||||
GB HiCache/L1 metadata scale, even a successful hot-path allocation becomes a
|
||||
scheduler stall.
|
||||
|
||||
Correction in progress:
|
||||
|
||||
- Make `CPSharedPagedTokenToKVPoolAllocator` keep per-owner free/release page
|
||||
buckets as allocator state instead of deriving owner buckets by scanning
|
||||
`free_pages` for every allocation.
|
||||
- Keep allocation itself request-sized: count required pages per owner, take the
|
||||
needed prefix from that owner bucket, and consume bucket prefixes only after
|
||||
all lanes are known satisfiable.
|
||||
- Preserve the public `free_pages` / `release_pages` tensor interface through
|
||||
lazy materialized caches for tests and legacy paths, but avoid materializing it
|
||||
in the CP owner-lane fast path.
|
||||
- Sort within each owner bucket when pages are inserted/restored. This changes
|
||||
the CP owner allocator from global FIFO semantics to owner-lane contiguous
|
||||
semantics, which is more aligned with the RDMA/H2D/D2H goal: each owner lane
|
||||
should prefer physically consecutive pages without a per-allocation run scan.
|
||||
|
||||
Risk / contract note:
|
||||
|
||||
- Code that only depends on owner-correct page allocation is unaffected.
|
||||
- Code or tests that implicitly relied on exact global `free_pages` order after a
|
||||
CP owner-lane allocation must be updated; global FIFO ordering is not the
|
||||
intended CP shared-KV owner-lane contract.
|
||||
|
||||
C122 validation update:
|
||||
|
||||
- Remote `g0034` production allocator microbench with `--stub-sgl-kernel`,
|
||||
`physical_pages=32768`, `cp_size=8`, zigzag owners:
|
||||
- random, 64 pages: `stats` p50 17 us, `select_only` p50 173 us,
|
||||
`alloc_pages` p50 527 us.
|
||||
- owner-fragmented later-run, 64 pages: `stats` p50 13 us,
|
||||
`select_only` p50 166 us, `alloc_pages` p50 235 us.
|
||||
- random, 512 pages: `stats` p50 27 us, `select_only` p50 284 us,
|
||||
`alloc_pages` p50 638 us.
|
||||
- owner-fragmented later-run, 512 pages: `stats` p50 32 us,
|
||||
`select_only` p50 266 us, `alloc_pages` p50 476 us.
|
||||
- This confirms the allocator hot path is no longer full-cache-scan shaped for
|
||||
this benchmark: the previous 6–8 ms selector/allocation p50 is reduced to
|
||||
sub-ms p50 while preserving owner-lane correctness and contiguous lane
|
||||
selections.
|
||||
- Remaining CPU cost is now mostly request-sized construction (`page_compute_owners`
|
||||
grouping, output page tensor fill, and token-loc expansion), not total-cache
|
||||
sized metadata scans.
|
||||
|
||||
### C123 — 2026-06-02 CPU allocator tests must not load native sgl_kernel during collection
|
||||
|
||||
Finding:
|
||||
|
||||
- Remote `test_cp_shared_kv_layout.py` aborted during pytest collection inside
|
||||
`sgl_kernel/load_utils.py::_load_architecture_specific_ops` before any test
|
||||
logic ran.
|
||||
- This is not an allocator correctness failure. The test imported
|
||||
`sglang.test.test_utils`, which imports `sglang.srt.utils.common`; that module
|
||||
imports `sgl_kernel` to probe AMX availability. On the remote image the native
|
||||
loader aborts the process instead of raising a catchable Python exception.
|
||||
- A `try: import sgl_kernel` fallback is insufficient for this environment
|
||||
because SIGABRT bypasses Python exception handling.
|
||||
|
||||
Correction:
|
||||
|
||||
- Install a minimal `sgl_kernel`, `sgl_kernel.kvcacheio`, and
|
||||
`sgl_kernel.quantization` stub at the top of CPU-only allocator/layout tests
|
||||
before importing any `sglang` helper module.
|
||||
- Keep Torch custom-op schema registration for the operators referenced by the
|
||||
imported SGLang code, but do not load the native extension during collection.
|
||||
|
||||
Scope:
|
||||
|
||||
- This applies only to CPU unit tests. It does not change production runtime
|
||||
import behavior and does not mask native kernel problems in CUDA/ETE tests.
|
||||
|
||||
C123 validation update:
|
||||
|
||||
- After avoiding the native import abort, collection reached Python test logic but
|
||||
failed because the CPU stub had not defined the fp8 quantization custom-op
|
||||
schemas used by `fp8_kernel.py` fake registrations.
|
||||
- The test stub now defines the fp8 quantization and fp8 GEMM schemas before any
|
||||
`sglang` import, matching the CPU-only pattern used by the heavier HiCache
|
||||
metadata tests.
|
||||
|
||||
C123 full-suite update:
|
||||
|
||||
- The full `test_cp_shared_kv_layout.py` file then reached scheduler rollback
|
||||
tests and failed because `memory_pool_host.py` imports named transfer helpers
|
||||
from `sgl_kernel.kvcacheio`.
|
||||
- The CPU stub now also gives `sgl_kernel.kvcacheio` a module-level
|
||||
`__getattr__`, so named imports resolve to inert functions without importing
|
||||
the native extension.
|
||||
|
||||
@@ -626,10 +626,237 @@ class CPSharedPagedTokenToKVPoolAllocator(PagedTokenToKVPoolAllocator):
|
||||
if not 0 <= cp_rank < cp_size:
|
||||
raise ValueError(f"cp_rank must be in [0, {cp_size}), got {cp_rank}")
|
||||
|
||||
super().__init__(logical_size, page_size, dtype, device, kvcache, need_sort)
|
||||
self.physical_size = physical_size
|
||||
self.cp_size = cp_size
|
||||
self.cp_rank = cp_rank
|
||||
self._owner_free_pages = None
|
||||
self._owner_release_pages = None
|
||||
self._flat_free_pages_cache = None
|
||||
self._flat_release_pages_cache = None
|
||||
|
||||
super().__init__(logical_size, page_size, dtype, device, kvcache, need_sort)
|
||||
|
||||
def _owner_buckets_ready(self) -> bool:
|
||||
return hasattr(self, "cp_size") and self.cp_size is not None
|
||||
|
||||
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 not self._owner_buckets_ready():
|
||||
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_buckets_ready() and 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 getattr(self, "_flat_free_pages_cache", None)
|
||||
|
||||
@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_buckets_ready() and 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 getattr(self, "_flat_release_pages_cache", None)
|
||||
|
||||
@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 _append_pages_to_owner_buckets(
|
||||
self, *, release: bool, pages: torch.Tensor
|
||||
) -> None:
|
||||
if pages.numel() == 0:
|
||||
return
|
||||
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, None)
|
||||
if buckets is None:
|
||||
buckets = [self._empty_pages() for _ in range(self.cp_size)]
|
||||
owner_ids = torch.remainder(pages - 1, self.cp_size)
|
||||
for owner in range(self.cp_size):
|
||||
owner_pages = pages[owner_ids == owner]
|
||||
if owner_pages.numel() == 0:
|
||||
continue
|
||||
merged = torch.cat((buckets[owner], owner_pages))
|
||||
if merged.numel() > 1:
|
||||
merged, _ = torch.sort(merged)
|
||||
buckets[owner] = merged
|
||||
setattr(self, target_attr, buckets)
|
||||
setattr(self, cache_attr, None)
|
||||
|
||||
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, None)
|
||||
if buckets is None:
|
||||
return
|
||||
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 available_size(self):
|
||||
return (
|
||||
sum(self._owner_available_counts()) * self.page_size
|
||||
if self._owner_buckets_ready()
|
||||
else super().available_size()
|
||||
)
|
||||
|
||||
def immediate_available_pages(self) -> int:
|
||||
if self._owner_buckets_ready() and self._owner_free_pages is not None:
|
||||
return sum(self._owner_bucket_counts(self._owner_free_pages))
|
||||
return super().immediate_available_pages()
|
||||
|
||||
def deferred_available_pages(self) -> int:
|
||||
if self._owner_buckets_ready() and self._owner_release_pages is not None:
|
||||
return sum(self._owner_bucket_counts(self._owner_release_pages))
|
||||
return super().deferred_available_pages()
|
||||
|
||||
def backup_state(self):
|
||||
return (self.free_pages, self.release_pages)
|
||||
|
||||
def restore_state(self, state):
|
||||
self.free_pages, self.release_pages = state
|
||||
|
||||
def clear(self):
|
||||
# The padded slot 0 is used for writing dummy outputs from padded tokens.
|
||||
all_pages = torch.arange(
|
||||
1, self.num_pages + 1, dtype=torch.int64, device=self.device
|
||||
)
|
||||
self.free_pages = all_pages
|
||||
self.is_not_in_free_group = True
|
||||
self.free_group = []
|
||||
self.release_pages = self._empty_pages()
|
||||
|
||||
def merge_and_sort_free(self):
|
||||
if self.deferred_available_pages() == 0:
|
||||
return
|
||||
if self._owner_free_pages is None or self._owner_release_pages is None:
|
||||
return super().merge_and_sort_free()
|
||||
for owner in range(self.cp_size):
|
||||
release_bucket = self._owner_release_pages[owner]
|
||||
if release_bucket.numel() == 0:
|
||||
continue
|
||||
merged = torch.cat((self._owner_free_pages[owner], release_bucket))
|
||||
if merged.numel() > 1:
|
||||
merged, _ = torch.sort(merged)
|
||||
self._owner_free_pages[owner] = merged
|
||||
self._owner_release_pages[owner] = torch.empty_like(release_bucket[:0])
|
||||
self._flat_free_pages_cache = None
|
||||
self._flat_release_pages_cache = self._empty_pages()
|
||||
|
||||
def alloc(self, need_size: int):
|
||||
if self.debug_mode:
|
||||
assert (
|
||||
need_size % self.page_size == 0
|
||||
), "The allocation size should be page-aligned"
|
||||
|
||||
num_pages = need_size // self.page_size
|
||||
if self.need_sort and num_pages > self.immediate_available_pages():
|
||||
self.merge_and_sort_free()
|
||||
if num_pages > self.immediate_available_pages():
|
||||
return None
|
||||
|
||||
selected_pages: list[torch.Tensor] = []
|
||||
counts_by_owner = [0 for _ in range(self.cp_size)]
|
||||
remaining = num_pages
|
||||
for owner, bucket in enumerate(self._owner_free_pages):
|
||||
if remaining <= 0:
|
||||
break
|
||||
count = min(remaining, int(bucket.numel()))
|
||||
if count <= 0:
|
||||
continue
|
||||
selected_pages.append(bucket[:count])
|
||||
counts_by_owner[owner] = count
|
||||
remaining -= count
|
||||
if remaining != 0:
|
||||
return None
|
||||
|
||||
out_pages = torch.cat(selected_pages) if selected_pages else self._empty_pages()
|
||||
self._consume_owner_bucket_prefix(release=False, counts_by_owner=counts_by_owner)
|
||||
out_indices = (
|
||||
out_pages[:, None] * self.page_size
|
||||
+ torch.arange(self.page_size, device=self.device)
|
||||
).reshape(-1)
|
||||
return out_indices
|
||||
|
||||
def free(self, free_index: torch.Tensor):
|
||||
if free_index.numel() == 0:
|
||||
return
|
||||
|
||||
if self.is_not_in_free_group:
|
||||
free_page_indices = torch.unique(free_index // self.page_size)
|
||||
self._append_pages_to_owner_buckets(
|
||||
release=self.need_sort, pages=free_page_indices
|
||||
)
|
||||
else:
|
||||
self.free_group.append(free_index)
|
||||
|
||||
if self.debug_mode and self._owner_free_pages is not None:
|
||||
free_pages = self.free_pages
|
||||
assert len(torch.unique(free_pages)) == len(free_pages)
|
||||
|
||||
def compute_owner_lane_stats(
|
||||
self,
|
||||
@@ -643,17 +870,7 @@ class CPSharedPagedTokenToKVPoolAllocator(PagedTokenToKVPoolAllocator):
|
||||
)
|
||||
required[owner] += 1
|
||||
|
||||
free_pages = self.free_pages
|
||||
if len(self.release_pages) > 0:
|
||||
free_pages = torch.cat((free_pages, self.release_pages))
|
||||
available = [
|
||||
int(
|
||||
(
|
||||
torch.remainder(free_pages - 1, self.cp_size) == owner
|
||||
).sum().item()
|
||||
)
|
||||
for owner in range(self.cp_size)
|
||||
]
|
||||
available = self._owner_available_counts()
|
||||
deficits = [
|
||||
max(0, required_count - available_count)
|
||||
for required_count, available_count in zip(required, available)
|
||||
@@ -692,108 +909,73 @@ class CPSharedPagedTokenToKVPoolAllocator(PagedTokenToKVPoolAllocator):
|
||||
|
||||
def _select_owner_free_pages_prefer_contiguous(
|
||||
self,
|
||||
owner_mask: torch.Tensor,
|
||||
owner_pages: 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
|
||||
)
|
||||
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, torch.Tensor, torch.Tensor]]:
|
||||
) -> Optional[tuple[torch.Tensor, list[int], list[int]]]:
|
||||
if not page_compute_owners:
|
||||
return (
|
||||
torch.empty((0,), dtype=torch.int64, device=self.device),
|
||||
torch.zeros_like(self.free_pages, dtype=torch.bool),
|
||||
torch.zeros_like(self.release_pages, dtype=torch.bool),
|
||||
[0 for _ in range(self.cp_size)],
|
||||
[0 for _ in range(self.cp_size)],
|
||||
)
|
||||
|
||||
required_by_owner = [0 for _ in range(self.cp_size)]
|
||||
for owner in page_compute_owners:
|
||||
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_mask = torch.zeros_like(self.free_pages, dtype=torch.bool)
|
||||
selected_release_mask = torch.zeros_like(self.release_pages, dtype=torch.bool)
|
||||
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
|
||||
|
||||
owner_mask = torch.remainder(self.free_pages - 1, self.cp_size) == owner
|
||||
selected_owner_free_mask = (
|
||||
self._select_owner_free_pages_prefer_contiguous(
|
||||
owner_mask, required_count
|
||||
)
|
||||
free_bucket = self._owner_free_pages[owner]
|
||||
selected_owner_pages = self._select_owner_free_pages_prefer_contiguous(
|
||||
free_bucket, required_count
|
||||
)
|
||||
selected_owner_pages = self.free_pages[selected_owner_free_mask]
|
||||
|
||||
remaining_count = required_count - selected_owner_pages.numel()
|
||||
free_count = int(selected_owner_pages.numel())
|
||||
remaining_count = required_count - free_count
|
||||
if remaining_count > 0:
|
||||
release_owner_mask = (
|
||||
torch.remainder(self.release_pages - 1, self.cp_size) == owner
|
||||
)
|
||||
selected_owner_release_mask = release_owner_mask & (
|
||||
torch.cumsum(release_owner_mask.to(torch.int64), dim=0)
|
||||
<= remaining_count
|
||||
)
|
||||
selected_owner_release_pages = self.release_pages[
|
||||
selected_owner_release_mask
|
||||
]
|
||||
if remaining_count > selected_owner_release_pages.numel():
|
||||
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_mask |= selected_owner_release_mask
|
||||
selected_release_counts[owner] = remaining_count
|
||||
|
||||
selected_free_counts[owner] = free_count
|
||||
lane_pages[owner] = selected_owner_pages
|
||||
selected_free_mask |= selected_owner_free_mask
|
||||
|
||||
selected_pages = []
|
||||
lane_offsets = [0 for _ in range(self.cp_size)]
|
||||
for owner in page_compute_owners:
|
||||
lane_offset = lane_offsets[owner]
|
||||
selected_pages.append(lane_pages[owner][lane_offset])
|
||||
lane_offsets[owner] = lane_offset + 1
|
||||
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 (
|
||||
torch.stack(selected_pages).to(torch.int64),
|
||||
selected_free_mask,
|
||||
selected_release_mask,
|
||||
selected_pages,
|
||||
selected_free_counts,
|
||||
selected_release_counts,
|
||||
)
|
||||
|
||||
def alloc_extend_compute_owner(
|
||||
@@ -831,7 +1013,7 @@ class CPSharedPagedTokenToKVPoolAllocator(PagedTokenToKVPoolAllocator):
|
||||
selected = self._select_compute_owner_pages(page_compute_owners)
|
||||
if selected is None:
|
||||
return None
|
||||
selected_pages, selected_free_mask, selected_release_mask = selected
|
||||
selected_pages, selected_free_counts, selected_release_counts = selected
|
||||
|
||||
out_indices = torch.empty(
|
||||
(extend_num_tokens,), dtype=torch.int64, device=self.device
|
||||
@@ -846,8 +1028,12 @@ class CPSharedPagedTokenToKVPoolAllocator(PagedTokenToKVPoolAllocator):
|
||||
self.device,
|
||||
)
|
||||
|
||||
self.free_pages = self.free_pages[~selected_free_mask]
|
||||
self.release_pages = self.release_pages[~selected_release_mask]
|
||||
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
|
||||
)
|
||||
|
||||
if self.debug_mode:
|
||||
assert len(torch.unique(out_indices)) == len(out_indices)
|
||||
@@ -893,15 +1079,19 @@ class CPSharedPagedTokenToKVPoolAllocator(PagedTokenToKVPoolAllocator):
|
||||
selected = self._select_compute_owner_pages(page_compute_owners)
|
||||
if selected is None:
|
||||
return None
|
||||
selected_pages, selected_free_mask, selected_release_mask = selected
|
||||
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.free_pages = self.free_pages[~selected_free_mask]
|
||||
self.release_pages = self.release_pages[~selected_release_mask]
|
||||
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
|
||||
)
|
||||
if self.debug_mode:
|
||||
assert torch.unique(out_indices).numel() == out_indices.numel()
|
||||
check_owners = torch.remainder(selected_pages - 1, self.cp_size)
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -1,20 +1,77 @@
|
||||
import sys
|
||||
import types
|
||||
import unittest
|
||||
from unittest.mock import patch
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
# This CPU unit file must not import the native sgl_kernel package during test
|
||||
# collection. Some remote images abort inside sgl_kernel's architecture-specific
|
||||
# loader instead of raising ImportError/RuntimeError, so a try/except fallback is
|
||||
# not sufficient here. Install a minimal stub before importing sglang helpers.
|
||||
if "sgl_kernel" not in sys.modules:
|
||||
sgl_kernel_stub = types.ModuleType("sgl_kernel")
|
||||
sgl_kernel_stub.__file__ = "sgl_kernel_stub.py"
|
||||
sgl_kernel_stub.__path__ = []
|
||||
|
||||
def _sgl_kernel_getattr(name):
|
||||
if name.startswith("__"):
|
||||
raise AttributeError(name)
|
||||
fn = lambda *args, **kwargs: None
|
||||
setattr(sgl_kernel_stub, name, fn)
|
||||
return fn
|
||||
|
||||
sgl_kernel_stub.__getattr__ = _sgl_kernel_getattr
|
||||
sys.modules["sgl_kernel"] = sgl_kernel_stub
|
||||
|
||||
if "sgl_kernel.kvcacheio" not in sys.modules:
|
||||
kvcacheio_stub = types.ModuleType("sgl_kernel.kvcacheio")
|
||||
kvcacheio_stub.__file__ = "sgl_kernel_kvcacheio_stub.py"
|
||||
|
||||
def _kvcacheio_getattr(name):
|
||||
if name.startswith("__"):
|
||||
raise AttributeError(name)
|
||||
fn = lambda *args, **kwargs: None
|
||||
setattr(kvcacheio_stub, name, fn)
|
||||
return fn
|
||||
|
||||
kvcacheio_stub.__getattr__ = _kvcacheio_getattr
|
||||
sys.modules["sgl_kernel.kvcacheio"] = kvcacheio_stub
|
||||
|
||||
if "sgl_kernel.quantization" not in sys.modules:
|
||||
quantization_stub = types.ModuleType("sgl_kernel.quantization")
|
||||
quantization_stub.__file__ = "sgl_kernel_quantization_stub.py"
|
||||
|
||||
def _quantization_getattr(name):
|
||||
if name.startswith("__"):
|
||||
raise AttributeError(name)
|
||||
fn = lambda *args, **kwargs: None
|
||||
setattr(quantization_stub, name, fn)
|
||||
return fn
|
||||
|
||||
quantization_stub.__getattr__ = _quantization_getattr
|
||||
sys.modules["sgl_kernel.quantization"] = quantization_stub
|
||||
|
||||
_sgl_kernel_lib = torch.library.Library("sgl_kernel", "FRAGMENT")
|
||||
try:
|
||||
_sgl_kernel_lib.define(
|
||||
for _schema in (
|
||||
"sgl_per_token_group_quant_8bit(Tensor input, Tensor(a!) output_q, Tensor(b!) output_s, int group_size, float eps, float fp8_min, float fp8_max, bool scale_ue8m0) -> ()",
|
||||
"sgl_per_token_group_quant_fp8(Tensor input, Tensor(a!) output_q, Tensor(b!) output_s, int group_size, float eps, float fp8_min, float fp8_max, bool scale_ue8m0) -> ()",
|
||||
"sgl_per_token_quant_fp8(Tensor input, Tensor(a!) output_q, Tensor(b!) output_s) -> ()",
|
||||
"fp8_scaled_mm(Tensor mat_a, Tensor mat_b, Tensor scales_a, Tensor scales_b, ScalarType out_dtype, Tensor? bias=None) -> Tensor",
|
||||
"fp8_blockwise_scaled_mm(Tensor mat_a, Tensor mat_b, Tensor scales_a, Tensor scales_b, ScalarType out_dtype) -> Tensor",
|
||||
(
|
||||
"moe_fused_gate(Tensor input_tensor, Tensor? bias, int num_expert_group, "
|
||||
"int topk_group, int topk, int num_fused_shared_experts, "
|
||||
"float routed_scaling_factor, bool apply_routed_scaling_factor_on_output) "
|
||||
"-> (Tensor, Tensor)"
|
||||
)
|
||||
except RuntimeError as exc:
|
||||
if "already" not in str(exc).lower() and "duplicate" not in str(exc).lower():
|
||||
raise
|
||||
),
|
||||
):
|
||||
try:
|
||||
_sgl_kernel_lib.define(_schema)
|
||||
except RuntimeError as exc:
|
||||
if "already" not in str(exc).lower() and "duplicate" not in str(exc).lower():
|
||||
raise
|
||||
|
||||
from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
@@ -392,7 +449,7 @@ 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):
|
||||
def test_owner_lane_selection_uses_sorted_contiguous_bucket_prefix(self):
|
||||
from sglang.srt.mem_cache.allocator import CPSharedPagedTokenToKVPoolAllocator
|
||||
|
||||
page_size = 64
|
||||
@@ -418,11 +475,13 @@ class TestCPSharedPagedAllocator(CustomTestCase):
|
||||
|
||||
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.tolist(), [1, 5, 9])
|
||||
self.assertEqual(
|
||||
((logical_pages - 1) % cp_size).tolist(),
|
||||
[0, 0, 0],
|
||||
)
|
||||
physical_pages = torch.div(logical_pages - 1, cp_size, rounding_mode="floor")
|
||||
self.assertTrue(torch.all(physical_pages[1:] - physical_pages[:-1] == 1))
|
||||
for selected_page in logical_pages.tolist():
|
||||
self.assertNotIn(selected_page, allocator.free_pages.tolist())
|
||||
|
||||
|
||||
Reference in New Issue
Block a user