Support swa allocator page size > 1 (#16296)

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(
+        self.dtype = dtype
-            size,
+        self.device = device
-            dtype,
+        self.page_size = page_size
-            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._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):
+    def size(self):
-        return self._size_full
+        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
+        assert self._kvcache.full_to_swa_index_mapping is not None
-        return self.full_to_swa_index_mapping[kv_indices].to(torch.int32)
+        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 = []
 

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()
         )
-
+        return (
-        kv_data_ptrs = full_kv_data_ptrs
+            full_kv_data_ptrs,
-        kv_data_lens = full_kv_data_lens
+            full_kv_data_lens,
-        kv_item_lens = full_kv_item_lens
+            full_kv_item_lens,
-
+        )
-        return kv_data_ptrs, kv_data_lens, 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:
+            if self.swa_loc is not None:
-                loc = self.translate_loc_from_full_to_swa(loc)
+                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

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:
-                    if self.is_hybrid_swa:
+                    self.token_to_kv_pool_allocator = SWATokenToKVPoolAllocator(
-                        self.token_to_kv_pool_allocator = SWATokenToKVPoolAllocator(
+                        self.full_max_total_num_tokens,
-                            self.full_max_total_num_tokens,
+                        self.swa_max_total_num_tokens,
-                            self.swa_max_total_num_tokens,
+                        page_size=self.page_size,
-                            dtype=self.kv_cache_dtype,
+                        dtype=self.kv_cache_dtype,
-                            device=self.device,
+                        device=self.device,
-                            kvcache=self.token_to_kv_pool,
+                        kvcache=self.token_to_kv_pool,
-                            need_sort=need_sort,
+                        need_sort=need_sort,
-                        )
+                    )
-                    else:
+                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:
+                    else:
-                    assert not self.is_hybrid_swa
+                        self.token_to_kv_pool_allocator = PagedTokenToKVPoolAllocator(
-                    self.token_to_kv_pool_allocator = PagedTokenToKVPoolAllocator(
+                            self.max_total_num_tokens,
-                        self.max_total_num_tokens,
+                            page_size=self.page_size,
-                        page_size=self.page_size,
+                            dtype=self.kv_cache_dtype,
-                        dtype=self.kv_cache_dtype,
+                            device=self.device,
-                        device=self.device,
+                            kvcache=self.token_to_kv_pool,
-                        kvcache=self.token_to_kv_pool,
+                            need_sort=need_sort,
-                        need_sort=need_sort,
+                        )
-                    )
+
         else:
             assert self.is_draft_worker
             if self.is_hybrid_swa:

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,
             ),