From 1af53f67b24aeb8e8e39e9326304c452a76c35e7 Mon Sep 17 00:00:00 2001 From: shuwenn <47200617+alphabetc1@users.noreply.github.com> Date: Thu, 4 Jun 2026 15:10:37 +0800 Subject: [PATCH] [mem_cache][2/N] refactor: move SWATokenToKVPoolAllocator to allocator/swa.py (#26676) --- python/sglang/srt/kv_canary/api.py | 2 +- .../srt/kv_canary/runner/canary_manager.py | 2 +- .../srt/kv_canary/runner/swa_divergence.py | 2 +- .../layers/attention/flashinfer_backend.py | 2 +- python/sglang/srt/managers/schedule_policy.py | 2 +- python/sglang/srt/mem_cache/allocator/swa.py | 369 ++++++++++++++++++ python/sglang/srt/mem_cache/chunk_cache.py | 2 +- python/sglang/srt/mem_cache/common.py | 2 +- .../sglang/srt/mem_cache/swa_memory_pool.py | 368 ----------------- .../sglang/srt/mem_cache/swa_radix_cache.py | 2 +- .../unified_cache_components/swa_component.py | 2 +- .../model_runner_kv_cache_mixin.py | 3 +- .../test_swa_alloc_extend_page_estimation.py | 2 +- .../mem_cache/test_swa_eviction_boundary.py | 3 +- .../test_swa_lock_release_lifecycle.py | 3 +- .../unit/mem_cache/test_swa_unittest.py | 3 +- .../test_unified_radix_cache_bench.py | 6 +- .../test_unified_radix_cache_unittest.py | 3 +- 18 files changed, 391 insertions(+), 387 deletions(-) create mode 100644 python/sglang/srt/mem_cache/allocator/swa.py diff --git a/python/sglang/srt/kv_canary/api.py b/python/sglang/srt/kv_canary/api.py index 45be28d46..2800c4614 100644 --- a/python/sglang/srt/kv_canary/api.py +++ b/python/sglang/srt/kv_canary/api.py @@ -11,7 +11,7 @@ from sglang.srt.kv_canary.perturb.config import PerturbConfig from sglang.srt.kv_canary.pool_patcher.api import attach_canary_buffers from sglang.srt.kv_canary.pool_patcher.utils import wrap_method from sglang.srt.kv_canary.runner.canary_manager import CanaryManager -from sglang.srt.mem_cache.swa_memory_pool import SWATokenToKVPoolAllocator +from sglang.srt.mem_cache.allocator.swa import SWATokenToKVPoolAllocator from sglang.srt.model_executor.forward_batch_info import ForwardBatch if TYPE_CHECKING: diff --git a/python/sglang/srt/kv_canary/runner/canary_manager.py b/python/sglang/srt/kv_canary/runner/canary_manager.py index 0b4656cfc..04c27eecf 100644 --- a/python/sglang/srt/kv_canary/runner/canary_manager.py +++ b/python/sglang/srt/kv_canary/runner/canary_manager.py @@ -31,9 +31,9 @@ from sglang.srt.kv_canary.state import CanaryDeviceState from sglang.srt.kv_canary.token_oracle.oracle_manager import TokenOracleManager if TYPE_CHECKING: + from sglang.srt.mem_cache.allocator.swa import SWATokenToKVPoolAllocator from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache from sglang.srt.mem_cache.memory_pool import ReqToTokenPool - from sglang.srt.mem_cache.swa_memory_pool import SWATokenToKVPoolAllocator from sglang.srt.model_executor.forward_batch_info import ForwardBatch logger = logging.getLogger(__name__) diff --git a/python/sglang/srt/kv_canary/runner/swa_divergence.py b/python/sglang/srt/kv_canary/runner/swa_divergence.py index 9cbcb816e..412681c26 100644 --- a/python/sglang/srt/kv_canary/runner/swa_divergence.py +++ b/python/sglang/srt/kv_canary/runner/swa_divergence.py @@ -13,8 +13,8 @@ from sglang.srt.kv_canary.buffer_group import CanaryBufferGroup, PoolKind from sglang.srt.kv_canary.runner.future_tensor import DelayedDeviceHostHandler if TYPE_CHECKING: + from sglang.srt.mem_cache.allocator.swa import SWATokenToKVPoolAllocator from sglang.srt.mem_cache.memory_pool import ReqToTokenPool - from sglang.srt.mem_cache.swa_memory_pool import SWATokenToKVPoolAllocator from sglang.srt.model_executor.forward_batch_info import ForwardBatch logger = logging.getLogger(__name__) diff --git a/python/sglang/srt/layers/attention/flashinfer_backend.py b/python/sglang/srt/layers/attention/flashinfer_backend.py index 93943f533..567ffadad 100644 --- a/python/sglang/srt/layers/attention/flashinfer_backend.py +++ b/python/sglang/srt/layers/attention/flashinfer_backend.py @@ -24,7 +24,7 @@ from sglang.srt.layers.attention.base_attn_backend import AttentionBackend from sglang.srt.layers.attention.utils import create_flashinfer_kv_indices_triton from sglang.srt.layers.dp_attention import get_attention_tp_size from sglang.srt.layers.radix_attention import AttentionType -from sglang.srt.mem_cache.swa_memory_pool import SWATokenToKVPoolAllocator +from sglang.srt.mem_cache.allocator.swa import SWATokenToKVPoolAllocator from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode from sglang.srt.speculative.spec_info import SpecInput from sglang.srt.utils import ( diff --git a/python/sglang/srt/managers/schedule_policy.py b/python/sglang/srt/managers/schedule_policy.py index 43975c57e..7e1a74b7a 100644 --- a/python/sglang/srt/managers/schedule_policy.py +++ b/python/sglang/srt/managers/schedule_policy.py @@ -38,6 +38,7 @@ from sglang.srt.dllm.config import DllmConfig from sglang.srt.layers.attention.dsa.utils import is_dsa_prefill_cp_in_seq_split from sglang.srt.layers.utils.cp_utils import is_prefill_context_parallel_enabled from sglang.srt.managers.schedule_batch import Req, ScheduleBatch +from sglang.srt.mem_cache.allocator.swa import SWATokenToKVPoolAllocator from sglang.srt.mem_cache.base_prefix_cache import ( BasePrefixCache, InitLoadBackParams, @@ -49,7 +50,6 @@ from sglang.srt.mem_cache.hisparse_memory_pool import ( DeepSeekV4HiSparseTokenToKVPoolAllocator, ) from sglang.srt.mem_cache.radix_cache import RadixCache, RadixKey, TreeNode -from sglang.srt.mem_cache.swa_memory_pool import SWATokenToKVPoolAllocator from sglang.srt.server_args import ServerArgs, get_global_server_args if TYPE_CHECKING: diff --git a/python/sglang/srt/mem_cache/allocator/swa.py b/python/sglang/srt/mem_cache/allocator/swa.py new file mode 100644 index 000000000..2d4a240be --- /dev/null +++ b/python/sglang/srt/mem_cache/allocator/swa.py @@ -0,0 +1,369 @@ +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: + 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() + + @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 + + if _is_npu: + self.full_to_swa_index_mapping[alloc_full_indices.to(torch.int64)] = ( + alloc_swa_indices.to(torch.int64) + ) + else: + self.full_to_swa_index_mapping[alloc_full_indices] = alloc_swa_indices + return alloc_full_indices + + 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 num_new_pages > self.full_attn_allocator.available_size() // self.page_size: + return None + if num_new_pages > self.swa_attn_allocator.available_size() // self.page_size: + 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 + + if _is_npu: + self.full_to_swa_index_mapping[alloc_full_indices.to(torch.int64)] = ( + alloc_swa_indices.to(torch.int64) + ) + else: + self.full_to_swa_index_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 num_full_pages > self.full_attn_allocator.available_size() // self.page_size: + return None + if num_swa_pages > self.swa_attn_allocator.available_size() // self.page_size: + 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, + ) + 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, + ) + assert alloc_swa_indices is not None + + self.full_to_swa_index_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]] = 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: + self.full_to_swa_index_mapping[alloc_full_indices.to(torch.int64)] = ( + 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() + if _is_npu: + self.full_to_swa_index_mapping[full_indices.to(torch.int64)] = ( + swa_indices.to(torch.int64) + ) + else: + self.full_to_swa_index_mapping[full_indices] = swa_indices + + def free_swa(self, free_index: torch.Tensor): + swa_indices = self.full_to_swa_index_mapping[free_index] + swa_indices = swa_indices[swa_indices > 0] + self.swa_attn_allocator.free(swa_indices) + self.full_to_swa_index_mapping[free_index] = 0 + + 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 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 + ) diff --git a/python/sglang/srt/mem_cache/chunk_cache.py b/python/sglang/srt/mem_cache/chunk_cache.py index 350a2ec8b..bd1351a38 100644 --- a/python/sglang/srt/mem_cache/chunk_cache.py +++ b/python/sglang/srt/mem_cache/chunk_cache.py @@ -7,6 +7,7 @@ from typing import TYPE_CHECKING, Any, Optional import torch +from sglang.srt.mem_cache.allocator.swa import SWATokenToKVPoolAllocator from sglang.srt.mem_cache.base_prefix_cache import ( BasePrefixCache, DecLockRefParams, @@ -22,7 +23,6 @@ from sglang.srt.mem_cache.base_prefix_cache import ( from sglang.srt.mem_cache.hisparse_memory_pool import ( DeepSeekV4HiSparseTokenToKVPoolAllocator, ) -from sglang.srt.mem_cache.swa_memory_pool import SWATokenToKVPoolAllocator if TYPE_CHECKING: from sglang.srt.managers.schedule_batch import Req diff --git a/python/sglang/srt/mem_cache/common.py b/python/sglang/srt/mem_cache/common.py index 49aed56c5..11a52acdc 100644 --- a/python/sglang/srt/mem_cache/common.py +++ b/python/sglang/srt/mem_cache/common.py @@ -6,9 +6,9 @@ from typing import TYPE_CHECKING import numpy as np import torch +from sglang.srt.mem_cache.allocator.swa import SWATokenToKVPoolAllocator from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache, EvictParams from sglang.srt.mem_cache.memory_pool import HybridReqToTokenPool, ReqToTokenPool -from sglang.srt.mem_cache.swa_memory_pool import SWATokenToKVPoolAllocator from sglang.srt.mem_cache.triton_ops.common import ( _get_last_loc_safe_kernel as _get_last_loc_safe_kernel, ) diff --git a/python/sglang/srt/mem_cache/swa_memory_pool.py b/python/sglang/srt/mem_cache/swa_memory_pool.py index 550bb53ac..ed98ececd 100644 --- a/python/sglang/srt/mem_cache/swa_memory_pool.py +++ b/python/sglang/srt/mem_cache/swa_memory_pool.py @@ -4,23 +4,9 @@ from typing import Dict, List, Optional, Tuple import torch from sglang.srt.layers.radix_attention import RadixAttention -from sglang.srt.mem_cache.allocator import ( - BaseTokenToKVPoolAllocator, - PagedTokenToKVPoolAllocator, - TokenToKVPoolAllocator, -) from sglang.srt.mem_cache.base_swa_memory_pool import BaseSWAKVPool from sglang.srt.mem_cache.memory_pool import KVCache, MHATokenToKVPool from sglang.srt.mem_cache.utils import maybe_init_custom_mem_pool -from sglang.srt.utils import is_npu -from sglang.srt.utils.common import get_num_new_pages - -_is_npu = is_npu() - -if _is_npu: - from sglang.srt.hardware_backend.npu.allocator_npu import ( - NPUPagedTokenToKVPoolAllocator, - ) logger = logging.getLogger(__name__) GB = 1024 * 1024 * 1024 @@ -280,357 +266,3 @@ class SWAKVPool(BaseSWAKVPool): swa_kv_cpu = self._filter_swa_cpu_copy(swa_kv_cpu, row_mask) self.swa_kv_pool.load_cpu_copy(swa_kv_cpu, swa_indices) - - -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() - - @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 - - if _is_npu: - self.full_to_swa_index_mapping[alloc_full_indices.to(torch.int64)] = ( - alloc_swa_indices.to(torch.int64) - ) - else: - self.full_to_swa_index_mapping[alloc_full_indices] = alloc_swa_indices - return alloc_full_indices - - 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 num_new_pages > self.full_attn_allocator.available_size() // self.page_size: - return None - if num_new_pages > self.swa_attn_allocator.available_size() // self.page_size: - 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 - - if _is_npu: - self.full_to_swa_index_mapping[alloc_full_indices.to(torch.int64)] = ( - alloc_swa_indices.to(torch.int64) - ) - else: - self.full_to_swa_index_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 num_full_pages > self.full_attn_allocator.available_size() // self.page_size: - return None - if num_swa_pages > self.swa_attn_allocator.available_size() // self.page_size: - 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, - ) - 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, - ) - assert alloc_swa_indices is not None - - self.full_to_swa_index_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]] = 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: - self.full_to_swa_index_mapping[alloc_full_indices.to(torch.int64)] = ( - 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() - if _is_npu: - self.full_to_swa_index_mapping[full_indices.to(torch.int64)] = ( - swa_indices.to(torch.int64) - ) - else: - self.full_to_swa_index_mapping[full_indices] = swa_indices - - def free_swa(self, free_index: torch.Tensor): - swa_indices = self.full_to_swa_index_mapping[free_index] - swa_indices = swa_indices[swa_indices > 0] - self.swa_attn_allocator.free(swa_indices) - self.full_to_swa_index_mapping[free_index] = 0 - - 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 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 - ) diff --git a/python/sglang/srt/mem_cache/swa_radix_cache.py b/python/sglang/srt/mem_cache/swa_radix_cache.py index dd76ad37d..cd428f86c 100644 --- a/python/sglang/srt/mem_cache/swa_radix_cache.py +++ b/python/sglang/srt/mem_cache/swa_radix_cache.py @@ -28,6 +28,7 @@ import torch from numpy import float64 from sglang.srt.environ import envs +from sglang.srt.mem_cache.allocator.swa import SWATokenToKVPoolAllocator from sglang.srt.mem_cache.base_prefix_cache import ( BasePrefixCache, DecLockRefParams, @@ -43,7 +44,6 @@ from sglang.srt.mem_cache.base_prefix_cache import ( from sglang.srt.mem_cache.cache_init_params import CacheInitParams from sglang.srt.mem_cache.events import KVCacheEventMixin from sglang.srt.mem_cache.radix_cache import RadixKey -from sglang.srt.mem_cache.swa_memory_pool import SWATokenToKVPoolAllocator from sglang.srt.mem_cache.utils import split_node_hash_value if TYPE_CHECKING: diff --git a/python/sglang/srt/mem_cache/unified_cache_components/swa_component.py b/python/sglang/srt/mem_cache/unified_cache_components/swa_component.py index d6b7cb6c3..6c1c52c8b 100644 --- a/python/sglang/srt/mem_cache/unified_cache_components/swa_component.py +++ b/python/sglang/srt/mem_cache/unified_cache_components/swa_component.py @@ -49,7 +49,7 @@ class SWAComponent(TreeComponent): """ def __init__(self, cache: UnifiedRadixCache, params: CacheInitParams): - from sglang.srt.mem_cache.swa_memory_pool import SWATokenToKVPoolAllocator + from sglang.srt.mem_cache.allocator.swa import SWATokenToKVPoolAllocator assert isinstance( cache.token_to_kv_pool_allocator, SWATokenToKVPoolAllocator diff --git a/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py b/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py index d291367d1..b3683bec2 100644 --- a/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py +++ b/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py @@ -17,6 +17,7 @@ from sglang.srt.mem_cache.allocator import ( PagedTokenToKVPoolAllocator, TokenToKVPoolAllocator, ) +from sglang.srt.mem_cache.allocator.swa import SWATokenToKVPoolAllocator from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool from sglang.srt.mem_cache.hisparse_memory_pool import ( DeepSeekV4HiSparseTokenToKVPoolAllocator, @@ -34,7 +35,7 @@ from sglang.srt.mem_cache.memory_pool import ( NoOpMHATokenToKVPool, ReqToTokenPool, ) -from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool, SWATokenToKVPoolAllocator +from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool from sglang.srt.utils.common import ( get_available_gpu_memory, is_float4_e2m1fn_x2, diff --git a/test/registered/unit/mem_cache/test_swa_alloc_extend_page_estimation.py b/test/registered/unit/mem_cache/test_swa_alloc_extend_page_estimation.py index 15ff77ea3..3700ca73b 100644 --- a/test/registered/unit/mem_cache/test_swa_alloc_extend_page_estimation.py +++ b/test/registered/unit/mem_cache/test_swa_alloc_extend_page_estimation.py @@ -11,7 +11,7 @@ from unittest.mock import MagicMock import torch -from sglang.srt.mem_cache.swa_memory_pool import SWATokenToKVPoolAllocator +from sglang.srt.mem_cache.allocator.swa import SWATokenToKVPoolAllocator from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.test_utils import CustomTestCase diff --git a/test/registered/unit/mem_cache/test_swa_eviction_boundary.py b/test/registered/unit/mem_cache/test_swa_eviction_boundary.py index ef643c8bb..c51fba5e9 100644 --- a/test/registered/unit/mem_cache/test_swa_eviction_boundary.py +++ b/test/registered/unit/mem_cache/test_swa_eviction_boundary.py @@ -19,9 +19,10 @@ from types import SimpleNamespace import torch from sglang.srt.managers.schedule_batch import ScheduleBatch +from sglang.srt.mem_cache.allocator.swa import SWATokenToKVPoolAllocator from sglang.srt.mem_cache.cache_init_params import CacheInitParams from sglang.srt.mem_cache.memory_pool import ReqToTokenPool -from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool, SWATokenToKVPoolAllocator +from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool from sglang.srt.mem_cache.swa_radix_cache import SWARadixCache from sglang.srt.utils import get_device from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci diff --git a/test/registered/unit/mem_cache/test_swa_lock_release_lifecycle.py b/test/registered/unit/mem_cache/test_swa_lock_release_lifecycle.py index 180e54c5f..f2d05f9b6 100644 --- a/test/registered/unit/mem_cache/test_swa_lock_release_lifecycle.py +++ b/test/registered/unit/mem_cache/test_swa_lock_release_lifecycle.py @@ -16,6 +16,7 @@ from array import array import torch +from sglang.srt.mem_cache.allocator.swa import SWATokenToKVPoolAllocator from sglang.srt.mem_cache.base_prefix_cache import ( DecLockRefParams, EvictParams, @@ -25,7 +26,7 @@ from sglang.srt.mem_cache.base_prefix_cache import ( from sglang.srt.mem_cache.cache_init_params import CacheInitParams from sglang.srt.mem_cache.memory_pool import ReqToTokenPool from sglang.srt.mem_cache.radix_cache import RadixKey -from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool, SWATokenToKVPoolAllocator +from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool from sglang.srt.mem_cache.swa_radix_cache import SWARadixCache from sglang.srt.utils import get_device from sglang.test.ci.ci_register import register_cuda_ci diff --git a/test/registered/unit/mem_cache/test_swa_unittest.py b/test/registered/unit/mem_cache/test_swa_unittest.py index aac705945..1a1ee245a 100644 --- a/test/registered/unit/mem_cache/test_swa_unittest.py +++ b/test/registered/unit/mem_cache/test_swa_unittest.py @@ -5,6 +5,7 @@ import torch from sglang.srt.disaggregation.kv_events import BlockRemoved, BlockStored from sglang.srt.environ import envs +from sglang.srt.mem_cache.allocator.swa import SWATokenToKVPoolAllocator from sglang.srt.mem_cache.base_prefix_cache import ( DecLockRefParams, EvictParams, @@ -16,7 +17,7 @@ from sglang.srt.mem_cache.cache_init_params import CacheInitParams from sglang.srt.mem_cache.common import available_and_evictable_str from sglang.srt.mem_cache.memory_pool import ReqToTokenPool from sglang.srt.mem_cache.radix_cache import RadixKey -from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool, SWATokenToKVPoolAllocator +from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool from sglang.srt.mem_cache.swa_radix_cache import SWARadixCache from sglang.srt.utils import get_device from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci diff --git a/test/registered/unit/mem_cache/test_unified_radix_cache_bench.py b/test/registered/unit/mem_cache/test_unified_radix_cache_bench.py index 7516bfa48..c7a478eb0 100644 --- a/test/registered/unit/mem_cache/test_unified_radix_cache_bench.py +++ b/test/registered/unit/mem_cache/test_unified_radix_cache_bench.py @@ -178,10 +178,8 @@ def create_bench_cache( # --- KV pool + allocator --- if has_swa: - from sglang.srt.mem_cache.swa_memory_pool import ( - SWAKVPool, - SWATokenToKVPoolAllocator, - ) + from sglang.srt.mem_cache.allocator.swa import SWATokenToKVPoolAllocator + from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool pool = SWAKVPool( size=kv_size, diff --git a/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py b/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py index 5edfbdb6e..af971d9af 100644 --- a/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py +++ b/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py @@ -22,6 +22,7 @@ from sglang.srt.environ import envs from sglang.srt.layers.attention.fla.chunk_delta_h import CHUNK_SIZE as FLA_CHUNK_SIZE from sglang.srt.managers.schedule_batch import Req from sglang.srt.mem_cache.allocator import TokenToKVPoolAllocator +from sglang.srt.mem_cache.allocator.swa import SWATokenToKVPoolAllocator from sglang.srt.mem_cache.base_prefix_cache import ( DecLockRefParams, EvictParams, @@ -41,7 +42,7 @@ from sglang.srt.mem_cache.memory_pool import ( ReqToTokenPool, ) from sglang.srt.mem_cache.radix_cache import RadixKey -from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool, SWATokenToKVPoolAllocator +from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool from sglang.srt.mem_cache.unified_cache_components.tree_component import ( CacheTransferPhase, ComponentType,