[mem_cache][1/N] refactor: split allocator.py into allocator/ subpackage (#26675)
This commit is contained in:
shuwenn
@@ -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",
|
||||
]
|
||||
@@ -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()
|
||||
+3
-148
@@ -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,
|
||||
@@ -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
|
||||
)
|
||||
Reference in New Issue
Block a user