[mem_cache][1/N] refactor: split allocator.py into allocator/ subpackage (#26675)

python/sglang/srt/mem_cache/allocator/__init__.py

@@ -0,0 +1,15 @@
+"""Token-to-KV-slot allocators. One file per allocation strategy."""
+
+from sglang.srt.mem_cache.allocator.base import BaseTokenToKVPoolAllocator
+from sglang.srt.mem_cache.allocator.paged import (
+    PagedTokenToKVPoolAllocator,
+    alloc_extend_naive,
+)
+from sglang.srt.mem_cache.allocator.token import TokenToKVPoolAllocator
+
+__all__ = [
+    "BaseTokenToKVPoolAllocator",
+    "PagedTokenToKVPoolAllocator",
+    "TokenToKVPoolAllocator",
+    "alloc_extend_naive",
+]

python/sglang/srt/mem_cache/allocator/base.py

@@ -0,0 +1,110 @@
+"""
+Copyright 2025 SGLang Team
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+
+from __future__ import annotations
+
+import abc
+from typing import TYPE_CHECKING
+
+import torch
+
+if TYPE_CHECKING:
+    from sglang.srt.mem_cache.memory_pool import KVCache
+
+
+class BaseTokenToKVPoolAllocator(abc.ABC):
+    @abc.abstractmethod
+    def __init__(
+        self,
+        size: int,
+        page_size: int,
+        dtype: torch.dtype,
+        device: str,
+        kvcache: KVCache,
+        need_sort: bool,
+    ):
+        self.size = size
+        self.page_size = page_size
+        self.dtype = dtype
+        self.device = device
+        self._kvcache = kvcache
+        self.need_sort = need_sort
+
+        self.free_pages = None
+        self.release_pages = None
+        self.is_not_in_free_group = True
+        self.free_group = []
+
+    @property
+    def size_full(self):
+        return self.size
+
+    def debug_print(self) -> str:
+        return ""
+
+    def available_size(self):
+        return (len(self.free_pages) + len(self.release_pages)) * self.page_size
+
+    def get_kvcache(self):
+        return self._kvcache
+
+    def restore_state(self, state):
+        self.free_pages, self.release_pages = state
+
+    def backup_state(self):
+        return (self.free_pages, self.release_pages)
+
+    def free_group_begin(self):
+        self.is_not_in_free_group = False
+        self.free_group = []
+
+    def free_group_end(self):
+        self.is_not_in_free_group = True
+        if self.free_group:
+            self.free(torch.cat(self.free_group))
+
+    def merge_and_sort_free(self):
+        if len(self.release_pages) > 0:
+            self.free_pages = torch.cat((self.free_pages, self.release_pages))
+            self.free_pages, _ = torch.sort(self.free_pages)
+            self.release_pages = torch.empty(
+                (0,), dtype=self.release_pages.dtype, device=self.device
+            )
+
+    def get_cpu_copy(self, indices, mamba_indices=None):
+        # FIXME: reuse the get_cpu_copy after paged allocator is implemented
+        raise NotImplementedError()
+
+    def load_cpu_copy(self, kv_cache_cpu, indices, mamba_indices=None):
+        # FIXME: reuse the load_cpu_copy after paged allocator is implemented
+        raise NotImplementedError()
+
+    def alloc_extend(self, *args, **kwargs):
+        raise NotImplementedError("alloc_extend is only for paged allocator")
+
+    def alloc_decode(self, *args, **kwargs):
+        raise NotImplementedError("alloc_decode is only for paged allocator")
+
+    @abc.abstractmethod
+    def clear(self):
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    def alloc(self, need_size: int):
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    def free(self, free_index: torch.Tensor):
+        raise NotImplementedError()

python/sglang/srt/mem_cache/allocator/paged.py

@@ -1,5 +1,3 @@
-from __future__ import annotations
-
 """
 Copyright 2025 SGLang Team
 Licensed under the Apache License, Version 2.0 (the "License");
@@ -15,168 +13,25 @@ See the License for the specific language governing permissions and
 limitations under the License.
 """
 
+from __future__ import annotations
+
 """
 Page-aligned memory pool.
 """
 
-import abc
 from typing import TYPE_CHECKING
 
 import torch
 import triton
 import triton.language as tl
 
+from sglang.srt.mem_cache.allocator.base import BaseTokenToKVPoolAllocator
 from sglang.srt.utils import get_bool_env_var, get_num_new_pages, next_power_of_2
 
 if TYPE_CHECKING:
     from sglang.srt.mem_cache.memory_pool import KVCache
 
 
-class BaseTokenToKVPoolAllocator(abc.ABC):
-    @abc.abstractmethod
-    def __init__(
-        self,
-        size: int,
-        page_size: int,
-        dtype: torch.dtype,
-        device: str,
-        kvcache: KVCache,
-        need_sort: bool,
-    ):
-        self.size = size
-        self.page_size = page_size
-        self.dtype = dtype
-        self.device = device
-        self._kvcache = kvcache
-        self.need_sort = need_sort
-
-        self.free_pages = None
-        self.release_pages = None
-        self.is_not_in_free_group = True
-        self.free_group = []
-
-    @property
-    def size_full(self):
-        return self.size
-
-    def debug_print(self) -> str:
-        return ""
-
-    def available_size(self):
-        return (len(self.free_pages) + len(self.release_pages)) * self.page_size
-
-    def get_kvcache(self):
-        return self._kvcache
-
-    def restore_state(self, state):
-        self.free_pages, self.release_pages = state
-
-    def backup_state(self):
-        return (self.free_pages, self.release_pages)
-
-    def free_group_begin(self):
-        self.is_not_in_free_group = False
-        self.free_group = []
-
-    def free_group_end(self):
-        self.is_not_in_free_group = True
-        if self.free_group:
-            self.free(torch.cat(self.free_group))
-
-    def merge_and_sort_free(self):
-        if len(self.release_pages) > 0:
-            self.free_pages = torch.cat((self.free_pages, self.release_pages))
-            self.free_pages, _ = torch.sort(self.free_pages)
-            self.release_pages = torch.empty(
-                (0,), dtype=self.release_pages.dtype, device=self.device
-            )
-
-    def get_cpu_copy(self, indices, mamba_indices=None):
-        # FIXME: reuse the get_cpu_copy after paged allocator is implemented
-        raise NotImplementedError()
-
-    def load_cpu_copy(self, kv_cache_cpu, indices, mamba_indices=None):
-        # FIXME: reuse the load_cpu_copy after paged allocator is implemented
-        raise NotImplementedError()
-
-    def alloc_extend(self, *args, **kwargs):
-        raise NotImplementedError("alloc_extend is only for paged allocator")
-
-    def alloc_decode(self, *args, **kwargs):
-        raise NotImplementedError("alloc_decode is only for paged allocator")
-
-    @abc.abstractmethod
-    def clear(self):
-        raise NotImplementedError()
-
-    @abc.abstractmethod
-    def alloc(self, need_size: int):
-        raise NotImplementedError()
-
-    @abc.abstractmethod
-    def free(self, free_index: torch.Tensor):
-        raise NotImplementedError()
-
-
-class TokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
-    """An allocator managing the indices to kv cache data."""
-
-    def __init__(
-        self,
-        size: int,
-        dtype: torch.dtype,
-        device: str,
-        kvcache: KVCache,
-        need_sort: bool,
-    ):
-        super().__init__(size, 1, dtype, device, kvcache, need_sort)
-        self.clear()
-
-    def clear(self):
-        # The padded slot 0 is used for writing dummy outputs from padded tokens.
-        self.free_pages = torch.arange(
-            1, self.size + 1, dtype=torch.int64, device=self.device
-        )
-        self.is_not_in_free_group = True
-        self.free_group = []
-        self.release_pages = torch.empty((0,), dtype=torch.int64, device=self.device)
-
-    def available_size(self):
-        # To avoid minor "len(free_pages) * 1" overhead
-        return len(self.free_pages) + len(self.release_pages)
-
-    def alloc(self, need_size: int):
-        if self.need_sort and need_size > len(self.free_pages):
-            self.merge_and_sort_free()
-
-        if need_size > len(self.free_pages):
-            return None
-
-        select_index = self.free_pages[:need_size]
-        self.free_pages = self.free_pages[need_size:]
-        return select_index
-
-    def free(self, free_index: torch.Tensor):
-        if free_index.numel() == 0:
-            return
-
-        if self.is_not_in_free_group:
-            if self.need_sort:
-                self.release_pages = torch.cat((self.release_pages, free_index))
-            else:
-                self.free_pages = torch.cat((self.free_pages, free_index))
-        else:
-            self.free_group.append(free_index)
-
-    def get_cpu_copy(self, indices, mamba_indices=None):
-        return self._kvcache.get_cpu_copy(indices, mamba_indices=mamba_indices)
-
-    def load_cpu_copy(self, kv_cache_cpu, indices, mamba_indices=None):
-        return self._kvcache.load_cpu_copy(
-            kv_cache_cpu, indices, mamba_indices=mamba_indices
-        )
-
-
 def alloc_extend_naive(
     prefix_lens,
     seq_lens,

python/sglang/srt/mem_cache/allocator/token.py

@@ -0,0 +1,84 @@
+"""
+Copyright 2025 SGLang Team
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import torch
+
+from sglang.srt.mem_cache.allocator.base import BaseTokenToKVPoolAllocator
+
+if TYPE_CHECKING:
+    from sglang.srt.mem_cache.memory_pool import KVCache
+
+
+class TokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
+    """An allocator managing the indices to kv cache data."""
+
+    def __init__(
+        self,
+        size: int,
+        dtype: torch.dtype,
+        device: str,
+        kvcache: KVCache,
+        need_sort: bool,
+    ):
+        super().__init__(size, 1, dtype, device, kvcache, need_sort)
+        self.clear()
+
+    def clear(self):
+        # The padded slot 0 is used for writing dummy outputs from padded tokens.
+        self.free_pages = torch.arange(
+            1, self.size + 1, dtype=torch.int64, device=self.device
+        )
+        self.is_not_in_free_group = True
+        self.free_group = []
+        self.release_pages = torch.empty((0,), dtype=torch.int64, device=self.device)
+
+    def available_size(self):
+        # To avoid minor "len(free_pages) * 1" overhead
+        return len(self.free_pages) + len(self.release_pages)
+
+    def alloc(self, need_size: int):
+        if self.need_sort and need_size > len(self.free_pages):
+            self.merge_and_sort_free()
+
+        if need_size > len(self.free_pages):
+            return None
+
+        select_index = self.free_pages[:need_size]
+        self.free_pages = self.free_pages[need_size:]
+        return select_index
+
+    def free(self, free_index: torch.Tensor):
+        if free_index.numel() == 0:
+            return
+
+        if self.is_not_in_free_group:
+            if self.need_sort:
+                self.release_pages = torch.cat((self.release_pages, free_index))
+            else:
+                self.free_pages = torch.cat((self.free_pages, free_index))
+        else:
+            self.free_group.append(free_index)
+
+    def get_cpu_copy(self, indices, mamba_indices=None):
+        return self._kvcache.get_cpu_copy(indices, mamba_indices=mamba_indices)
+
+    def load_cpu_copy(self, kv_cache_cpu, indices, mamba_indices=None):
+        return self._kvcache.load_cpu_copy(
+            kv_cache_cpu, indices, mamba_indices=mamba_indices
+        )