chore: import upstream snapshot with attribution

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

@@ -0,0 +1,116 @@
+"""
+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")
+
+    def resize(self, config) -> None:
+        self.size = config.max_total_num_tokens
+        if self.page_size > 1:
+            self.num_pages = config.max_total_num_tokens // self.page_size
+        self.clear()
+
+    @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()