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
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user