import torch from sglang.srt.mem_cache.allocator.base import BaseTokenToKVPoolAllocator from sglang.srt.mem_cache.allocator.paged import PagedTokenToKVPoolAllocator from sglang.srt.mem_cache.allocator.token import TokenToKVPoolAllocator from sglang.srt.mem_cache.base_swa_memory_pool import BaseSWAKVPool from sglang.srt.utils import is_npu from sglang.srt.utils.common import get_num_new_pages _is_npu = is_npu() if _is_npu: import torch_npu from sglang.srt.hardware_backend.npu.allocator_npu import ( NPUPagedTokenToKVPoolAllocator, ) class SWATokenToKVPoolAllocator(BaseTokenToKVPoolAllocator): """Allocator for SWA hybrid KV cache.""" def __init__( self, size: int, size_swa: int, page_size: int, dtype: torch.dtype, device: str, kvcache: BaseSWAKVPool, need_sort: bool, ): assert isinstance(kvcache, BaseSWAKVPool) self._size_full = size self._size_swa = size_swa self.dtype = dtype self.device = device self.page_size = page_size full_kv_pool = getattr(kvcache, "full_kv_pool", None) swa_kv_pool = getattr(kvcache, "swa_kv_pool", None) if page_size == 1: self.full_attn_allocator = TokenToKVPoolAllocator( size, dtype, device, full_kv_pool, need_sort, ) self.swa_attn_allocator = TokenToKVPoolAllocator( size_swa, dtype, device, swa_kv_pool, need_sort, ) else: if _is_npu: PagedTokenToKVPoolAllocatorClass = NPUPagedTokenToKVPoolAllocator else: PagedTokenToKVPoolAllocatorClass = PagedTokenToKVPoolAllocator self.full_attn_allocator = PagedTokenToKVPoolAllocatorClass( size, page_size, dtype, device, full_kv_pool, need_sort, ) self.swa_attn_allocator = PagedTokenToKVPoolAllocatorClass( size_swa, page_size, dtype, device, 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._kvcache = kvcache self.clear() self._kvcache.register_mapping(self.full_to_swa_index_mapping) def available_size(self): return min( self.full_attn_allocator.available_size(), self.swa_attn_allocator.available_size(), ) def full_available_size(self): return self.full_attn_allocator.available_size() def swa_available_size(self): return self.swa_attn_allocator.available_size() # Slot-conservation views for the leak invariant. On the non-shared allocator # the static budget IS physical (conserve == physical); the shared composite # overrides these with the static-cap view. def _conserve_full_available_size(self): return self.full_available_size() def _conserve_swa_available_size(self): return self.swa_available_size() @property def size(self): 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()}, " msg += ( f"#full-attn-available-size: {self.full_attn_allocator.available_size()}, " ) return msg def get_kvcache(self): return self._kvcache def translate_loc_from_full_to_swa(self, kv_indices: torch.Tensor): assert self._kvcache.full_to_swa_index_mapping is not None 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(): return None 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.set_full_to_swa_mapping(alloc_full_indices, alloc_swa_indices) return alloc_full_indices def new_pages_available(self, num_full_pages: int, num_swa_pages: int) -> bool: return ( num_full_pages <= self.full_attn_allocator.available_size() // self.page_size and num_swa_pages <= self.swa_attn_allocator.available_size() // self.page_size ) 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_new_pages = get_num_new_pages( seq_lens=seq_lens_cpu, page_size=self.page_size, prefix_lens=prefix_lens_cpu ) if not self.new_pages_available(num_new_pages, num_new_pages): 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, num_new_pages=num_new_pages, ) 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, num_new_pages=num_new_pages, ) assert alloc_full_indices is not None assert alloc_swa_indices is not None self.set_full_to_swa_mapping(alloc_full_indices, alloc_swa_indices) return alloc_full_indices def alloc_extend_swa_tail( 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, swa_tail_len: int, ): """Allocate full KV for the whole extend and SWA KV only for the tail. This is used by disaggregated decode preallocation: decode receives full prompt KV for full-attention layers, but only the sliding-window state is transferred for SWA layers. """ assert self.page_size > 1 assert len(seq_lens_cpu) == 1, "SWA tail allocation currently supports bs=1" assert len(prefix_lens_cpu) == 1 assert 0 <= swa_tail_len <= extend_num_tokens num_full_pages = get_num_new_pages( seq_lens=seq_lens_cpu, page_size=self.page_size, prefix_lens=prefix_lens_cpu ) num_swa_pages = (swa_tail_len + self.page_size - 1) // self.page_size if not self.new_pages_available(num_full_pages, num_swa_pages): return None alloc_full_indices = self.full_attn_allocator.alloc_extend( prefix_lens, prefix_lens_cpu, seq_lens, seq_lens_cpu, last_loc, extend_num_tokens, num_new_pages=num_full_pages, ) assert alloc_full_indices is not None if swa_tail_len == 0: return alloc_full_indices device = self.device swa_prefix_lens = torch.zeros((1,), dtype=torch.int64, device=device) swa_prefix_lens_cpu = torch.zeros((1,), dtype=torch.int64) swa_seq_lens = torch.tensor([swa_tail_len], dtype=torch.int64, device=device) swa_seq_lens_cpu = torch.tensor([swa_tail_len], dtype=torch.int64) swa_last_loc = torch.tensor([-1], dtype=torch.int64, device=device) alloc_swa_indices = self.swa_attn_allocator.alloc_extend( swa_prefix_lens, swa_prefix_lens_cpu, swa_seq_lens, swa_seq_lens_cpu, swa_last_loc, swa_tail_len, num_new_pages=num_swa_pages, ) assert alloc_swa_indices is not None self.set_full_to_swa_mapping( alloc_full_indices[-swa_tail_len:], alloc_swa_indices ) if swa_tail_len < extend_num_tokens: self.full_to_swa_index_mapping[ alloc_full_indices[:-swa_tail_len].to(torch.int64) ] = 0 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 if _is_npu: indices_2d = alloc_full_indices.to(torch.int64).unsqueeze(-1) torch_npu.npu_scatter_nd_update_( self.full_to_swa_index_mapping, indices_2d, alloc_swa_indices.to(torch.int64), ) else: 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) else: self.free_group.append(free_index) assert ( self.full_attn_allocator.available_size() <= self.full_attn_allocator.size ) assert self.swa_attn_allocator.available_size() <= self.swa_attn_allocator.size def set_full_to_swa_mapping( self, full_indices: torch.Tensor, swa_indices: torch.Tensor ) -> None: """Write full_to_swa_index_mapping[full_indices[i]] = swa_indices[i]. Used by HiCache load-back path to rebuild the mapping after FULL and SWA device alloc. """ if full_indices.numel() == 0: return assert full_indices.numel() == swa_indices.numel() full_indices = full_indices.to(torch.int64) swa_indices = swa_indices.to(self.full_to_swa_index_mapping.dtype) self.full_to_swa_index_mapping[full_indices] = swa_indices def free_swa(self, free_index: torch.Tensor): if free_index.numel() == 0: return if self.page_size == 1: mapping_indices = free_index else: mapping_indices = self._expand_to_full_pages(free_index) swa_indices = self.full_to_swa_index_mapping[mapping_indices] swa_indices = swa_indices[swa_indices > 0] self.swa_attn_allocator.free(swa_indices) self.full_to_swa_index_mapping[mapping_indices] = 0 def _expand_to_full_pages(self, indices: torch.Tensor) -> torch.Tensor: pages = torch.unique(indices // self.page_size) page_offsets = torch.arange( self.page_size, dtype=indices.dtype, device=indices.device ) return (pages[:, None] * self.page_size + page_offsets[None, :]).reshape(-1) def backup_state(self): return [ self.full_attn_allocator.backup_state(), self.swa_attn_allocator.backup_state(), ] def restore_state(self, state): assert len(state) == 2 self.full_attn_allocator.restore_state(state[0]) self.swa_attn_allocator.restore_state(state[1]) def resize(self, config) -> None: size_full = int(config.full_max_total_num_tokens) size_swa = int(config.swa_max_total_num_tokens) self._size_full = size_full self._size_swa = size_swa for alloc, sz in ( (self.full_attn_allocator, size_full), (self.swa_attn_allocator, size_swa), ): alloc.size = int(sz) if self.page_size > 1: alloc.num_pages = int(sz) // self.page_size self.clear() def clear(self): self.swa_attn_allocator.clear() self.full_attn_allocator.clear() # 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 = [] 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 ) class PureSWATokenToKVPoolAllocator(SWATokenToKVPoolAllocator): """Single-pool allocator for models whose every layer is sliding-window attention.""" def __init__( self, size_swa: int, page_size: int, dtype: torch.dtype, device: str, kvcache: BaseSWAKVPool, need_sort: bool, ): assert page_size == 1 assert isinstance(kvcache, BaseSWAKVPool) self.page_size = page_size self.dtype = dtype self.device = device self.need_sort = need_sort self._size_full = self._size_swa = size_swa self.swa_attn_allocator = TokenToKVPoolAllocator( size_swa, dtype, device, kvcache.swa_kv_pool, need_sort, ) self.full_attn_allocator = self.swa_attn_allocator self.full_to_swa_index_mapping = torch.cat( [ torch.arange(size_swa + page_size, dtype=torch.int64, device=device), torch.tensor([-1], dtype=torch.int64, device=device), ] ) self.free_pages = None self.release_pages = None self.is_not_in_free_group = True self.free_group = [] self._kvcache = kvcache self.swa_attn_allocator.clear() self._kvcache.register_mapping(self.full_to_swa_index_mapping) def available_size(self): return self.swa_attn_allocator.available_size() def full_available_size(self): return self.swa_attn_allocator.available_size() def swa_available_size(self): return self.swa_attn_allocator.available_size() def new_pages_available(self, num_full_pages: int, num_swa_pages: int) -> bool: avail = self.swa_attn_allocator.available_size() // self.page_size return num_full_pages <= avail and num_swa_pages <= avail def translate_loc_from_full_to_swa(self, kv_indices: torch.Tensor): return kv_indices def alloc(self, need_size: int): assert self.page_size == 1 return self.swa_attn_allocator.alloc(need_size) def alloc_extend(self, *args, **kwargs): raise NotImplementedError( "PureSWATokenToKVPoolAllocator does not support page_size > 1." ) def alloc_decode(self, *args, **kwargs): raise NotImplementedError( "PureSWATokenToKVPoolAllocator does not support page_size > 1." ) def alloc_extend_swa_tail(self, *args, **kwargs): raise NotImplementedError( "PureSWATokenToKVPoolAllocator does not support page_size > 1." ) def free(self, free_index: torch.Tensor): if free_index.numel() == 0: return if self.is_not_in_free_group: self.swa_attn_allocator.free(free_index[free_index > 0]) else: self.free_group.append(free_index) assert self.swa_attn_allocator.available_size() <= self.swa_attn_allocator.size def free_swa(self, free_index: torch.Tensor): if free_index.numel() == 0: return self.swa_attn_allocator.free(free_index[free_index > 0]) 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)) self.free_group = [] def backup_state(self): return self.swa_attn_allocator.backup_state() def restore_state(self, state): self.swa_attn_allocator.restore_state(state) def clear(self): self.swa_attn_allocator.clear() self.is_not_in_free_group = True self.free_group = []