diff --git a/benchmark/hicache/bench_cp_hicache_allocator_overhead.py b/benchmark/hicache/bench_cp_hicache_allocator_overhead.py index cf956efc6..3d4e1a7b1 100644 --- a/benchmark/hicache/bench_cp_hicache_allocator_overhead.py +++ b/benchmark/hicache/bench_cp_hicache_allocator_overhead.py @@ -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", diff --git a/docs/advanced_features/nsa_prefill_cp_page_aligned_cache_contract.md b/docs/advanced_features/nsa_prefill_cp_page_aligned_cache_contract.md index 81994802e..908413b26 100644 --- a/docs/advanced_features/nsa_prefill_cp_page_aligned_cache_contract.md +++ b/docs/advanced_features/nsa_prefill_cp_page_aligned_cache_contract.md @@ -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. diff --git a/python/sglang/srt/mem_cache/allocator.py b/python/sglang/srt/mem_cache/allocator.py index 25e278df4..7e8c9eefb 100644 --- a/python/sglang/srt/mem_cache/allocator.py +++ b/python/sglang/srt/mem_cache/allocator.py @@ -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) diff --git a/test/registered/unit/benchmark/test_cp_hicache_allocator_bench.py b/test/registered/unit/benchmark/test_cp_hicache_allocator_bench.py index 2b8b1ef56..1c98b7a50 100644 --- a/test/registered/unit/benchmark/test_cp_hicache_allocator_bench.py +++ b/test/registered/unit/benchmark/test_cp_hicache_allocator_bench.py @@ -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 diff --git a/test/registered/unit/mem_cache/test_cp_shared_kv_layout.py b/test/registered/unit/mem_cache/test_cp_shared_kv_layout.py index 8757ef4a7..3b3127279 100644 --- a/test/registered/unit/mem_cache/test_cp_shared_kv_layout.py +++ b/test/registered/unit/mem_cache/test_cp_shared_kv_layout.py @@ -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())