diff --git a/python/sglang/srt/mem_cache/allocator/__init__.py b/python/sglang/srt/mem_cache/allocator/__init__.py new file mode 100644 index 000000000..ddf427ec7 --- /dev/null +++ b/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", +] diff --git a/python/sglang/srt/mem_cache/allocator/base.py b/python/sglang/srt/mem_cache/allocator/base.py new file mode 100644 index 000000000..949d92d73 --- /dev/null +++ b/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() diff --git a/python/sglang/srt/mem_cache/allocator.py b/python/sglang/srt/mem_cache/allocator/paged.py similarity index 73% rename from python/sglang/srt/mem_cache/allocator.py rename to python/sglang/srt/mem_cache/allocator/paged.py index 7ec8bd226..0d623f159 100755 --- a/python/sglang/srt/mem_cache/allocator.py +++ b/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, diff --git a/python/sglang/srt/mem_cache/allocator/token.py b/python/sglang/srt/mem_cache/allocator/token.py new file mode 100644 index 000000000..990e468dd --- /dev/null +++ b/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 + )