460 lines
15 KiB
Python
460 lines
15 KiB
Python
import logging
|
|
import weakref
|
|
from typing import Dict, List, Optional, Tuple
|
|
|
|
import torch
|
|
|
|
from sglang.srt.layers.radix_attention import RadixAttention
|
|
from sglang.srt.mem_cache.allocator import (
|
|
BaseTokenToKVPoolAllocator,
|
|
PagedTokenToKVPoolAllocator,
|
|
TokenToKVPoolAllocator,
|
|
)
|
|
from sglang.srt.mem_cache.memory_pool import KVCache, MHATokenToKVPool
|
|
from sglang.srt.mem_cache.utils import maybe_init_custom_mem_pool
|
|
|
|
logger = logging.getLogger(__name__)
|
|
GB = 1024 * 1024 * 1024
|
|
|
|
|
|
class SWAKVPool(KVCache):
|
|
"""KV cache with separate pools for full and SWA attention layers."""
|
|
|
|
def __init__(
|
|
self,
|
|
size: int,
|
|
size_swa: int,
|
|
page_size: int,
|
|
dtype: torch.dtype,
|
|
head_num: int,
|
|
head_dim: int,
|
|
swa_attention_layer_ids: List[int],
|
|
full_attention_layer_ids: List[int],
|
|
enable_kvcache_transpose: bool,
|
|
device: str,
|
|
token_to_kv_pool_class: KVCache = MHATokenToKVPool,
|
|
**kwargs,
|
|
):
|
|
self.size = size
|
|
self.size_swa = size_swa
|
|
self.dtype = dtype
|
|
self.head_num = head_num
|
|
self.head_dim = head_dim
|
|
self.device = device
|
|
self.swa_layer_nums = len(swa_attention_layer_ids)
|
|
self.full_layer_nums = len(full_attention_layer_ids)
|
|
self.start_layer = 0
|
|
self.page_size = page_size
|
|
self.swa_loc = None
|
|
|
|
kwargs["page_size"] = page_size
|
|
kwargs["enable_memory_saver"] = False
|
|
kwargs["head_num"] = head_num
|
|
kwargs["head_dim"] = head_dim
|
|
kwargs["device"] = device
|
|
# TODO MHATransposedTokenToKVPool if enable_kvcache_transpose is True
|
|
assert not enable_kvcache_transpose
|
|
|
|
# for disagg with nvlink
|
|
self.enable_custom_mem_pool, self.custom_mem_pool, _ = (
|
|
maybe_init_custom_mem_pool(device=self.device)
|
|
)
|
|
|
|
self.swa_kv_pool = token_to_kv_pool_class(
|
|
size=size_swa,
|
|
dtype=dtype,
|
|
layer_num=self.swa_layer_nums,
|
|
**kwargs,
|
|
)
|
|
kwargs.pop("swa_head_num", None)
|
|
kwargs.pop("swa_head_dim", None)
|
|
kwargs.pop("swa_v_head_dim", None)
|
|
self.full_kv_pool = token_to_kv_pool_class(
|
|
size=size,
|
|
dtype=dtype,
|
|
layer_num=self.full_layer_nums,
|
|
**kwargs,
|
|
)
|
|
# {layer_id: (index, is_swa_layer)}
|
|
self.layers_mapping: Dict[int, Tuple[int, bool]] = {}
|
|
for full_attn_layer_id, global_layer_id in enumerate(full_attention_layer_ids):
|
|
self.layers_mapping[global_layer_id] = (full_attn_layer_id, False)
|
|
for swa_layer_id, global_layer_id in enumerate(swa_attention_layer_ids):
|
|
self.layers_mapping[global_layer_id] = (swa_layer_id, True)
|
|
self.full_to_swa_index_mapping: Optional[torch.Tensor] = None
|
|
|
|
k_size, v_size = self.get_kv_size_bytes()
|
|
self.mem_usage = (k_size + v_size) / GB
|
|
logger.info(
|
|
f"SWAKVPool mem usage: {self.mem_usage:.2f} GB, swa size: {self.size_swa}, full size: {self.size}"
|
|
)
|
|
|
|
def register_mapping(self, full_to_swa_index_mapping: torch.Tensor):
|
|
self.full_to_swa_index_mapping = full_to_swa_index_mapping
|
|
|
|
def get_kv_size_bytes(self):
|
|
k_size, v_size = self.full_kv_pool.get_kv_size_bytes()
|
|
k_size_swa, v_size_swa = self.swa_kv_pool.get_kv_size_bytes()
|
|
return k_size + k_size_swa, v_size + v_size_swa
|
|
|
|
def get_contiguous_buf_infos(self):
|
|
full_kv_data_ptrs, full_kv_data_lens, full_kv_item_lens = (
|
|
self.full_kv_pool.get_contiguous_buf_infos()
|
|
)
|
|
return (
|
|
full_kv_data_ptrs,
|
|
full_kv_data_lens,
|
|
full_kv_item_lens,
|
|
)
|
|
|
|
def get_state_buf_infos(self):
|
|
swa_kv_data_ptrs, swa_kv_data_lens, swa_kv_item_lens = (
|
|
self.swa_kv_pool.get_contiguous_buf_infos()
|
|
)
|
|
|
|
return swa_kv_data_ptrs, swa_kv_data_lens, swa_kv_item_lens
|
|
|
|
def get_key_buffer(self, layer_id: int):
|
|
layer_id_pool, is_swa_layer = self.layers_mapping[layer_id]
|
|
if is_swa_layer:
|
|
return self.swa_kv_pool.get_key_buffer(layer_id_pool)
|
|
else:
|
|
return self.full_kv_pool.get_key_buffer(layer_id_pool)
|
|
|
|
def get_value_buffer(self, layer_id: int):
|
|
layer_id_pool, is_swa_layer = self.layers_mapping[layer_id]
|
|
if is_swa_layer:
|
|
return self.swa_kv_pool.get_value_buffer(layer_id_pool)
|
|
else:
|
|
return self.full_kv_pool.get_value_buffer(layer_id_pool)
|
|
|
|
def get_kv_buffer(self, layer_id: int):
|
|
layer_id_pool, is_swa_layer = self.layers_mapping[layer_id]
|
|
if is_swa_layer:
|
|
return self.swa_kv_pool.get_kv_buffer(layer_id_pool)
|
|
else:
|
|
return self.full_kv_pool.get_kv_buffer(layer_id_pool)
|
|
|
|
def set_swa_loc(self, loc: torch.Tensor):
|
|
self.swa_loc = loc
|
|
|
|
def translate_loc_from_full_to_swa(self, kv_indices: torch.Tensor):
|
|
assert self.full_to_swa_index_mapping is not None
|
|
|
|
# Note: kv_indices could have -1 values (from alloc_extend), which will be mapped to -1
|
|
# since the last item of full_to_swa_index_mapping is -1.
|
|
return self.full_to_swa_index_mapping[kv_indices].to(torch.int32)
|
|
|
|
def set_kv_buffer(
|
|
self,
|
|
layer: RadixAttention,
|
|
loc: torch.Tensor,
|
|
cache_k: torch.Tensor,
|
|
cache_v: torch.Tensor,
|
|
k_scale: float = 1.0,
|
|
v_scale: float = 1.0,
|
|
):
|
|
|
|
layer_id = layer.layer_id
|
|
layer_id_pool, is_swa_layer = self.layers_mapping[layer_id]
|
|
if is_swa_layer:
|
|
if self.swa_loc is not None:
|
|
loc = self.swa_loc
|
|
else:
|
|
if self.full_to_swa_index_mapping is not None:
|
|
loc = self.translate_loc_from_full_to_swa(loc)
|
|
|
|
self.swa_kv_pool.set_kv_buffer(
|
|
None,
|
|
loc,
|
|
cache_k,
|
|
cache_v,
|
|
k_scale,
|
|
v_scale,
|
|
layer_id_override=layer_id_pool,
|
|
)
|
|
else:
|
|
self.full_kv_pool.set_kv_buffer(
|
|
None,
|
|
loc,
|
|
cache_k,
|
|
cache_v,
|
|
k_scale,
|
|
v_scale,
|
|
layer_id_override=layer_id_pool,
|
|
)
|
|
|
|
def move_kv_cache(self, tgt_loc: torch.Tensor, src_loc: torch.Tensor):
|
|
self.full_kv_pool.move_kv_cache(tgt_loc, src_loc)
|
|
tgt_loc_swa = self.translate_loc_from_full_to_swa(tgt_loc)
|
|
src_loc_swa = self.translate_loc_from_full_to_swa(src_loc)
|
|
self.swa_kv_pool.move_kv_cache(tgt_loc_swa, src_loc_swa)
|
|
|
|
def get_cpu_copy(self, indices):
|
|
# For SWA, we need to copy KV cache from both full and SWA pools
|
|
# The indices are for the full pool, and we use mapping to get SWA indices
|
|
full_kv_cpu = self.full_kv_pool.get_cpu_copy(indices)
|
|
|
|
# Get SWA indices through the mapping
|
|
# Note: SWA allocation always creates 1:1 mapping, so no need to filter
|
|
if self.full_to_swa_index_mapping is not None:
|
|
swa_indices = self.full_to_swa_index_mapping[indices]
|
|
swa_kv_cpu = self.swa_kv_pool.get_cpu_copy(swa_indices)
|
|
else:
|
|
swa_kv_cpu = None
|
|
|
|
return {"full": full_kv_cpu, "swa": swa_kv_cpu}
|
|
|
|
def load_cpu_copy(self, kv_cache_cpu, indices):
|
|
# Load KV cache back from CPU to both full and SWA pools
|
|
# Note: indices here are NEW indices (newly allocated), different from get_cpu_copy indices
|
|
full_kv_cpu = kv_cache_cpu["full"]
|
|
swa_kv_cpu = kv_cache_cpu["swa"]
|
|
|
|
# Load full KV cache to the new indices
|
|
self.full_kv_pool.load_cpu_copy(full_kv_cpu, indices)
|
|
|
|
# Load SWA KV cache if it exists
|
|
if swa_kv_cpu is not None and self.full_to_swa_index_mapping is not None:
|
|
swa_indices = self.full_to_swa_index_mapping[indices]
|
|
self.swa_kv_pool.load_cpu_copy(swa_kv_cpu, swa_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)
|