Move swa memory pool to a seperate file (#16347)
This commit is contained in:
Ke Bao
@@ -20,14 +20,12 @@ Page-aligned memory pool.
|
||||
"""
|
||||
|
||||
import abc
|
||||
import weakref
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import torch
|
||||
import triton
|
||||
import triton.language as tl
|
||||
|
||||
from sglang.srt.mem_cache.memory_pool import SWAKVPool
|
||||
from sglang.srt.utils import get_bool_env_var, get_num_new_pages, next_power_of_2
|
||||
|
||||
if TYPE_CHECKING:
|
||||
@@ -173,245 +171,6 @@ class TokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
|
||||
return self._kvcache.load_cpu_copy(kv_cache_cpu, indices)
|
||||
|
||||
|
||||
class SWATokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
|
||||
"""Allocator for SWA hybrid KV cache."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
size: int,
|
||||
size_swa: int,
|
||||
page_size: int,
|
||||
dtype: torch.dtype,
|
||||
device: str,
|
||||
kvcache: SWAKVPool,
|
||||
need_sort: bool,
|
||||
):
|
||||
assert isinstance(kvcache, SWAKVPool)
|
||||
self._size_full = size
|
||||
self._size_swa = size_swa
|
||||
self.dtype = dtype
|
||||
self.device = device
|
||||
self.page_size = page_size
|
||||
|
||||
if page_size == 1:
|
||||
self.full_attn_allocator = TokenToKVPoolAllocator(
|
||||
size,
|
||||
dtype,
|
||||
device,
|
||||
kvcache.full_kv_pool,
|
||||
need_sort,
|
||||
)
|
||||
self.swa_attn_allocator = TokenToKVPoolAllocator(
|
||||
size_swa,
|
||||
dtype,
|
||||
device,
|
||||
kvcache.swa_kv_pool,
|
||||
need_sort,
|
||||
)
|
||||
else:
|
||||
self.full_attn_allocator = PagedTokenToKVPoolAllocator(
|
||||
size,
|
||||
page_size,
|
||||
dtype,
|
||||
device,
|
||||
kvcache.full_kv_pool,
|
||||
need_sort,
|
||||
)
|
||||
self.swa_attn_allocator = PagedTokenToKVPoolAllocator(
|
||||
size_swa,
|
||||
page_size,
|
||||
dtype,
|
||||
device,
|
||||
kvcache.swa_kv_pool,
|
||||
need_sort,
|
||||
)
|
||||
# Note: append one more item of value -1 in the end so -1 maps to -1.
|
||||
# It is needed for the last_loc in alloc_extend, where the first full_last_loc
|
||||
# is -1, and we need to map it to swa_last_loc -1 as well.
|
||||
self.full_to_swa_index_mapping = torch.cat(
|
||||
[
|
||||
torch.zeros(
|
||||
size + self.page_size,
|
||||
dtype=torch.int64,
|
||||
device=device,
|
||||
),
|
||||
torch.tensor([-1], dtype=torch.int64, device=device),
|
||||
]
|
||||
)
|
||||
|
||||
self.need_sort = need_sort
|
||||
self.free_pages = None
|
||||
self.release_pages = None
|
||||
self.is_not_in_free_group = True
|
||||
self.free_group = []
|
||||
|
||||
self.clear()
|
||||
self._kvcache = kvcache
|
||||
self._kvcache.register_mapping(weakref.proxy(self.full_to_swa_index_mapping))
|
||||
|
||||
def available_size(self):
|
||||
# Note: use full_available_size() and swa_available_size() instead.
|
||||
raise NotImplementedError()
|
||||
|
||||
def full_available_size(self):
|
||||
return self.full_attn_allocator.available_size()
|
||||
|
||||
def swa_available_size(self):
|
||||
return self.swa_attn_allocator.available_size()
|
||||
|
||||
@property
|
||||
def size(self):
|
||||
return min(self._size_full, self._size_swa)
|
||||
|
||||
@property
|
||||
def size_swa(self):
|
||||
return self._size_swa
|
||||
|
||||
@property
|
||||
def size_full(self):
|
||||
return self._size_full
|
||||
|
||||
def debug_print(self) -> str:
|
||||
msg = ""
|
||||
msg += f"#swa-available-size: {self.swa_attn_allocator.available_size()}, "
|
||||
msg += (
|
||||
f"#full-attn-available-size: {self.full_attn_allocator.available_size()}, "
|
||||
)
|
||||
return msg
|
||||
|
||||
def get_kvcache(self):
|
||||
return self._kvcache
|
||||
|
||||
def translate_loc_from_full_to_swa(self, kv_indices: torch.Tensor):
|
||||
assert self._kvcache.full_to_swa_index_mapping is not None
|
||||
return self._kvcache.translate_loc_from_full_to_swa(kv_indices)
|
||||
|
||||
def alloc(self, need_size: int):
|
||||
assert self.page_size == 1
|
||||
if need_size > self.full_attn_allocator.available_size():
|
||||
return None
|
||||
if need_size > self.swa_attn_allocator.available_size():
|
||||
return None
|
||||
|
||||
alloc_full_indices = self.full_attn_allocator.alloc(need_size)
|
||||
alloc_swa_indices = self.swa_attn_allocator.alloc(need_size)
|
||||
assert alloc_full_indices is not None
|
||||
assert alloc_swa_indices is not None
|
||||
|
||||
self.full_to_swa_index_mapping[alloc_full_indices] = alloc_swa_indices
|
||||
return alloc_full_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, # 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.full_attn_allocator.available_size():
|
||||
return None
|
||||
if num_tokens > self.swa_attn_allocator.available_size():
|
||||
return None
|
||||
|
||||
swa_last_loc = self.translate_loc_from_full_to_swa(last_loc)
|
||||
|
||||
alloc_full_indices = self.full_attn_allocator.alloc_extend(
|
||||
prefix_lens,
|
||||
prefix_lens_cpu,
|
||||
seq_lens,
|
||||
seq_lens_cpu,
|
||||
last_loc,
|
||||
extend_num_tokens,
|
||||
)
|
||||
alloc_swa_indices = self.swa_attn_allocator.alloc_extend(
|
||||
prefix_lens,
|
||||
prefix_lens_cpu,
|
||||
seq_lens,
|
||||
seq_lens_cpu,
|
||||
swa_last_loc,
|
||||
extend_num_tokens,
|
||||
)
|
||||
assert alloc_full_indices is not None
|
||||
assert alloc_swa_indices is not None
|
||||
|
||||
self.full_to_swa_index_mapping[alloc_full_indices] = alloc_swa_indices
|
||||
|
||||
return alloc_full_indices
|
||||
|
||||
def alloc_decode(
|
||||
self,
|
||||
seq_lens: torch.Tensor,
|
||||
seq_lens_cpu: torch.Tensor,
|
||||
last_loc: torch.Tensor, # last_loc for full layers
|
||||
):
|
||||
assert self.page_size > 1
|
||||
swa_last_loc = self.translate_loc_from_full_to_swa(last_loc)
|
||||
|
||||
alloc_full_indices = self.full_attn_allocator.alloc_decode(
|
||||
seq_lens, seq_lens_cpu, last_loc
|
||||
)
|
||||
alloc_swa_indices = self.swa_attn_allocator.alloc_decode(
|
||||
seq_lens, seq_lens_cpu, swa_last_loc
|
||||
)
|
||||
|
||||
if alloc_full_indices is None or alloc_swa_indices is None:
|
||||
return None
|
||||
|
||||
self.full_to_swa_index_mapping[alloc_full_indices] = alloc_swa_indices
|
||||
|
||||
return alloc_full_indices
|
||||
|
||||
def free(self, free_index: torch.Tensor):
|
||||
if free_index.numel() == 0:
|
||||
return
|
||||
|
||||
# NOTE: the API is not idempotent.
|
||||
if self.is_not_in_free_group:
|
||||
self.full_attn_allocator.free(free_index)
|
||||
self.free_swa(free_index)
|
||||
else:
|
||||
self.free_group.append(free_index)
|
||||
assert (
|
||||
self.full_attn_allocator.available_size() <= self.full_attn_allocator.size
|
||||
)
|
||||
assert self.swa_attn_allocator.available_size() <= self.swa_attn_allocator.size
|
||||
|
||||
def free_swa(self, free_index: torch.Tensor):
|
||||
swa_indices = self.full_to_swa_index_mapping[free_index]
|
||||
swa_indices = swa_indices[swa_indices > 0]
|
||||
self.swa_attn_allocator.free(swa_indices)
|
||||
self.full_to_swa_index_mapping[free_index] = 0
|
||||
|
||||
def backup_state(self):
|
||||
return [
|
||||
self.full_attn_allocator.backup_state(),
|
||||
self.swa_attn_allocator.backup_state(),
|
||||
]
|
||||
|
||||
def restore_state(self, state):
|
||||
assert len(state) == 2
|
||||
self.full_attn_allocator.restore_state(state[0])
|
||||
self.swa_attn_allocator.restore_state(state[1])
|
||||
|
||||
def clear(self):
|
||||
self.swa_attn_allocator.clear()
|
||||
self.full_attn_allocator.clear()
|
||||
# Note: the last item is -1, we don't clear it, see the comment in __init__
|
||||
self.full_to_swa_index_mapping[:-1].fill_(0)
|
||||
self.is_not_in_free_group = True
|
||||
self.free_group = []
|
||||
|
||||
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)
|
||||
|
||||
|
||||
@triton.jit
|
||||
def alloc_extend_kernel(
|
||||
pre_lens_ptr,
|
||||
|
||||
Reference in New Issue
Block a user