342 lines
11 KiB
Python
342 lines
11 KiB
Python
# mapping on device memory, host memory and memory allocator
|
|
|
|
import weakref
|
|
from typing import Optional
|
|
|
|
import torch
|
|
from sgl_kernel.kvcacheio import transfer_kv_all_layer_mla
|
|
|
|
from sglang.srt.layers.radix_attention import RadixAttention
|
|
from sglang.srt.mem_cache.allocator import (
|
|
BaseTokenToKVPoolAllocator,
|
|
PagedTokenToKVPoolAllocator,
|
|
)
|
|
from sglang.srt.mem_cache.memory_pool import NSATokenToKVPool
|
|
|
|
|
|
class HiSparseNSATokenToKVPool(NSATokenToKVPool):
|
|
def __init__(
|
|
self,
|
|
size: int,
|
|
page_size: int,
|
|
kv_lora_rank: int,
|
|
dtype: torch.dtype,
|
|
qk_rope_head_dim: int,
|
|
layer_num: int,
|
|
device: str,
|
|
index_head_dim: int,
|
|
enable_memory_saver: bool,
|
|
kv_cache_dim: int,
|
|
start_layer: Optional[int] = None,
|
|
end_layer: Optional[int] = None,
|
|
host_to_device_ratio: int = 2,
|
|
):
|
|
super().__init__(
|
|
size=size,
|
|
page_size=page_size,
|
|
kv_lora_rank=kv_lora_rank,
|
|
dtype=dtype,
|
|
qk_rope_head_dim=qk_rope_head_dim,
|
|
layer_num=layer_num,
|
|
device=device,
|
|
index_head_dim=index_head_dim,
|
|
enable_memory_saver=enable_memory_saver,
|
|
kv_cache_dim=kv_cache_dim,
|
|
start_layer=start_layer,
|
|
end_layer=end_layer,
|
|
index_buf_size=size * host_to_device_ratio,
|
|
)
|
|
self.bytes_per_token = self.kv_cache_dim * self.dtype.itemsize
|
|
|
|
def register_mapping(self, full_to_hisparse_device_index_mapping: torch.Tensor):
|
|
self.full_to_hisparse_device_index_mapping = (
|
|
full_to_hisparse_device_index_mapping
|
|
)
|
|
|
|
def translate_loc_to_hisparse_device(self, compressed_indices: torch.Tensor):
|
|
return self.full_to_hisparse_device_index_mapping[compressed_indices].to(
|
|
torch.int32
|
|
)
|
|
|
|
def _translate_loc_to_hisparse_device(self, compressed_indices: torch.Tensor):
|
|
return self.full_to_hisparse_device_index_mapping[compressed_indices]
|
|
|
|
def set_kv_buffer(
|
|
self,
|
|
layer: RadixAttention,
|
|
loc: torch.Tensor,
|
|
cache_k: torch.Tensor,
|
|
cache_v: torch.Tensor,
|
|
):
|
|
loc = self.translate_loc_to_hisparse_device(loc)
|
|
super().set_kv_buffer(layer, loc, cache_k, cache_v)
|
|
|
|
def set_mla_kv_buffer(
|
|
self,
|
|
layer: RadixAttention,
|
|
loc: torch.Tensor,
|
|
cache_k_nope: torch.Tensor,
|
|
cache_k_rope: torch.Tensor,
|
|
):
|
|
loc = self.translate_loc_to_hisparse_device(loc)
|
|
super().set_mla_kv_buffer(layer, loc, cache_k_nope, cache_k_rope)
|
|
|
|
def get_mla_kv_buffer(
|
|
self,
|
|
layer: RadixAttention,
|
|
loc: torch.Tensor,
|
|
dst_dtype: Optional[torch.dtype] = None,
|
|
):
|
|
loc = self.translate_loc_to_hisparse_device(loc)
|
|
return super().get_mla_kv_buffer(layer, loc, dst_dtype)
|
|
|
|
def transfer_values_on_device(self, dst_indices, src_indices):
|
|
transfer_kv_all_layer_mla(
|
|
src_layers=self.data_ptrs,
|
|
dst_layers=self.data_ptrs,
|
|
src_indices=src_indices,
|
|
dst_indices=dst_indices,
|
|
item_size=self.bytes_per_token,
|
|
num_layers=self.layer_num,
|
|
)
|
|
|
|
def get_cpu_copy(self, indices):
|
|
raise NotImplementedError("HiSparseDevicePool does not support get_cpu_copy")
|
|
|
|
def load_cpu_copy(self, kv_cache_cpu, indices):
|
|
raise NotImplementedError("HiSparseDevicePool does not support load_cpu_copy")
|
|
|
|
|
|
class HiSparseTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
|
|
def __init__(
|
|
self,
|
|
size: int,
|
|
page_size: int,
|
|
dtype: torch.dtype,
|
|
device: torch.device,
|
|
kvcache: NSATokenToKVPool,
|
|
need_sort: bool,
|
|
host_to_device_ratio: int = 2,
|
|
):
|
|
self._kvcache = kvcache
|
|
self._size_full = size * host_to_device_ratio
|
|
self._size_hisparse = size
|
|
self.dtype = dtype
|
|
self.device = device
|
|
self.page_size = page_size
|
|
self.need_sort = need_sort
|
|
|
|
self.logical_attn_allocator = PagedTokenToKVPoolAllocator(
|
|
self._size_full,
|
|
self.page_size,
|
|
self.dtype,
|
|
self.device,
|
|
kvcache,
|
|
need_sort,
|
|
)
|
|
|
|
self.hisparse_attn_allocator = PagedTokenToKVPoolAllocator(
|
|
self._size_hisparse,
|
|
self.page_size,
|
|
self.dtype,
|
|
self.device,
|
|
kvcache,
|
|
need_sort,
|
|
)
|
|
|
|
self.full_to_hisparse_device_index_mapping = torch.cat(
|
|
[
|
|
torch.zeros(
|
|
self._size_full + self.page_size,
|
|
dtype=torch.int64,
|
|
device=self.device,
|
|
),
|
|
torch.tensor([-1], dtype=torch.int64, device=self.device),
|
|
]
|
|
)
|
|
|
|
self.free_pages = None
|
|
self.release_pages = None
|
|
self.is_not_in_free_group = True
|
|
self.free_group = []
|
|
self.clear()
|
|
|
|
self._kvcache.register_mapping(
|
|
weakref.proxy(self.full_to_hisparse_device_index_mapping)
|
|
)
|
|
|
|
@property
|
|
def size_full(self) -> int:
|
|
return self._size_full
|
|
|
|
def available_size(self) -> int:
|
|
return min(
|
|
self.logical_attn_allocator.available_size(),
|
|
self.hisparse_attn_allocator.available_size(),
|
|
)
|
|
|
|
def alloc(self, need_size: int):
|
|
raise NotImplementedError(
|
|
"Page size = 1 is not supported in HiSparse allocator"
|
|
)
|
|
|
|
def alloc_device_buffer(self, allocated_indices, need_size: int):
|
|
assert need_size % self.page_size == 0
|
|
# clear original reference and isolate the buffer from outside addressing, allocate new buffer if needed
|
|
hisparse_indices = self.full_to_hisparse_device_index_mapping[allocated_indices]
|
|
self.full_to_hisparse_device_index_mapping[allocated_indices] = 0
|
|
if len(hisparse_indices) >= need_size:
|
|
buffer_indices = hisparse_indices[:need_size]
|
|
self.free_hisparse_indices(hisparse_indices[need_size:])
|
|
else:
|
|
# page alignment, claiming the residual space for an incomplete page
|
|
page_residual_length = len(hisparse_indices) % self.page_size
|
|
if page_residual_length != 0:
|
|
hisparse_indices = torch.cat(
|
|
[
|
|
hisparse_indices,
|
|
torch.arange(
|
|
hisparse_indices[-1] + 1,
|
|
hisparse_indices[-1]
|
|
+ self.page_size
|
|
- page_residual_length
|
|
+ 1,
|
|
device=self.device,
|
|
),
|
|
]
|
|
)
|
|
extra_indices = self.hisparse_attn_allocator.alloc(
|
|
need_size - len(hisparse_indices)
|
|
)
|
|
assert (
|
|
extra_indices is not None
|
|
), "Hisparse allocation failed in alloc_device_buffer"
|
|
buffer_indices = torch.cat([hisparse_indices, extra_indices])
|
|
return buffer_indices
|
|
|
|
def free_hisparse_indices(self, buffer_indices: torch.Tensor):
|
|
# disable free group mechanism for device buffer free
|
|
self.hisparse_attn_allocator.is_not_in_free_group = True
|
|
self.hisparse_attn_allocator.free(buffer_indices[buffer_indices > 0])
|
|
|
|
def get_last_loc_hisparse_device(self, last_locs: torch.Tensor):
|
|
hisparse_last_locs = self._kvcache._translate_loc_to_hisparse_device(last_locs)
|
|
return hisparse_last_locs
|
|
|
|
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, # last_loc for full layers
|
|
extend_num_tokens: int,
|
|
):
|
|
assert self.page_size > 1
|
|
num_tokens = extend_num_tokens + len(seq_lens) * self.page_size
|
|
|
|
if num_tokens > self.available_size():
|
|
return None
|
|
|
|
logical_indices = self.logical_attn_allocator.alloc_extend(
|
|
prefix_lens,
|
|
prefix_lens_cpu,
|
|
seq_lens,
|
|
seq_lens_cpu,
|
|
last_loc,
|
|
extend_num_tokens,
|
|
)
|
|
assert logical_indices is not None, "Logical allocation failed in alloc_extend"
|
|
|
|
hisparse_last_loc = self.get_last_loc_hisparse_device(last_loc)
|
|
hisparse_indices = self.hisparse_attn_allocator.alloc_extend(
|
|
prefix_lens,
|
|
prefix_lens_cpu,
|
|
seq_lens,
|
|
seq_lens_cpu,
|
|
hisparse_last_loc,
|
|
len(logical_indices),
|
|
)
|
|
assert (
|
|
hisparse_indices is not None
|
|
), "Hisparse allocation failed in alloc_extend"
|
|
|
|
self.full_to_hisparse_device_index_mapping[logical_indices] = hisparse_indices
|
|
|
|
return logical_indices
|
|
|
|
def alloc_decode(
|
|
self,
|
|
seq_lens: torch.Tensor,
|
|
seq_lens_cpu: torch.Tensor,
|
|
last_loc: torch.Tensor, # last_loc for full layers
|
|
):
|
|
logical_indices = self.logical_attn_allocator.alloc_decode(
|
|
seq_lens, seq_lens_cpu, last_loc
|
|
)
|
|
|
|
return logical_indices
|
|
|
|
def alloc_decode_debug(
|
|
self,
|
|
seq_lens: torch.Tensor,
|
|
seq_lens_cpu: torch.Tensor,
|
|
last_loc: torch.Tensor, # last_loc for full layers
|
|
):
|
|
logical_indices = self.logical_attn_allocator.alloc_decode(
|
|
seq_lens, seq_lens_cpu, last_loc
|
|
)
|
|
|
|
hisparse_last_loc = self.get_last_loc_hisparse_device(last_loc)
|
|
hisparse_indices = self.hisparse_attn_allocator.alloc_decode(
|
|
seq_lens,
|
|
seq_lens_cpu,
|
|
hisparse_last_loc,
|
|
)
|
|
|
|
if logical_indices is None or hisparse_indices is None:
|
|
return None
|
|
|
|
self.full_to_hisparse_device_index_mapping[logical_indices] = hisparse_indices
|
|
|
|
return logical_indices
|
|
|
|
def free_hisparse(self, free_indices: torch.Tensor):
|
|
hisparse_indices = self._kvcache._translate_loc_to_hisparse_device(free_indices)
|
|
hisparse_indices = hisparse_indices[hisparse_indices > 0]
|
|
self.free_hisparse_indices(hisparse_indices)
|
|
self.full_to_hisparse_device_index_mapping[free_indices] = 0
|
|
|
|
def clear(self):
|
|
self.logical_attn_allocator.clear()
|
|
self.hisparse_attn_allocator.clear()
|
|
|
|
# Note: the last item is -1, we don't clear it, see the comment in __init__
|
|
self.full_to_hisparse_device_index_mapping[:-1].fill_(0)
|
|
self.is_not_in_free_group = True
|
|
self.free_group = []
|
|
|
|
def free_group_begin(self):
|
|
return
|
|
|
|
def free_group_end(self):
|
|
return
|
|
|
|
def free(self, free_index: torch.Tensor):
|
|
if free_index.numel() == 0:
|
|
return
|
|
|
|
if self.is_not_in_free_group:
|
|
self.logical_attn_allocator.free(free_index)
|
|
self.free_hisparse(free_index)
|
|
else:
|
|
self.free_group.append(free_index)
|
|
assert (
|
|
self.logical_attn_allocator.available_size()
|
|
<= self.logical_attn_allocator.size
|
|
)
|
|
assert (
|
|
self.hisparse_attn_allocator.available_size()
|
|
<= self.hisparse_attn_allocator.size
|
|
)
|