diff --git a/python/sglang/srt/mem_cache/allocator.py b/python/sglang/srt/mem_cache/allocator.py index e8edb50da..dfc34ffe9 100644 --- a/python/sglang/srt/mem_cache/allocator.py +++ b/python/sglang/srt/mem_cache/allocator.py @@ -20,6 +20,7 @@ Page-aligned memory pool. """ import abc +import weakref from typing import TYPE_CHECKING import torch @@ -179,39 +180,77 @@ class SWATokenToKVPoolAllocator(BaseTokenToKVPoolAllocator): self, size: int, size_swa: int, + page_size: int, dtype: torch.dtype, device: str, kvcache: SWAKVPool, need_sort: bool, ): - super().__init__(size, 1, dtype, device, kvcache, need_sort) assert isinstance(kvcache, SWAKVPool) self._size_full = size self._size_swa = size_swa - 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, - ) - self.full_to_swa_index_mapping = torch.empty( - size + size_swa + 1, - dtype=torch.int64, - device=device, - ) - self.clear() + self.dtype = dtype + self.device = device + self.page_size = page_size - self._kvcache.full_to_swa_index_mapping = self.full_to_swa_index_mapping + 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): @@ -221,13 +260,17 @@ class SWATokenToKVPoolAllocator(BaseTokenToKVPoolAllocator): return self.swa_attn_allocator.available_size() @property - def size_full(self): - return self._size_full + 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()}, " @@ -240,10 +283,11 @@ class SWATokenToKVPoolAllocator(BaseTokenToKVPoolAllocator): return self._kvcache def translate_loc_from_full_to_swa(self, kv_indices: torch.Tensor): - assert self.full_to_swa_index_mapping is not None - return self.full_to_swa_index_mapping[kv_indices].to(torch.int32) + 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(): @@ -251,12 +295,81 @@ class SWATokenToKVPoolAllocator(BaseTokenToKVPoolAllocator): 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) @@ -287,7 +400,8 @@ class SWATokenToKVPoolAllocator(BaseTokenToKVPoolAllocator): def clear(self): self.swa_attn_allocator.clear() self.full_attn_allocator.clear() - self.full_to_swa_index_mapping.fill_(0) + # 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 = [] diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index b75d3be86..820a21db5 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/python/sglang/srt/mem_cache/memory_pool.py @@ -1284,6 +1284,7 @@ class SWAKVPool(KVCache): self, size: int, size_swa: int, + page_size: int, dtype: torch.dtype, head_num: int, head_dim: int, @@ -1303,9 +1304,10 @@ class SWAKVPool(KVCache): 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 = 1 + self.page_size = page_size + self.swa_loc = None - kwargs["page_size"] = 1 + kwargs["page_size"] = page_size kwargs["enable_memory_saver"] = False kwargs["head_num"] = head_num kwargs["head_dim"] = head_dim @@ -1347,6 +1349,9 @@ class SWAKVPool(KVCache): 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() @@ -1356,12 +1361,11 @@ class SWAKVPool(KVCache): full_kv_data_ptrs, full_kv_data_lens, full_kv_item_lens = ( self.full_kv_pool.get_contiguous_buf_infos() ) - - kv_data_ptrs = full_kv_data_ptrs - kv_data_lens = full_kv_data_lens - kv_item_lens = full_kv_item_lens - - return kv_data_ptrs, kv_data_lens, kv_item_lens + 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 = ( @@ -1391,8 +1395,14 @@ class SWAKVPool(KVCache): 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( @@ -1408,8 +1418,12 @@ class SWAKVPool(KVCache): layer_id = layer.layer_id layer_id_pool, is_swa_layer = self.layers_mapping[layer_id] if is_swa_layer: - if self.full_to_swa_index_mapping is not None: - loc = self.translate_loc_from_full_to_swa(loc) + 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, @@ -1430,6 +1444,12 @@ class SWAKVPool(KVCache): 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 diff --git a/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py b/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py index f4b06a594..d9645ebc8 100644 --- a/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py +++ b/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py @@ -517,6 +517,7 @@ class ModelRunnerKVCacheMixin: self.token_to_kv_pool = SWAKVPool( size=self.full_max_total_num_tokens, size_swa=self.swa_max_total_num_tokens, + page_size=self.page_size, dtype=self.kv_cache_dtype, head_num=self.model_config.get_num_kv_heads( get_attention_tp_size() @@ -614,17 +615,18 @@ class ModelRunnerKVCacheMixin: need_sort=need_sort, ) else: - if self.page_size == 1: - if self.is_hybrid_swa: - self.token_to_kv_pool_allocator = SWATokenToKVPoolAllocator( - self.full_max_total_num_tokens, - self.swa_max_total_num_tokens, - dtype=self.kv_cache_dtype, - device=self.device, - kvcache=self.token_to_kv_pool, - need_sort=need_sort, - ) - else: + if self.is_hybrid_swa: + self.token_to_kv_pool_allocator = SWATokenToKVPoolAllocator( + self.full_max_total_num_tokens, + self.swa_max_total_num_tokens, + page_size=self.page_size, + dtype=self.kv_cache_dtype, + device=self.device, + kvcache=self.token_to_kv_pool, + need_sort=need_sort, + ) + else: + if self.page_size == 1: self.token_to_kv_pool_allocator = TokenToKVPoolAllocator( self.max_total_num_tokens, dtype=self.kv_cache_dtype, @@ -632,16 +634,16 @@ class ModelRunnerKVCacheMixin: kvcache=self.token_to_kv_pool, need_sort=need_sort, ) - else: - assert not self.is_hybrid_swa - self.token_to_kv_pool_allocator = PagedTokenToKVPoolAllocator( - self.max_total_num_tokens, - page_size=self.page_size, - dtype=self.kv_cache_dtype, - device=self.device, - kvcache=self.token_to_kv_pool, - need_sort=need_sort, - ) + else: + self.token_to_kv_pool_allocator = PagedTokenToKVPoolAllocator( + self.max_total_num_tokens, + page_size=self.page_size, + dtype=self.kv_cache_dtype, + device=self.device, + kvcache=self.token_to_kv_pool, + need_sort=need_sort, + ) + else: assert self.is_draft_worker if self.is_hybrid_swa: diff --git a/test/srt/test_swa_unittest.py b/test/srt/test_swa_unittest.py index 2d01f90bd..ef8a69d68 100644 --- a/test/srt/test_swa_unittest.py +++ b/test/srt/test_swa_unittest.py @@ -21,6 +21,7 @@ class TestSWA(unittest.TestCase): def test_swa_memory_pool(self): size = 16 size_swa = 16 + page_size = 1 head_num = 8 head_dim = 128 num_layers = 48 @@ -35,6 +36,7 @@ class TestSWA(unittest.TestCase): pool = SWAKVPool( size=size, size_swa=size_swa, + page_size=page_size, dtype=dtype, head_num=head_num, head_dim=head_dim, @@ -46,6 +48,7 @@ class TestSWA(unittest.TestCase): alloc = SWATokenToKVPoolAllocator( size=size, size_swa=size_swa, + page_size=page_size, dtype=dtype, device=device, kvcache=pool, @@ -69,6 +72,7 @@ class TestSWA(unittest.TestCase): max_context_len = 128 kv_size = 128 kv_size_swa = 64 + page_size = 1 sliding_window_size = 4 head_num = 8 head_dim = 128 @@ -92,6 +96,7 @@ class TestSWA(unittest.TestCase): kv_pool = SWAKVPool( size=kv_size, size_swa=kv_size_swa, + page_size=page_size, dtype=dtype, head_num=head_num, head_dim=head_dim, @@ -104,6 +109,7 @@ class TestSWA(unittest.TestCase): allocator = SWATokenToKVPoolAllocator( size=kv_size, size_swa=kv_size_swa, + page_size=page_size, dtype=dtype, device=device, kvcache=kv_pool, @@ -115,7 +121,7 @@ class TestSWA(unittest.TestCase): req_to_token_pool=req_to_token_pool, token_to_kv_pool_allocator=allocator, disable=False, - page_size=1, + page_size=page_size, ), sliding_window_size=sliding_window_size, ) @@ -202,6 +208,7 @@ class TestSWA(unittest.TestCase): max_context_len = 128 kv_size = 128 kv_size_swa = 64 + page_size = 1 sliding_window_size = 4 head_num = 8 head_dim = 128 @@ -225,6 +232,7 @@ class TestSWA(unittest.TestCase): kv_pool = SWAKVPool( size=kv_size, size_swa=kv_size_swa, + page_size=page_size, dtype=dtype, head_num=head_num, head_dim=head_dim, @@ -237,6 +245,7 @@ class TestSWA(unittest.TestCase): allocator = SWATokenToKVPoolAllocator( size=kv_size, size_swa=kv_size_swa, + page_size=page_size, dtype=dtype, device=device, kvcache=kv_pool, @@ -247,7 +256,7 @@ class TestSWA(unittest.TestCase): params=CacheInitParams( req_to_token_pool=req_to_token_pool, token_to_kv_pool_allocator=allocator, - page_size=1, + page_size=page_size, disable=False, is_eagle=True, ),