Move swa memory pool to a seperate file (#16347)

python/sglang/srt/disaggregation/decode.py

@@ -58,8 +58,8 @@ from sglang.srt.mem_cache.memory_pool import (
     KVCache,
     NSATokenToKVPool,
     ReqToTokenPool,
-    SWAKVPool,
 )
+from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool
 from sglang.srt.tracing.trace import trace_event_batch, trace_slice_end
 from sglang.srt.utils import get_int_env_var
 from sglang.srt.utils.torch_memory_saver_adapter import TorchMemorySaverAdapter

python/sglang/srt/disaggregation/prefill.py

@@ -49,11 +49,8 @@ from sglang.srt.managers.schedule_batch import (
     ScheduleBatch,
 )
 from sglang.srt.mem_cache.common import release_kv_cache
-from sglang.srt.mem_cache.memory_pool import (
+from sglang.srt.mem_cache.memory_pool import HybridLinearKVPool, NSATokenToKVPool
-    HybridLinearKVPool,
+from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool
-    NSATokenToKVPool,
-    SWAKVPool,
-)
 from sglang.srt.tracing.trace import trace_event_batch, trace_slice, trace_slice_end
 
 if TYPE_CHECKING:

python/sglang/srt/layers/attention/flashattention_backend.py

@@ -11,7 +11,7 @@ import triton.language as tl
 from sglang.srt.configs.model_config import AttentionArch
 from sglang.srt.layers.attention.base_attn_backend import AttentionBackend
 from sglang.srt.layers.radix_attention import AttentionType
-from sglang.srt.mem_cache.memory_pool import SWAKVPool
+from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
 from sglang.srt.server_args import get_global_server_args
 from sglang.srt.speculative.spec_info import SpecInput

python/sglang/srt/layers/attention/flashinfer_backend.py

@@ -22,7 +22,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.allocator import SWATokenToKVPoolAllocator
+from sglang.srt.mem_cache.swa_memory_pool 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 (

python/sglang/srt/managers/schedule_batch.py

@@ -58,10 +58,7 @@ from sglang.srt.disaggregation.utils import DisaggregationMode
 from sglang.srt.distributed.parallel_state import get_tensor_model_parallel_rank
 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.mem_cache.allocator import (
+from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator
-    BaseTokenToKVPoolAllocator,
-    SWATokenToKVPoolAllocator,
-)
 from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache
 from sglang.srt.mem_cache.common import (
     alloc_for_decode,
@@ -72,6 +69,7 @@ from sglang.srt.mem_cache.common import (
 from sglang.srt.mem_cache.mamba_radix_cache import MambaRadixCache
 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 SWATokenToKVPoolAllocator
 from sglang.srt.metrics.collector import (
     DPCooperationInfo,
     SchedulerMetricsCollector,

python/sglang/srt/managers/schedule_policy.py

@@ -28,10 +28,10 @@ import torch
 
 from sglang.srt.layers.attention.nsa.utils import is_nsa_prefill_cp_in_seq_split
 from sglang.srt.managers.schedule_batch import Req, ScheduleBatch
-from sglang.srt.mem_cache.allocator import SWATokenToKVPoolAllocator
 from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache
 from sglang.srt.mem_cache.mamba_radix_cache import MambaRadixCache
 from sglang.srt.mem_cache.radix_cache import RadixCache, RadixKey, TreeNode
+from sglang.srt.mem_cache.swa_memory_pool import SWATokenToKVPoolAllocator
 from sglang.srt.mem_cache.swa_radix_cache import SWARadixCache
 from sglang.srt.server_args import ServerArgs
 

python/sglang/srt/mem_cache/allocator.py

@@ -20,14 +20,12 @@ Page-aligned memory pool.
 """
 
 import abc
-import weakref
 from typing import TYPE_CHECKING
 
 import torch
 import triton
 import triton.language as tl
 
-from sglang.srt.mem_cache.memory_pool import SWAKVPool
 from sglang.srt.utils import get_bool_env_var, get_num_new_pages, next_power_of_2
 
 if TYPE_CHECKING:
@@ -173,245 +171,6 @@ class TokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
         return self._kvcache.load_cpu_copy(kv_cache_cpu, 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: SWAKVPool,
-        need_sort: bool,
-    ):
-        assert isinstance(kvcache, SWAKVPool)
-        self._size_full = size
-        self._size_swa = size_swa
-        self.dtype = dtype
-        self.device = device
-        self.page_size = page_size
-
-        if page_size == 1:
-            self.full_attn_allocator = TokenToKVPoolAllocator(
-                size,
-                dtype,
-                device,
-                kvcache.full_kv_pool,
-                need_sort,
-            )
-            self.swa_attn_allocator = TokenToKVPoolAllocator(
-                size_swa,
-                dtype,
-                device,
-                kvcache.swa_kv_pool,
-                need_sort,
-            )
-        else:
-            self.full_attn_allocator = PagedTokenToKVPoolAllocator(
-                size,
-                page_size,
-                dtype,
-                device,
-                kvcache.full_kv_pool,
-                need_sort,
-            )
-            self.swa_attn_allocator = PagedTokenToKVPoolAllocator(
-                size_swa,
-                page_size,
-                dtype,
-                device,
-                kvcache.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.clear()
-        self._kvcache = kvcache
-        self._kvcache.register_mapping(weakref.proxy(self.full_to_swa_index_mapping))
-
-    def available_size(self):
-        # Note: use full_available_size() and swa_available_size() instead.
-        raise NotImplementedError()
-
-    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
-
-        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_tokens = extend_num_tokens + len(seq_lens) * self.page_size
-        if num_tokens > self.full_attn_allocator.available_size():
-            return None
-        if num_tokens > self.swa_attn_allocator.available_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,
-        )
-        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,
-        )
-        assert alloc_full_indices is not None
-        assert alloc_swa_indices is not None
-
-        self.full_to_swa_index_mapping[alloc_full_indices] = alloc_swa_indices
-
-        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
-
-        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 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):
-        return self._kvcache.get_cpu_copy(indices)
-
-    def load_cpu_copy(self, kv_cache_cpu, indices):
-        return self._kvcache.load_cpu_copy(kv_cache_cpu, indices)
-
-
 @triton.jit
 def alloc_extend_kernel(
     pre_lens_ptr,

python/sglang/srt/mem_cache/chunk_cache.py

@@ -7,8 +7,8 @@ from typing import TYPE_CHECKING, Any, Optional
 
 import torch
 
-from sglang.srt.mem_cache.allocator import SWATokenToKVPoolAllocator
 from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache, MatchResult
+from sglang.srt.mem_cache.swa_memory_pool import SWATokenToKVPoolAllocator
 
 if TYPE_CHECKING:
     from sglang.srt.managers.schedule_batch import Req

python/sglang/srt/mem_cache/common.py

@@ -7,11 +7,11 @@ import torch
 import triton
 import triton.language as tl
 
-from sglang.srt.mem_cache.allocator import SWATokenToKVPoolAllocator
 from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache
 from sglang.srt.mem_cache.chunk_cache import ChunkCache, SWAChunkCache
 from sglang.srt.mem_cache.mamba_radix_cache import MambaRadixCache
 from sglang.srt.mem_cache.memory_pool import HybridReqToTokenPool, ReqToTokenPool
+from sglang.srt.mem_cache.swa_memory_pool import SWATokenToKVPoolAllocator
 from sglang.srt.server_args import get_global_server_args
 from sglang.srt.utils import support_triton
 from sglang.srt.utils.common import ceil_align

python/sglang/srt/mem_cache/memory_pool.py

@@ -40,7 +40,7 @@ KVCache actually holds the physical kv cache.
 import abc
 import logging
 from contextlib import contextmanager, nullcontext
-from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Union
+from typing import TYPE_CHECKING, List, Optional, Tuple, Union
 
 import numpy as np
 import torch
@@ -1277,209 +1277,6 @@ class HybridLinearKVPool(KVCache):
             return self.full_kv_pool.get_mla_kv_buffer(layer, loc, dst_dtype)
 
 
-class SWAKVPool(KVCache):
-    """KV cache with separate pools for full and SWA attention layers."""
-
-    def __init__(
-        self,
-        size: int,
-        size_swa: int,
-        page_size: int,
-        dtype: torch.dtype,
-        head_num: int,
-        head_dim: int,
-        swa_attention_layer_ids: List[int],
-        full_attention_layer_ids: List[int],
-        enable_kvcache_transpose: bool,
-        device: str,
-        token_to_kv_pool_class: KVCache = MHATokenToKVPool,
-        **kwargs,
-    ):
-        self.size = size
-        self.size_swa = size_swa
-        self.dtype = dtype
-        self.head_num = head_num
-        self.head_dim = head_dim
-        self.device = device
-        self.swa_layer_nums = len(swa_attention_layer_ids)
-        self.full_layer_nums = len(full_attention_layer_ids)
-        self.start_layer = 0
-        self.page_size = page_size
-        self.swa_loc = None
-
-        kwargs["page_size"] = page_size
-        kwargs["enable_memory_saver"] = False
-        kwargs["head_num"] = head_num
-        kwargs["head_dim"] = head_dim
-        kwargs["device"] = device
-        # TODO MHATransposedTokenToKVPool if enable_kvcache_transpose is True
-        assert not enable_kvcache_transpose
-
-        # for disagg with nvlink
-        self.enable_custom_mem_pool, self.custom_mem_pool, _ = (
-            maybe_init_custom_mem_pool(device=self.device)
-        )
-
-        self.swa_kv_pool = token_to_kv_pool_class(
-            size=size_swa,
-            dtype=dtype,
-            layer_num=self.swa_layer_nums,
-            **kwargs,
-        )
-        kwargs.pop("swa_head_num", None)
-        kwargs.pop("swa_head_dim", None)
-        kwargs.pop("swa_v_head_dim", None)
-        self.full_kv_pool = token_to_kv_pool_class(
-            size=size,
-            dtype=dtype,
-            layer_num=self.full_layer_nums,
-            **kwargs,
-        )
-        # {layer_id: (index, is_swa_layer)}
-        self.layers_mapping: Dict[int, Tuple[int, bool]] = {}
-        for full_attn_layer_id, global_layer_id in enumerate(full_attention_layer_ids):
-            self.layers_mapping[global_layer_id] = (full_attn_layer_id, False)
-        for swa_layer_id, global_layer_id in enumerate(swa_attention_layer_ids):
-            self.layers_mapping[global_layer_id] = (swa_layer_id, True)
-        self.full_to_swa_index_mapping: Optional[torch.Tensor] = None
-
-        k_size, v_size = self.get_kv_size_bytes()
-        self.mem_usage = (k_size + v_size) / GB
-        logger.info(
-            f"SWAKVPool mem usage: {self.mem_usage:.2f} GB, swa size: {self.size_swa}, full size: {self.size}"
-        )
-
-    def register_mapping(self, full_to_swa_index_mapping: torch.Tensor):
-        self.full_to_swa_index_mapping = full_to_swa_index_mapping
-
-    def get_kv_size_bytes(self):
-        k_size, v_size = self.full_kv_pool.get_kv_size_bytes()
-        k_size_swa, v_size_swa = self.swa_kv_pool.get_kv_size_bytes()
-        return k_size + k_size_swa, v_size + v_size_swa
-
-    def get_contiguous_buf_infos(self):
-        full_kv_data_ptrs, full_kv_data_lens, full_kv_item_lens = (
-            self.full_kv_pool.get_contiguous_buf_infos()
-        )
-        return (
-            full_kv_data_ptrs,
-            full_kv_data_lens,
-            full_kv_item_lens,
-        )
-
-    def get_state_buf_infos(self):
-        swa_kv_data_ptrs, swa_kv_data_lens, swa_kv_item_lens = (
-            self.swa_kv_pool.get_contiguous_buf_infos()
-        )
-
-        return swa_kv_data_ptrs, swa_kv_data_lens, swa_kv_item_lens
-
-    def get_key_buffer(self, layer_id: int):
-        layer_id_pool, is_swa_layer = self.layers_mapping[layer_id]
-        if is_swa_layer:
-            return self.swa_kv_pool.get_key_buffer(layer_id_pool)
-        else:
-            return self.full_kv_pool.get_key_buffer(layer_id_pool)
-
-    def get_value_buffer(self, layer_id: int):
-        layer_id_pool, is_swa_layer = self.layers_mapping[layer_id]
-        if is_swa_layer:
-            return self.swa_kv_pool.get_value_buffer(layer_id_pool)
-        else:
-            return self.full_kv_pool.get_value_buffer(layer_id_pool)
-
-    def get_kv_buffer(self, layer_id: int):
-        layer_id_pool, is_swa_layer = self.layers_mapping[layer_id]
-        if is_swa_layer:
-            return self.swa_kv_pool.get_kv_buffer(layer_id_pool)
-        else:
-            return self.full_kv_pool.get_kv_buffer(layer_id_pool)
-
-    def set_swa_loc(self, loc: torch.Tensor):
-        self.swa_loc = loc
-
-    def translate_loc_from_full_to_swa(self, kv_indices: torch.Tensor):
-        assert self.full_to_swa_index_mapping is not None
-
-        # Note: kv_indices could have -1 values (from alloc_extend), which will be mapped to -1
-        # since the last item of full_to_swa_index_mapping is -1.
-        return self.full_to_swa_index_mapping[kv_indices].to(torch.int32)
-
-    def set_kv_buffer(
-        self,
-        layer: RadixAttention,
-        loc: torch.Tensor,
-        cache_k: torch.Tensor,
-        cache_v: torch.Tensor,
-        k_scale: float = 1.0,
-        v_scale: float = 1.0,
-    ):
-
-        layer_id = layer.layer_id
-        layer_id_pool, is_swa_layer = self.layers_mapping[layer_id]
-        if is_swa_layer:
-            if self.swa_loc is not None:
-                loc = self.swa_loc
-            else:
-                if self.full_to_swa_index_mapping is not None:
-                    loc = self.translate_loc_from_full_to_swa(loc)
-
-            self.swa_kv_pool.set_kv_buffer(
-                None,
-                loc,
-                cache_k,
-                cache_v,
-                k_scale,
-                v_scale,
-                layer_id_override=layer_id_pool,
-            )
-        else:
-            self.full_kv_pool.set_kv_buffer(
-                None,
-                loc,
-                cache_k,
-                cache_v,
-                k_scale,
-                v_scale,
-                layer_id_override=layer_id_pool,
-            )
-
-    def move_kv_cache(self, tgt_loc: torch.Tensor, src_loc: torch.Tensor):
-        self.full_kv_pool.move_kv_cache(tgt_loc, src_loc)
-        tgt_loc_swa = self.translate_loc_from_full_to_swa(tgt_loc)
-        src_loc_swa = self.translate_loc_from_full_to_swa(src_loc)
-        self.swa_kv_pool.move_kv_cache(tgt_loc_swa, src_loc_swa)
-
-    def get_cpu_copy(self, indices):
-        # For SWA, we need to copy KV cache from both full and SWA pools
-        # The indices are for the full pool, and we use mapping to get SWA indices
-        full_kv_cpu = self.full_kv_pool.get_cpu_copy(indices)
-
-        # Get SWA indices through the mapping
-        # Note: SWA allocation always creates 1:1 mapping, so no need to filter
-        if self.full_to_swa_index_mapping is not None:
-            swa_indices = self.full_to_swa_index_mapping[indices]
-            swa_kv_cpu = self.swa_kv_pool.get_cpu_copy(swa_indices)
-        else:
-            swa_kv_cpu = None
-
-        return {"full": full_kv_cpu, "swa": swa_kv_cpu}
-
-    def load_cpu_copy(self, kv_cache_cpu, indices):
-        # Load KV cache back from CPU to both full and SWA pools
-        # Note: indices here are NEW indices (newly allocated), different from get_cpu_copy indices
-        full_kv_cpu = kv_cache_cpu["full"]
-        swa_kv_cpu = kv_cache_cpu["swa"]
-
-        # Load full KV cache to the new indices
-        self.full_kv_pool.load_cpu_copy(full_kv_cpu, indices)
-
-        # Load SWA KV cache if it exists
-        if swa_kv_cpu is not None and self.full_to_swa_index_mapping is not None:
-            swa_indices = self.full_to_swa_index_mapping[indices]
-            self.swa_kv_pool.load_cpu_copy(swa_kv_cpu, swa_indices)
-
-
 class MLATokenToKVPool(KVCache):
     def __init__(
         self,

python/sglang/srt/mem_cache/swa_memory_pool.py

@@ -0,0 +1,459 @@
+import logging
+import weakref
+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.memory_pool import KVCache, MHATokenToKVPool
+from sglang.srt.mem_cache.utils import maybe_init_custom_mem_pool
+
+logger = logging.getLogger(__name__)
+GB = 1024 * 1024 * 1024
+
+
+class SWAKVPool(KVCache):
+    """KV cache with separate pools for full and SWA attention layers."""
+
+    def __init__(
+        self,
+        size: int,
+        size_swa: int,
+        page_size: int,
+        dtype: torch.dtype,
+        head_num: int,
+        head_dim: int,
+        swa_attention_layer_ids: List[int],
+        full_attention_layer_ids: List[int],
+        enable_kvcache_transpose: bool,
+        device: str,
+        token_to_kv_pool_class: KVCache = MHATokenToKVPool,
+        **kwargs,
+    ):
+        self.size = size
+        self.size_swa = size_swa
+        self.dtype = dtype
+        self.head_num = head_num
+        self.head_dim = head_dim
+        self.device = device
+        self.swa_layer_nums = len(swa_attention_layer_ids)
+        self.full_layer_nums = len(full_attention_layer_ids)
+        self.start_layer = 0
+        self.page_size = page_size
+        self.swa_loc = None
+
+        kwargs["page_size"] = page_size
+        kwargs["enable_memory_saver"] = False
+        kwargs["head_num"] = head_num
+        kwargs["head_dim"] = head_dim
+        kwargs["device"] = device
+        # TODO MHATransposedTokenToKVPool if enable_kvcache_transpose is True
+        assert not enable_kvcache_transpose
+
+        # for disagg with nvlink
+        self.enable_custom_mem_pool, self.custom_mem_pool, _ = (
+            maybe_init_custom_mem_pool(device=self.device)
+        )
+
+        self.swa_kv_pool = token_to_kv_pool_class(
+            size=size_swa,
+            dtype=dtype,
+            layer_num=self.swa_layer_nums,
+            **kwargs,
+        )
+        kwargs.pop("swa_head_num", None)
+        kwargs.pop("swa_head_dim", None)
+        kwargs.pop("swa_v_head_dim", None)
+        self.full_kv_pool = token_to_kv_pool_class(
+            size=size,
+            dtype=dtype,
+            layer_num=self.full_layer_nums,
+            **kwargs,
+        )
+        # {layer_id: (index, is_swa_layer)}
+        self.layers_mapping: Dict[int, Tuple[int, bool]] = {}
+        for full_attn_layer_id, global_layer_id in enumerate(full_attention_layer_ids):
+            self.layers_mapping[global_layer_id] = (full_attn_layer_id, False)
+        for swa_layer_id, global_layer_id in enumerate(swa_attention_layer_ids):
+            self.layers_mapping[global_layer_id] = (swa_layer_id, True)
+        self.full_to_swa_index_mapping: Optional[torch.Tensor] = None
+
+        k_size, v_size = self.get_kv_size_bytes()
+        self.mem_usage = (k_size + v_size) / GB
+        logger.info(
+            f"SWAKVPool mem usage: {self.mem_usage:.2f} GB, swa size: {self.size_swa}, full size: {self.size}"
+        )
+
+    def register_mapping(self, full_to_swa_index_mapping: torch.Tensor):
+        self.full_to_swa_index_mapping = full_to_swa_index_mapping
+
+    def get_kv_size_bytes(self):
+        k_size, v_size = self.full_kv_pool.get_kv_size_bytes()
+        k_size_swa, v_size_swa = self.swa_kv_pool.get_kv_size_bytes()
+        return k_size + k_size_swa, v_size + v_size_swa
+
+    def get_contiguous_buf_infos(self):
+        full_kv_data_ptrs, full_kv_data_lens, full_kv_item_lens = (
+            self.full_kv_pool.get_contiguous_buf_infos()
+        )
+        return (
+            full_kv_data_ptrs,
+            full_kv_data_lens,
+            full_kv_item_lens,
+        )
+
+    def get_state_buf_infos(self):
+        swa_kv_data_ptrs, swa_kv_data_lens, swa_kv_item_lens = (
+            self.swa_kv_pool.get_contiguous_buf_infos()
+        )
+
+        return swa_kv_data_ptrs, swa_kv_data_lens, swa_kv_item_lens
+
+    def get_key_buffer(self, layer_id: int):
+        layer_id_pool, is_swa_layer = self.layers_mapping[layer_id]
+        if is_swa_layer:
+            return self.swa_kv_pool.get_key_buffer(layer_id_pool)
+        else:
+            return self.full_kv_pool.get_key_buffer(layer_id_pool)
+
+    def get_value_buffer(self, layer_id: int):
+        layer_id_pool, is_swa_layer = self.layers_mapping[layer_id]
+        if is_swa_layer:
+            return self.swa_kv_pool.get_value_buffer(layer_id_pool)
+        else:
+            return self.full_kv_pool.get_value_buffer(layer_id_pool)
+
+    def get_kv_buffer(self, layer_id: int):
+        layer_id_pool, is_swa_layer = self.layers_mapping[layer_id]
+        if is_swa_layer:
+            return self.swa_kv_pool.get_kv_buffer(layer_id_pool)
+        else:
+            return self.full_kv_pool.get_kv_buffer(layer_id_pool)
+
+    def set_swa_loc(self, loc: torch.Tensor):
+        self.swa_loc = loc
+
+    def translate_loc_from_full_to_swa(self, kv_indices: torch.Tensor):
+        assert self.full_to_swa_index_mapping is not None
+
+        # Note: kv_indices could have -1 values (from alloc_extend), which will be mapped to -1
+        # since the last item of full_to_swa_index_mapping is -1.
+        return self.full_to_swa_index_mapping[kv_indices].to(torch.int32)
+
+    def set_kv_buffer(
+        self,
+        layer: RadixAttention,
+        loc: torch.Tensor,
+        cache_k: torch.Tensor,
+        cache_v: torch.Tensor,
+        k_scale: float = 1.0,
+        v_scale: float = 1.0,
+    ):
+
+        layer_id = layer.layer_id
+        layer_id_pool, is_swa_layer = self.layers_mapping[layer_id]
+        if is_swa_layer:
+            if self.swa_loc is not None:
+                loc = self.swa_loc
+            else:
+                if self.full_to_swa_index_mapping is not None:
+                    loc = self.translate_loc_from_full_to_swa(loc)
+
+            self.swa_kv_pool.set_kv_buffer(
+                None,
+                loc,
+                cache_k,
+                cache_v,
+                k_scale,
+                v_scale,
+                layer_id_override=layer_id_pool,
+            )
+        else:
+            self.full_kv_pool.set_kv_buffer(
+                None,
+                loc,
+                cache_k,
+                cache_v,
+                k_scale,
+                v_scale,
+                layer_id_override=layer_id_pool,
+            )
+
+    def move_kv_cache(self, tgt_loc: torch.Tensor, src_loc: torch.Tensor):
+        self.full_kv_pool.move_kv_cache(tgt_loc, src_loc)
+        tgt_loc_swa = self.translate_loc_from_full_to_swa(tgt_loc)
+        src_loc_swa = self.translate_loc_from_full_to_swa(src_loc)
+        self.swa_kv_pool.move_kv_cache(tgt_loc_swa, src_loc_swa)
+
+    def get_cpu_copy(self, indices):
+        # For SWA, we need to copy KV cache from both full and SWA pools
+        # The indices are for the full pool, and we use mapping to get SWA indices
+        full_kv_cpu = self.full_kv_pool.get_cpu_copy(indices)
+
+        # Get SWA indices through the mapping
+        # Note: SWA allocation always creates 1:1 mapping, so no need to filter
+        if self.full_to_swa_index_mapping is not None:
+            swa_indices = self.full_to_swa_index_mapping[indices]
+            swa_kv_cpu = self.swa_kv_pool.get_cpu_copy(swa_indices)
+        else:
+            swa_kv_cpu = None
+
+        return {"full": full_kv_cpu, "swa": swa_kv_cpu}
+
+    def load_cpu_copy(self, kv_cache_cpu, indices):
+        # Load KV cache back from CPU to both full and SWA pools
+        # Note: indices here are NEW indices (newly allocated), different from get_cpu_copy indices
+        full_kv_cpu = kv_cache_cpu["full"]
+        swa_kv_cpu = kv_cache_cpu["swa"]
+
+        # Load full KV cache to the new indices
+        self.full_kv_pool.load_cpu_copy(full_kv_cpu, indices)
+
+        # Load SWA KV cache if it exists
+        if swa_kv_cpu is not None and self.full_to_swa_index_mapping is not None:
+            swa_indices = self.full_to_swa_index_mapping[indices]
+            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: SWAKVPool,
+        need_sort: bool,
+    ):
+        assert isinstance(kvcache, SWAKVPool)
+        self._size_full = size
+        self._size_swa = size_swa
+        self.dtype = dtype
+        self.device = device
+        self.page_size = page_size
+
+        if page_size == 1:
+            self.full_attn_allocator = TokenToKVPoolAllocator(
+                size,
+                dtype,
+                device,
+                kvcache.full_kv_pool,
+                need_sort,
+            )
+            self.swa_attn_allocator = TokenToKVPoolAllocator(
+                size_swa,
+                dtype,
+                device,
+                kvcache.swa_kv_pool,
+                need_sort,
+            )
+        else:
+            self.full_attn_allocator = PagedTokenToKVPoolAllocator(
+                size,
+                page_size,
+                dtype,
+                device,
+                kvcache.full_kv_pool,
+                need_sort,
+            )
+            self.swa_attn_allocator = PagedTokenToKVPoolAllocator(
+                size_swa,
+                page_size,
+                dtype,
+                device,
+                kvcache.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.clear()
+        self._kvcache = kvcache
+        self._kvcache.register_mapping(weakref.proxy(self.full_to_swa_index_mapping))
+
+    def available_size(self):
+        # Note: use full_available_size() and swa_available_size() instead.
+        raise NotImplementedError()
+
+    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
+
+        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_tokens = extend_num_tokens + len(seq_lens) * self.page_size
+        if num_tokens > self.full_attn_allocator.available_size():
+            return None
+        if num_tokens > self.swa_attn_allocator.available_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,
+        )
+        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,
+        )
+        assert alloc_full_indices is not None
+        assert alloc_swa_indices is not None
+
+        self.full_to_swa_index_mapping[alloc_full_indices] = alloc_swa_indices
+
+        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
+
+        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 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):
+        return self._kvcache.get_cpu_copy(indices)
+
+    def load_cpu_copy(self, kv_cache_cpu, indices):
+        return self._kvcache.load_cpu_copy(kv_cache_cpu, indices)

python/sglang/srt/mem_cache/swa_radix_cache.py

@@ -28,7 +28,6 @@ from typing import TYPE_CHECKING, List, Optional, Tuple
 import torch
 from numpy import float64
 
-from sglang.srt.mem_cache.allocator import SWATokenToKVPoolAllocator
 from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache, MatchResult
 from sglang.srt.mem_cache.cache_init_params import CacheInitParams
 from sglang.srt.mem_cache.radix_cache import (
@@ -37,6 +36,7 @@ from sglang.srt.mem_cache.radix_cache import (
     _key_match_paged,
     get_child_key,
 )
+from sglang.srt.mem_cache.swa_memory_pool import SWATokenToKVPoolAllocator
 from sglang.srt.mem_cache.utils import convert_to_bigram_key
 
 if TYPE_CHECKING:

python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py

@@ -10,7 +10,6 @@ from sglang.srt.distributed.parallel_state import get_world_group
 from sglang.srt.layers.dp_attention import get_attention_tp_size
 from sglang.srt.mem_cache.allocator import (
     PagedTokenToKVPoolAllocator,
-    SWATokenToKVPoolAllocator,
     TokenToKVPoolAllocator,
 )
 from sglang.srt.mem_cache.memory_pool import (
@@ -23,8 +22,8 @@ from sglang.srt.mem_cache.memory_pool import (
     MLATokenToKVPoolFP4,
     NSATokenToKVPool,
     ReqToTokenPool,
-    SWAKVPool,
 )
+from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool, SWATokenToKVPoolAllocator
 from sglang.srt.utils.common import (
     get_available_gpu_memory,
     is_float4_e2m1fn_x2,

test/srt/test_swa_unittest.py

@@ -2,10 +2,10 @@ import unittest
 
 import torch
 
-from sglang.srt.mem_cache.allocator import SWAKVPool, 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.radix_cache import RadixKey
+from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool, SWATokenToKVPoolAllocator
 from sglang.srt.mem_cache.swa_radix_cache import SWARadixCache