Reduce CP HiCache L2 allocator scan cost

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

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

Constraint: CP HiCache host allocations are page-shaped, but existing callers may still read free_slots directly.
Rejected: Sort and scan free_slots on each alloc_contiguous_preferred call | measured ms-level CPU overhead on 220GB-equivalent metadata.
Rejected: Remove free_slots compatibility | storage/tests still rely on the public tensor surface.
Confidence: medium
Scope-risk: moderate
Directive: Do not reintroduce per-allocation full free_slots scans on HostKVCache; preserve page-extent metadata or benchmark before changing allocator shape.
Tested: Local py_compile for memory_pool_host.py, allocator benchmark, and related tests.
Tested: Local test_cp_hicache_allocator_bench.py 10 passed.
Tested: Remote g0034 test_hicache_controller_cp.py 67 passed; test_cp_hicache_allocator_bench.py 10 passed.
Tested: Remote 220GB-equivalent host_churn benchmark: contiguous path reduced from ms-level to ~30-292us p50 depending on fragmentation.
Not-tested: Full CUDA ETE run after allocator change.
Not-tested: Production long-run fragmentation behavior under live traffic.

python/sglang/srt/mem_cache/memory_pool_host.py

@@ -1,4 +1,6 @@
 import abc
+import bisect
+import heapq
 import logging
 import threading
 from collections import defaultdict
@@ -314,6 +316,249 @@ class HostKVCache(abc.ABC):
         """
         raise NotImplementedError()
 
+    @property
+    def free_slots(self) -> torch.Tensor:
+        if getattr(self, "_free_slots_dirty", False):
+            self._free_slots = self._materialize_free_slots_from_extents()
+            self._free_slots_dirty = False
+        return self._free_slots
+
+    @free_slots.setter
+    def free_slots(self, value: torch.Tensor) -> None:
+        self._free_slots = value.cpu().to(dtype=torch.int64).contiguous()
+        self._rebuild_free_extent_index_from_slots(self._free_slots)
+
+    def _reset_free_extent_index(self, *, enabled: bool, token_count: int) -> None:
+        self._free_extent_index_enabled = bool(enabled)
+        self._free_token_count = int(token_count)
+        self._free_extents_by_start: dict[int, int] = {}
+        self._free_extent_starts: list[int] = []
+        self._free_extent_heap: list[tuple[int, int]] = []
+        self._free_slots_dirty = False
+
+    def _rebuild_free_extent_index_from_slots(self, free_slots: torch.Tensor) -> None:
+        token_count = int(free_slots.numel())
+        self._reset_free_extent_index(enabled=False, token_count=token_count)
+        page_size = int(self.page_size)
+        if token_count == 0:
+            self._free_extent_index_enabled = True
+            return
+        if token_count % page_size != 0:
+            logger.warning(
+                "[HiCache-L2-allocator] disabling extent index for non-page-shaped free_slots: tokens=%d page_size=%d",
+                token_count,
+                page_size,
+            )
+            return
+
+        page_slots = free_slots.view(-1, page_size)
+        offsets = torch.arange(page_size, dtype=page_slots.dtype)
+        if not bool(torch.all(page_slots == (page_slots[:, :1] + offsets)).item()):
+            logger.warning(
+                "[HiCache-L2-allocator] disabling extent index for non-contiguous page slots"
+            )
+            return
+
+        self._free_extent_index_enabled = True
+        pages = torch.div(page_slots[:, 0], page_size, rounding_mode="floor").tolist()
+        if not pages:
+            return
+        run_start = int(pages[0])
+        run_len = 1
+        prev = run_start
+        for page in pages[1:]:
+            page = int(page)
+            if page == prev + 1:
+                run_len += 1
+            else:
+                self._insert_free_extent(run_start, run_len)
+                run_start = page
+                run_len = 1
+            prev = page
+        self._insert_free_extent(run_start, run_len)
+
+    def _materialize_free_slots_from_extents(self) -> torch.Tensor:
+        page_size = int(self.page_size)
+        chunks: list[torch.Tensor] = []
+        offsets = torch.arange(page_size, dtype=torch.int64)
+        for start in self._free_extent_starts:
+            length = self._free_extents_by_start.get(start, 0)
+            if length <= 0:
+                continue
+            pages = torch.arange(start, start + length, dtype=torch.int64)
+            chunks.append((pages[:, None] * page_size + offsets[None, :]).reshape(-1))
+        if not chunks:
+            return torch.empty((0,), dtype=torch.int64)
+        return torch.cat(chunks).contiguous()
+
+    def _insert_free_extent(self, start_page: int, page_count: int) -> None:
+        if page_count <= 0:
+            return
+        start_page = int(start_page)
+        page_count = int(page_count)
+        end_page = start_page + page_count
+
+        pos = bisect.bisect_left(self._free_extent_starts, start_page)
+        if pos > 0:
+            prev_start = self._free_extent_starts[pos - 1]
+            prev_len = self._free_extents_by_start[prev_start]
+            prev_end = prev_start + prev_len
+            if prev_end > start_page:
+                raise RuntimeError(
+                    "[HiCache-L2-allocator] double free or overlapping host extent"
+                )
+            if prev_end == start_page:
+                start_page = prev_start
+                page_count += prev_len
+                del self._free_extents_by_start[prev_start]
+                del self._free_extent_starts[pos - 1]
+                pos -= 1
+        if pos < len(self._free_extent_starts):
+            next_start = self._free_extent_starts[pos]
+            next_len = self._free_extents_by_start[next_start]
+            if end_page > next_start:
+                raise RuntimeError(
+                    "[HiCache-L2-allocator] double free or overlapping host extent"
+                )
+            if end_page == next_start:
+                page_count += next_len
+                del self._free_extents_by_start[next_start]
+                del self._free_extent_starts[pos]
+
+        self._free_extents_by_start[start_page] = page_count
+        bisect.insort(self._free_extent_starts, start_page)
+        heapq.heappush(self._free_extent_heap, (-page_count, start_page))
+
+    def _pop_largest_free_extent(self) -> Optional[tuple[int, int]]:
+        while self._free_extent_heap:
+            neg_len, start = heapq.heappop(self._free_extent_heap)
+            length = -int(neg_len)
+            current = self._free_extents_by_start.get(start)
+            if current == length:
+                del self._free_extents_by_start[start]
+                pos = bisect.bisect_left(self._free_extent_starts, start)
+                if (
+                    pos < len(self._free_extent_starts)
+                    and self._free_extent_starts[pos] == start
+                ):
+                    del self._free_extent_starts[pos]
+                return start, length
+        return None
+
+    def _peek_largest_free_extent(self) -> Optional[tuple[int, int]]:
+        while self._free_extent_heap:
+            neg_len, start = self._free_extent_heap[0]
+            length = -int(neg_len)
+            if self._free_extents_by_start.get(start) == length:
+                return start, length
+            heapq.heappop(self._free_extent_heap)
+        return None
+
+    def _page_ids_to_token_indices(self, page_ids: list[int]) -> torch.Tensor:
+        if not page_ids:
+            return torch.empty((0,), dtype=torch.int64)
+        page_size = int(self.page_size)
+        pages = torch.tensor(page_ids, dtype=torch.int64)
+        offsets = torch.arange(page_size, dtype=torch.int64)
+        return (pages[:, None] * page_size + offsets[None, :]).reshape(-1).contiguous()
+
+    def _alloc_from_free_extents(
+        self, need_size: int, *, require_single_run: bool
+    ) -> Optional[torch.Tensor]:
+        page_size = int(self.page_size)
+        need_pages = int(need_size) // page_size
+        if need_size > self.available_size():
+            return None
+        if need_pages == 0:
+            return torch.empty((0,), dtype=torch.int64)
+
+        if require_single_run:
+            largest = self._peek_largest_free_extent()
+            if largest is None or largest[1] < need_pages:
+                return None
+            start, length = self._pop_largest_free_extent()
+            selected_pages: list[int] = []
+            selected_pages.extend(range(start, start + need_pages))
+            remaining = length - need_pages
+            if remaining > 0:
+                self._insert_free_extent(start + need_pages, remaining)
+        else:
+            selected_pages = self._alloc_fragmented_from_free_extents(need_pages)
+
+        self._free_token_count -= need_size
+        self._free_slots_dirty = True
+        return self._page_ids_to_token_indices(selected_pages)
+
+    def _alloc_fragmented_from_free_extents(self, need_pages: int) -> list[int]:
+        remaining_pages = int(need_pages)
+        selected_pages: list[int] = []
+        consumed_count = 0
+        replacement: Optional[tuple[int, int]] = None
+
+        for start in self._free_extent_starts:
+            if remaining_pages <= 0:
+                break
+            length = self._free_extents_by_start[start]
+            take = min(length, remaining_pages)
+            selected_pages.extend(range(start, start + take))
+            remaining_pages -= take
+            if take == length:
+                consumed_count += 1
+                continue
+            replacement = (start + take, length - take)
+            break
+
+        if remaining_pages > 0:
+            raise RuntimeError(
+                "[HiCache-L2-allocator] extent index underflow during fragmented allocation"
+            )
+
+        consumed_starts = self._free_extent_starts[:consumed_count]
+        for start in consumed_starts:
+            del self._free_extents_by_start[start]
+        del self._free_extent_starts[:consumed_count]
+
+        if replacement is not None:
+            old_start = self._free_extent_starts[0]
+            del self._free_extents_by_start[old_start]
+            self._free_extent_starts[0] = replacement[0]
+            self._free_extents_by_start[replacement[0]] = replacement[1]
+            heapq.heappush(self._free_extent_heap, (-replacement[1], replacement[0]))
+
+        return selected_pages
+
+    def _page_runs_from_token_indices(
+        self, indices: torch.Tensor
+    ) -> Optional[list[tuple[int, int]]]:
+        indices = indices.cpu().to(dtype=torch.int64).contiguous()
+        token_count = int(indices.numel())
+        page_size = int(self.page_size)
+        if token_count == 0:
+            return []
+        if token_count % page_size != 0:
+            return None
+        page_slots = indices.view(-1, page_size)
+        offsets = torch.arange(page_size, dtype=page_slots.dtype)
+        if not bool(torch.all(page_slots == (page_slots[:, :1] + offsets)).item()):
+            return None
+
+        pages = torch.div(page_slots[:, 0], page_size, rounding_mode="floor").tolist()
+        runs: list[tuple[int, int]] = []
+        run_start = int(pages[0])
+        run_len = 1
+        prev = run_start
+        for page in pages[1:]:
+            page = int(page)
+            if page == prev + 1:
+                run_len += 1
+            else:
+                runs.append((run_start, run_len))
+                run_start = page
+                run_len = 1
+            prev = page
+        runs.append((run_start, run_len))
+        return runs
+
     @synchronized
     def clear(self):
         # Initialize memory states and tracking structures.
@@ -323,6 +568,8 @@ class HostKVCache(abc.ABC):
         self.free_slots = torch.arange(self.size, dtype=torch.int64)
 
     def available_size(self):
+        if hasattr(self, "_free_token_count"):
+            return int(self._free_token_count)
         return len(self.free_slots)
 
     @synchronized
@@ -332,6 +579,10 @@ class HostKVCache(abc.ABC):
         ), "The requested size should be a multiple of the page size."
         if need_size > self.available_size():
             return None
+        if getattr(self, "_free_extent_index_enabled", False):
+            return self._alloc_from_free_extents(
+                need_size, require_single_run=False
+            )
 
         select_index = self.free_slots[:need_size]
         self.free_slots = self.free_slots[need_size:]
@@ -356,6 +607,15 @@ class HostKVCache(abc.ABC):
             return None
         if need_size == 0:
             return self.alloc(need_size)
+        if getattr(self, "_free_extent_index_enabled", False):
+            select_index = self._alloc_from_free_extents(
+                need_size, require_single_run=True
+            )
+            if select_index is not None:
+                return select_index
+            return self._alloc_from_free_extents(
+                need_size, require_single_run=False
+            )
 
         fifo_prefix = self.free_slots[:need_size]
         expected_prefix = fifo_prefix[:1] + torch.arange(
@@ -422,6 +682,20 @@ class HostKVCache(abc.ABC):
 
     @synchronized
     def free(self, indices: torch.Tensor) -> int:
+        if getattr(self, "_free_extent_index_enabled", False):
+            runs = self._page_runs_from_token_indices(indices)
+            if runs is not None:
+                for start, length in runs:
+                    self._insert_free_extent(start, length)
+                self._free_token_count += int(indices.numel())
+                self._free_slots_dirty = True
+                return len(indices)
+            logger.warning(
+                "[HiCache-L2-allocator] disabling extent index for non-page-shaped free indices: tokens=%d page_size=%d",
+                int(indices.numel()),
+                int(self.page_size),
+            )
+            self._free_extent_index_enabled = False
         self.free_slots = torch.cat([self.free_slots, indices.cpu()])
         return len(indices)