leavelet
1d630def95
In CP=8 + NSA-shared-KV + HiCache disagg-prefill, cache-hit prefill produced incoherent decode output. Cold prefill on CP was correct; pure CP without HiCache was correct. The bug lived at the HiCache load_cp / device-alloc interface. Root cause: cache_controller.load_cp called the plain mem_pool_device_allocator.alloc(logical_len), which returns logical pages with no CP owner-pattern preservation. Cold prefill instead uses alloc_extend_compute_owner with a zigzag owner pattern from build_in_seq_page_compute_owners. The saved CpHiCacheNodeMetadata.owned_positions records WHICH POSITIONS in the write-time alloc were owned by this rank. At load time, those same positions are applied to a new alloc whose per-position owner pattern is arbitrary -- each rank loads its host bytes into physical slots whose corresponding logical page is owned by a DIFFERENT rank. Attention's materialize_shared_token_kv_buffer reads from the owner's physical slot, which was never loaded. Result: garbage. Fix: - CpHiCacheNodeMetadata gains two required fields: page_owners (int8 per logical page, identical on all CP ranks) and page_size. __post_init__ validates; split() bisects page_owners by page index with a page-alignment check. - _write_cp derives page_owners from device_indices (page-first slot of each page -> logical page id -> layout.owner_for_logical_pages) and stores in both metadata-construction sites (zero-owned and normal). - New CPSharedPagedTokenToKVPoolAllocator.alloc_pages_with_owners() reuses _select_compute_owner_pages (with its tai-kernel fast path) and returns page-contiguous token locs whose per-page owner sequence equals the input. - load_cp now concats page_owners across nodes_to_load and calls alloc_pages_with_owners. On None (lane exhausted) the caller hits the retry-with-eviction path; further failure returns None and degrades to cache miss. No silent fallback to plain alloc -- that recreated the bug. - load_back retry path now calls _evict_for_compute_owner_lanes (module-top import) instead of plain evict(); this targets the deficit lane and gives the next alloc attempt a chance to satisfy it. - envs import moved to module top in cache_controller.py per code-review feedback. Removed an over-defensive owned_check.all().item() in load_cp that would have re-introduced the host-sync anti-pattern 97a9f850c removed -- the invariant is already guaranteed by alloc_pages_with_owners. Tests: 40 existing CpHiCacheNodeMetadata constructions migrated to pass the new required fields. 9 new metadata tests (validators + split page-alignment). 10 new allocator tests in test_alloc_pages_with_owners.py covering input-order preservation, lane exhaustion, release_pages fallback, debug-mode invariant. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
818 lines
28 KiB
Python
818 lines
28 KiB
Python
from __future__ import annotations
|
|
|
|
"""
|
|
Copyright 2025 SGLang Team
|
|
Licensed under the Apache License, Version 2.0 (the "License");
|
|
you may not use this file except in compliance with the License.
|
|
You may obtain a copy of the License at
|
|
|
|
http://www.apache.org/licenses/LICENSE-2.0
|
|
|
|
Unless required by applicable law or agreed to in writing, software
|
|
distributed under the License is distributed on an "AS IS" BASIS,
|
|
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
See the License for the specific language governing permissions and
|
|
limitations under the License.
|
|
"""
|
|
|
|
"""
|
|
Page-aligned memory pool.
|
|
"""
|
|
|
|
import abc
|
|
from typing import TYPE_CHECKING, List, Optional
|
|
|
|
import torch
|
|
import triton
|
|
|
|
from sglang.srt.environ import envs
|
|
import triton.language as tl
|
|
|
|
from sglang.srt.utils import get_bool_env_var, get_num_new_pages, next_power_of_2
|
|
|
|
_SORT_NVTX_ENABLED = envs.SGLANG_DEBUG_SORT_NVTX.get()
|
|
|
|
if TYPE_CHECKING:
|
|
from sglang.srt.mem_cache.memory_pool import KVCache
|
|
|
|
|
|
def _debug_sort_nvtx_enabled() -> bool:
|
|
return _SORT_NVTX_ENABLED
|
|
|
|
|
|
class BaseTokenToKVPoolAllocator(abc.ABC):
|
|
@abc.abstractmethod
|
|
def __init__(
|
|
self,
|
|
size: int,
|
|
page_size: int,
|
|
dtype: torch.dtype,
|
|
device: str,
|
|
kvcache: KVCache,
|
|
need_sort: bool,
|
|
):
|
|
self.size = size
|
|
self.page_size = page_size
|
|
self.dtype = dtype
|
|
self.device = device
|
|
self._kvcache = kvcache
|
|
self.need_sort = need_sort
|
|
|
|
self.free_pages = None
|
|
self.release_pages = None
|
|
self.is_not_in_free_group = True
|
|
self.free_group = []
|
|
|
|
def debug_print(self) -> str:
|
|
return ""
|
|
|
|
def available_size(self):
|
|
return (len(self.free_pages) + len(self.release_pages)) * self.page_size
|
|
|
|
def immediate_available_pages(self) -> int:
|
|
if self.free_pages is None:
|
|
return 0
|
|
return len(self.free_pages)
|
|
|
|
def deferred_available_pages(self) -> int:
|
|
if self.release_pages is None:
|
|
return 0
|
|
return len(self.release_pages)
|
|
|
|
def immediate_available_size(self) -> int:
|
|
if self.free_pages is None:
|
|
return self.available_size()
|
|
return self.immediate_available_pages() * self.page_size
|
|
|
|
def deferred_available_size(self) -> int:
|
|
if self.release_pages is None:
|
|
return 0
|
|
return self.deferred_available_pages() * self.page_size
|
|
|
|
def allocator_state_str(self) -> str:
|
|
if self.free_pages is None:
|
|
return f"allocator_available_size={self.available_size()}"
|
|
return (
|
|
"allocator_available_size="
|
|
f"{self.available_size()} "
|
|
f"(allocator_free_size={self.immediate_available_size()} "
|
|
f"[free_pages={self.immediate_available_pages()}] + "
|
|
f"allocator_release_size={self.deferred_available_size()} "
|
|
f"[release_pages={self.deferred_available_pages()}])"
|
|
)
|
|
|
|
def get_kvcache(self):
|
|
return self._kvcache
|
|
|
|
def restore_state(self, state):
|
|
self.free_pages, self.release_pages = state
|
|
|
|
def backup_state(self):
|
|
return (self.free_pages, self.release_pages)
|
|
|
|
def free_group_begin(self):
|
|
self.is_not_in_free_group = False
|
|
self.free_group = []
|
|
|
|
def free_group_end(self):
|
|
self.is_not_in_free_group = True
|
|
if self.free_group:
|
|
self.free(torch.cat(self.free_group))
|
|
|
|
def merge_and_sort_free(self):
|
|
if len(self.release_pages) > 0:
|
|
num_free_pages = len(self.free_pages)
|
|
num_release_pages = len(self.release_pages)
|
|
self.free_pages = torch.cat((self.free_pages, self.release_pages))
|
|
if _debug_sort_nvtx_enabled():
|
|
torch.cuda.nvtx.range_push(
|
|
f"KV_ALLOCATOR:merge_and_sort_free:torch.sort free_pages={num_free_pages} release_pages={num_release_pages}"
|
|
)
|
|
try:
|
|
self.free_pages, _ = torch.sort(self.free_pages)
|
|
finally:
|
|
torch.cuda.nvtx.range_pop()
|
|
else:
|
|
self.free_pages, _ = torch.sort(self.free_pages)
|
|
self.release_pages = torch.empty(
|
|
(0,), dtype=self.release_pages.dtype, device=self.device
|
|
)
|
|
|
|
def get_cpu_copy(self, *args, **kwargs):
|
|
# FIXME: reuse the get_cpu_copy after paged allocator is implemented
|
|
raise NotImplementedError()
|
|
|
|
def load_cpu_copy(self, *args, **kwargs):
|
|
# FIXME: reuse the load_cpu_copy after paged allocator is implemented
|
|
raise NotImplementedError()
|
|
|
|
def alloc_extend(self, *args, **kwargs):
|
|
raise NotImplementedError("alloc_extend is only for paged allocator")
|
|
|
|
def alloc_decode(self, *args, **kwargs):
|
|
raise NotImplementedError("alloc_decode is only for paged allocator")
|
|
|
|
@abc.abstractmethod
|
|
def clear(self):
|
|
raise NotImplementedError()
|
|
|
|
@abc.abstractmethod
|
|
def alloc(self, need_size: int):
|
|
raise NotImplementedError()
|
|
|
|
@abc.abstractmethod
|
|
def free(self, free_index: torch.Tensor):
|
|
raise NotImplementedError()
|
|
|
|
|
|
class TokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
|
|
"""An allocator managing the indices to kv cache data."""
|
|
|
|
def __init__(
|
|
self,
|
|
size: int,
|
|
dtype: torch.dtype,
|
|
device: str,
|
|
kvcache: KVCache,
|
|
need_sort: bool,
|
|
):
|
|
super().__init__(size, 1, dtype, device, kvcache, need_sort)
|
|
self.clear()
|
|
|
|
def clear(self):
|
|
# The padded slot 0 is used for writing dummy outputs from padded tokens.
|
|
self.free_pages = torch.arange(
|
|
1, self.size + 1, dtype=torch.int64, device=self.device
|
|
)
|
|
self.is_not_in_free_group = True
|
|
self.free_group = []
|
|
self.release_pages = torch.empty((0,), dtype=torch.int64, device=self.device)
|
|
|
|
def available_size(self):
|
|
# To avoid minor "len(free_pages) * 1" overhead
|
|
return len(self.free_pages) + len(self.release_pages)
|
|
|
|
def alloc(self, need_size: int):
|
|
if self.need_sort and need_size > len(self.free_pages):
|
|
self.merge_and_sort_free()
|
|
|
|
if need_size > len(self.free_pages):
|
|
return None
|
|
|
|
select_index = self.free_pages[:need_size]
|
|
self.free_pages = self.free_pages[need_size:]
|
|
return select_index
|
|
|
|
def free(self, free_index: torch.Tensor):
|
|
if free_index.numel() == 0:
|
|
return
|
|
|
|
if self.is_not_in_free_group:
|
|
if self.need_sort:
|
|
self.release_pages = torch.cat((self.release_pages, free_index))
|
|
else:
|
|
self.free_pages = torch.cat((self.free_pages, free_index))
|
|
else:
|
|
self.free_group.append(free_index)
|
|
|
|
def get_cpu_copy(self, indices):
|
|
return self._kvcache.get_cpu_copy(indices)
|
|
|
|
def load_cpu_copy(self, kv_cache_cpu, indices):
|
|
return self._kvcache.load_cpu_copy(kv_cache_cpu, indices)
|
|
|
|
|
|
def alloc_extend_naive(
|
|
prefix_lens,
|
|
seq_lens,
|
|
last_loc,
|
|
free_pages,
|
|
out_indices,
|
|
page_size,
|
|
device,
|
|
):
|
|
extend_lens = seq_lens - prefix_lens
|
|
end_pos = torch.cumsum(extend_lens, 0)
|
|
start_pos = end_pos - extend_lens
|
|
num_new_pages = (seq_lens + page_size - 1) // page_size - (
|
|
prefix_lens + page_size - 1
|
|
) // page_size
|
|
num_full_new_pages = (seq_lens) // page_size - (
|
|
prefix_lens + page_size - 1
|
|
) // page_size
|
|
need_page = num_new_pages - num_full_new_pages
|
|
end_new_pages = torch.cumsum(num_new_pages, 0)
|
|
start_new_pages = end_new_pages - num_new_pages
|
|
pos_in_page = torch.arange(page_size, device=device, dtype=torch.int32)
|
|
for i in range(len(prefix_lens)):
|
|
num1 = (
|
|
min(
|
|
seq_lens[i],
|
|
(prefix_lens[i] + page_size - 1) // page_size * page_size,
|
|
)
|
|
- prefix_lens[i]
|
|
)
|
|
if num1:
|
|
out_indices[start_pos[i] : start_pos[i] + num1] = (
|
|
last_loc[i] + 1 + pos_in_page[:num1].view(-1)
|
|
)
|
|
|
|
if prefix_lens[i] + num1 == seq_lens[i]:
|
|
continue
|
|
|
|
num2 = (
|
|
seq_lens[i] // page_size - (prefix_lens[i] + page_size - 1) // page_size
|
|
) * page_size
|
|
if num2:
|
|
pages = (
|
|
free_pages[start_new_pages[i] : end_new_pages[i] - need_page[i]]
|
|
* page_size
|
|
)
|
|
out_indices[start_pos[i] + num1 : start_pos[i] + num1 + num2] = (
|
|
pages.view(-1, 1) + pos_in_page.view(1, -1)
|
|
).view(-1)
|
|
|
|
if prefix_lens[i] + num1 + num2 == seq_lens[i]:
|
|
continue
|
|
|
|
num3 = seq_lens[i] - seq_lens[i] // page_size * page_size
|
|
if num3:
|
|
out_indices[end_pos[i] - num3 : end_pos[i]] = (
|
|
free_pages[end_new_pages[i] - 1] * page_size + pos_in_page[:num3]
|
|
).view(-1)
|
|
|
|
|
|
@triton.jit
|
|
def alloc_extend_kernel(
|
|
pre_lens_ptr,
|
|
seq_lens_ptr,
|
|
last_loc_ptr,
|
|
free_page_ptr,
|
|
out_indices,
|
|
bs_upper: tl.constexpr,
|
|
page_size: tl.constexpr,
|
|
):
|
|
pid = tl.program_id(0)
|
|
|
|
load_offset = tl.arange(0, bs_upper)
|
|
seq_lens = tl.load(seq_lens_ptr + load_offset, mask=load_offset <= pid)
|
|
pre_lens = tl.load(pre_lens_ptr + load_offset, mask=load_offset <= pid)
|
|
extend_lens = seq_lens - pre_lens
|
|
|
|
seq_len = tl.load(seq_lens_ptr + pid)
|
|
pre_len = tl.load(pre_lens_ptr + pid)
|
|
extend_len = seq_len - pre_len
|
|
|
|
sum_extend_lens = tl.sum(extend_lens)
|
|
output_start_loc = sum_extend_lens - extend_len
|
|
|
|
num_pages_after = (seq_lens + page_size - 1) // page_size
|
|
num_pages_before = (pre_lens + page_size - 1) // page_size
|
|
num_new_pages = num_pages_after - num_pages_before
|
|
|
|
num_page_start_loc_self = (seq_len + page_size - 1) // page_size - (
|
|
pre_len + page_size - 1
|
|
) // page_size
|
|
sum_num_new_pages = tl.sum(num_new_pages)
|
|
new_page_start_loc = sum_num_new_pages - num_page_start_loc_self
|
|
|
|
# Part 1: fill the old partial page
|
|
last_loc = tl.load(last_loc_ptr + pid)
|
|
num_part1 = (
|
|
min(seq_len, (pre_len + page_size - 1) // page_size * page_size) - pre_len
|
|
)
|
|
offset_one_page = tl.arange(0, page_size)
|
|
tl.store(
|
|
out_indices + output_start_loc + offset_one_page,
|
|
last_loc + 1 + offset_one_page,
|
|
mask=offset_one_page < num_part1,
|
|
)
|
|
if pre_len + num_part1 == seq_len:
|
|
return
|
|
|
|
# Part 2: fill the new full pages using a dynamic blocked loop.
|
|
# The loop bound is derived from num_part2 (runtime value), so Triton
|
|
# generates a real loop instead of unrolling — no constexpr dependency
|
|
# on extend size and only one kernel compilation.
|
|
num_part2 = (
|
|
seq_len // page_size * page_size
|
|
- (pre_len + page_size - 1) // page_size * page_size
|
|
)
|
|
BLOCK_EXTEND: tl.constexpr = 4096
|
|
num_blocks = (num_part2 + BLOCK_EXTEND - 1) // BLOCK_EXTEND
|
|
for block_id in range(num_blocks):
|
|
offset_in_block = tl.arange(0, BLOCK_EXTEND)
|
|
offset = block_id * BLOCK_EXTEND + offset_in_block
|
|
mask = offset < num_part2
|
|
page_start = tl.load(
|
|
free_page_ptr + new_page_start_loc + offset // page_size,
|
|
mask=mask,
|
|
)
|
|
tl.store(
|
|
out_indices + output_start_loc + num_part1 + offset,
|
|
page_start * page_size + offset % page_size,
|
|
mask=mask,
|
|
)
|
|
if pre_len + num_part1 + num_part2 == seq_len:
|
|
return
|
|
|
|
# Part 3: fill the new partial page
|
|
num_part3 = seq_len - seq_len // page_size * page_size
|
|
start_loc = tl.load(
|
|
free_page_ptr + new_page_start_loc + num_page_start_loc_self - 1
|
|
)
|
|
tl.store(
|
|
out_indices + output_start_loc + num_part1 + num_part2 + offset_one_page,
|
|
start_loc * page_size + offset_one_page,
|
|
mask=offset_one_page < num_part3,
|
|
)
|
|
|
|
|
|
@triton.jit
|
|
def alloc_decode_kernel(
|
|
seq_lens_ptr,
|
|
last_loc_ptr,
|
|
free_page_ptr,
|
|
out_indices,
|
|
bs_upper: tl.constexpr,
|
|
page_size: tl.constexpr,
|
|
):
|
|
pid = tl.program_id(0)
|
|
|
|
load_offset = tl.arange(0, bs_upper)
|
|
seq_lens = tl.load(seq_lens_ptr + load_offset, mask=load_offset <= pid)
|
|
pre_lens = tl.where(load_offset <= pid, seq_lens - 1, seq_lens)
|
|
|
|
seq_len = tl.load(seq_lens_ptr + pid)
|
|
pre_len = seq_len - 1
|
|
|
|
num_pages_after = (seq_lens + page_size - 1) // page_size
|
|
num_pages_before = (pre_lens + page_size - 1) // page_size
|
|
num_new_pages = num_pages_after - num_pages_before
|
|
|
|
num_page_start_loc_self = (seq_len + page_size - 1) // page_size - (
|
|
pre_len + page_size - 1
|
|
) // page_size
|
|
sum_num_new_pages = tl.sum(num_new_pages)
|
|
new_page_start_loc = sum_num_new_pages - num_page_start_loc_self
|
|
|
|
if num_page_start_loc_self == 0:
|
|
last_loc = tl.load(last_loc_ptr + pid)
|
|
tl.store(out_indices + pid, last_loc + 1)
|
|
else:
|
|
page = tl.load(free_page_ptr + new_page_start_loc)
|
|
tl.store(out_indices + pid, page * page_size)
|
|
|
|
|
|
class PagedTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
|
|
"""
|
|
An allocator managing the indices to kv cache data.
|
|
|
|
This class has the same interface as `TokenToKVPoolAllocator` but the output
|
|
of one request is always page-aligned.
|
|
|
|
TODO: fuse last_loc into the kernel.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
size: int,
|
|
page_size: int,
|
|
dtype: torch.dtype,
|
|
device: str,
|
|
kvcache: KVCache,
|
|
need_sort: bool,
|
|
):
|
|
super().__init__(size, page_size, dtype, device, kvcache, need_sort)
|
|
self.num_pages = size // page_size
|
|
self.debug_mode = get_bool_env_var("SGLANG_DEBUG_MEMORY_POOL")
|
|
self.clear()
|
|
|
|
def alloc(self, need_size: int):
|
|
# page-aligned allocation, returning contiguous indices of pages
|
|
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 > len(self.free_pages):
|
|
self.merge_and_sort_free()
|
|
if num_pages > len(self.free_pages):
|
|
return None
|
|
|
|
out_pages = self.free_pages[:num_pages]
|
|
self.free_pages = self.free_pages[num_pages:]
|
|
|
|
out_indices = (
|
|
out_pages[:, None] * self.page_size
|
|
+ torch.arange(self.page_size, device=self.device)
|
|
).reshape(-1)
|
|
|
|
return out_indices
|
|
|
|
def alloc_extend(
|
|
self,
|
|
prefix_lens: torch.Tensor,
|
|
prefix_lens_cpu: torch.Tensor,
|
|
seq_lens: torch.Tensor,
|
|
seq_lens_cpu: torch.Tensor,
|
|
last_loc: torch.Tensor,
|
|
extend_num_tokens: int,
|
|
):
|
|
if self.debug_mode:
|
|
assert torch.all(
|
|
(last_loc + 1) % self.page_size == prefix_lens % self.page_size
|
|
)
|
|
|
|
bs = len(prefix_lens)
|
|
num_new_pages = get_num_new_pages(
|
|
seq_lens=seq_lens_cpu,
|
|
page_size=self.page_size,
|
|
prefix_lens=prefix_lens_cpu,
|
|
)
|
|
if self.need_sort and num_new_pages > len(self.free_pages):
|
|
self.merge_and_sort_free()
|
|
if num_new_pages > len(self.free_pages):
|
|
return None
|
|
|
|
out_indices = torch.empty(
|
|
(extend_num_tokens,), dtype=torch.int64, device=self.device
|
|
)
|
|
|
|
alloc_extend_kernel[(bs,)](
|
|
prefix_lens,
|
|
seq_lens,
|
|
last_loc,
|
|
self.free_pages,
|
|
out_indices,
|
|
next_power_of_2(bs),
|
|
self.page_size,
|
|
)
|
|
|
|
if self.debug_mode:
|
|
assert len(torch.unique(out_indices)) == len(out_indices)
|
|
|
|
self.free_pages = self.free_pages[num_new_pages:]
|
|
return out_indices
|
|
|
|
def alloc_decode(
|
|
self,
|
|
seq_lens: torch.Tensor,
|
|
seq_lens_cpu: torch.Tensor,
|
|
last_loc: torch.Tensor,
|
|
):
|
|
if self.debug_mode:
|
|
assert torch.all(
|
|
(last_loc + 2) % self.page_size == seq_lens % self.page_size
|
|
)
|
|
|
|
bs = len(seq_lens)
|
|
num_new_pages = get_num_new_pages(
|
|
seq_lens=seq_lens_cpu,
|
|
page_size=self.page_size,
|
|
decode=True,
|
|
)
|
|
if self.need_sort and num_new_pages > len(self.free_pages):
|
|
self.merge_and_sort_free()
|
|
if num_new_pages > len(self.free_pages):
|
|
return None
|
|
|
|
out_indices = torch.empty((bs,), dtype=torch.int64, device=self.device)
|
|
alloc_decode_kernel[(bs,)](
|
|
seq_lens,
|
|
last_loc,
|
|
self.free_pages,
|
|
out_indices,
|
|
next_power_of_2(bs),
|
|
self.page_size,
|
|
)
|
|
|
|
if self.debug_mode:
|
|
assert len(torch.unique(out_indices)) == len(out_indices)
|
|
|
|
self.free_pages = self.free_pages[num_new_pages:]
|
|
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)
|
|
if self.need_sort:
|
|
self.release_pages = torch.cat((free_page_indices, self.release_pages))
|
|
else:
|
|
self.free_pages = torch.cat((free_page_indices, self.free_pages))
|
|
else:
|
|
self.free_group.append(free_index)
|
|
|
|
if self.debug_mode:
|
|
assert len(torch.unique(self.free_pages)) == len(self.free_pages)
|
|
|
|
def clear(self):
|
|
# The padded slot 0 is used for writing dummy outputs from padded tokens.
|
|
self.free_pages = torch.arange(
|
|
1, self.num_pages + 1, dtype=torch.int64, device=self.device
|
|
)
|
|
self.is_not_in_free_group = True
|
|
self.free_group = []
|
|
self.release_pages = torch.empty((0,), dtype=torch.int64, device=self.device)
|
|
|
|
def get_cpu_copy(self, indices):
|
|
return self._kvcache.get_cpu_copy(indices)
|
|
|
|
def load_cpu_copy(self, kv_cache_cpu, indices):
|
|
return self._kvcache.load_cpu_copy(kv_cache_cpu, indices)
|
|
|
|
|
|
class CPSharedPagedTokenToKVPoolAllocator(PagedTokenToKVPoolAllocator):
|
|
"""Paged allocator that returns CP-group logical KV locations.
|
|
|
|
The allocator tracks logical pages visible to scheduler/radix/req_to_token.
|
|
The physical KV pool is smaller; logical-to-physical translation happens at
|
|
actual KV buffer access time through CpSharedKVLayout.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
logical_size: int,
|
|
physical_size: int,
|
|
page_size: int,
|
|
dtype: torch.dtype,
|
|
device: str,
|
|
kvcache: KVCache,
|
|
need_sort: bool,
|
|
cp_size: int,
|
|
cp_rank: int,
|
|
):
|
|
if logical_size % page_size != 0:
|
|
raise ValueError("logical_size must be page aligned")
|
|
if physical_size % page_size != 0:
|
|
raise ValueError("physical_size must be page aligned")
|
|
if logical_size != physical_size * cp_size:
|
|
raise ValueError(
|
|
"logical_size must equal physical_size * cp_size, got "
|
|
f"{logical_size=} {physical_size=} {cp_size=}"
|
|
)
|
|
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
|
|
|
|
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
|
|
|
|
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)
|
|
]
|
|
deficits = [
|
|
max(0, required_count - available_count)
|
|
for required_count, available_count in zip(required, available)
|
|
]
|
|
return required, available, deficits
|
|
|
|
def _select_compute_owner_pages(
|
|
self,
|
|
page_compute_owners: List[int],
|
|
) -> Optional[tuple[torch.Tensor, torch.Tensor, torch.Tensor]]:
|
|
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),
|
|
)
|
|
|
|
required_by_owner = [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_by_owner[owner] += 1
|
|
|
|
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)
|
|
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 = owner_mask & (
|
|
torch.cumsum(owner_mask.to(torch.int64), dim=0) <= required_count
|
|
)
|
|
selected_owner_pages = self.free_pages[selected_owner_free_mask]
|
|
|
|
remaining_count = required_count - selected_owner_pages.numel()
|
|
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():
|
|
return None
|
|
selected_owner_pages = torch.cat(
|
|
(selected_owner_pages, selected_owner_release_pages)
|
|
)
|
|
selected_release_mask |= selected_owner_release_mask
|
|
|
|
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
|
|
|
|
return (
|
|
torch.stack(selected_pages).to(torch.int64),
|
|
selected_free_mask,
|
|
selected_release_mask,
|
|
)
|
|
|
|
def alloc_extend_compute_owner(
|
|
self,
|
|
prefix_lens: torch.Tensor,
|
|
prefix_lens_cpu: torch.Tensor,
|
|
seq_lens: torch.Tensor,
|
|
seq_lens_cpu: torch.Tensor,
|
|
last_loc: torch.Tensor,
|
|
extend_num_tokens: int,
|
|
page_compute_owners: List[int],
|
|
):
|
|
"""Allocate extend KV locs so logical page owner matches CP compute rank.
|
|
|
|
The returned logical `out_cache_loc` is still full-order and identical on
|
|
every CP rank. Only the chosen logical page ids change: each newly
|
|
allocated request page comes from the modulo-owner lane that will compute
|
|
and directly persist that page.
|
|
"""
|
|
|
|
if len(prefix_lens_cpu) != 1 or len(seq_lens_cpu) != 1:
|
|
raise ValueError("compute-owner allocation supports batch size 1 only")
|
|
|
|
num_new_pages = get_num_new_pages(
|
|
seq_lens=seq_lens_cpu,
|
|
page_size=self.page_size,
|
|
prefix_lens=prefix_lens_cpu,
|
|
)
|
|
if num_new_pages != len(page_compute_owners):
|
|
raise ValueError(
|
|
"compute-owner page count mismatch: "
|
|
f"{num_new_pages=} page_compute_owners={len(page_compute_owners)}"
|
|
)
|
|
|
|
selected = self._select_compute_owner_pages(page_compute_owners)
|
|
if selected is None:
|
|
return None
|
|
selected_pages, selected_free_mask, selected_release_mask = selected
|
|
|
|
out_indices = torch.empty(
|
|
(extend_num_tokens,), dtype=torch.int64, device=self.device
|
|
)
|
|
alloc_extend_naive(
|
|
prefix_lens_cpu,
|
|
seq_lens_cpu,
|
|
last_loc,
|
|
selected_pages,
|
|
out_indices,
|
|
self.page_size,
|
|
self.device,
|
|
)
|
|
|
|
self.free_pages = self.free_pages[~selected_free_mask]
|
|
self.release_pages = self.release_pages[~selected_release_mask]
|
|
|
|
if self.debug_mode:
|
|
assert len(torch.unique(out_indices)) == len(out_indices)
|
|
selected_owners = torch.remainder(selected_pages - 1, self.cp_size)
|
|
expected_owners = torch.tensor(
|
|
page_compute_owners,
|
|
dtype=selected_owners.dtype,
|
|
device=selected_owners.device,
|
|
)
|
|
assert torch.equal(selected_owners, expected_owners)
|
|
|
|
return out_indices
|
|
|
|
def alloc_pages_with_owners(
|
|
self,
|
|
page_compute_owners: List[int],
|
|
) -> Optional[torch.Tensor]:
|
|
"""Allocate ``len(page_compute_owners)`` whole logical pages where
|
|
page ``i``'s owner equals ``page_compute_owners[i]``.
|
|
|
|
Returns a flat int64 tensor of logical token locs of length
|
|
``len(owners) * page_size``, page-contiguous in input order, or
|
|
``None`` when an owner lane is exhausted (caller handles
|
|
eviction-retry / cache-miss).
|
|
|
|
Used by CP HiCache ``load_cp`` to reproduce the write-time owner
|
|
pattern recorded in ``CpHiCacheNodeMetadata.page_owners``. Without
|
|
this, plain ``alloc()`` would return pages with an arbitrary owner
|
|
pattern, the saved per-position owner mask would index the wrong
|
|
slots, and each rank would load its host bytes into physical slots
|
|
whose corresponding logical page is owned by another rank → attention
|
|
reads garbage.
|
|
|
|
Distinguished from ``alloc_extend_compute_owner``: that variant
|
|
expects extend semantics (prefix_lens / seq_lens / last_loc) and
|
|
invokes the ``alloc_extend_naive`` Triton kernel to splice the
|
|
partial last page of a prefix with new pages. HiCache reload has
|
|
no prefix and no partial page — host pages are always page-aligned
|
|
— so a pure page-fresh allocation is enough.
|
|
"""
|
|
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_mask, selected_release_mask = 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]
|
|
if self.debug_mode:
|
|
assert torch.unique(out_indices).numel() == out_indices.numel()
|
|
check_owners = torch.remainder(selected_pages - 1, self.cp_size)
|
|
expected = torch.tensor(
|
|
page_compute_owners,
|
|
dtype=check_owners.dtype,
|
|
device=check_owners.device,
|
|
)
|
|
assert torch.equal(check_owners, expected)
|
|
return out_indices
|