From 76bc07a33588075d73a9f53e7156f9d3c59d8478 Mon Sep 17 00:00:00 2001 From: Ke Bao Date: Sun, 4 Jan 2026 22:39:30 +0800 Subject: [PATCH] Move swa memory pool to a seperate file (#16347) --- python/sglang/srt/disaggregation/decode.py | 2 +- python/sglang/srt/disaggregation/prefill.py | 7 +- .../attention/flashattention_backend.py | 2 +- .../layers/attention/flashinfer_backend.py | 2 +- python/sglang/srt/managers/schedule_batch.py | 6 +- python/sglang/srt/managers/schedule_policy.py | 2 +- python/sglang/srt/mem_cache/allocator.py | 241 --------- python/sglang/srt/mem_cache/chunk_cache.py | 2 +- python/sglang/srt/mem_cache/common.py | 2 +- python/sglang/srt/mem_cache/memory_pool.py | 205 +------- .../sglang/srt/mem_cache/swa_memory_pool.py | 459 ++++++++++++++++++ .../sglang/srt/mem_cache/swa_radix_cache.py | 2 +- .../model_runner_kv_cache_mixin.py | 3 +- test/srt/test_swa_unittest.py | 2 +- 14 files changed, 473 insertions(+), 464 deletions(-) create mode 100644 python/sglang/srt/mem_cache/swa_memory_pool.py diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index 199885244..ad87b2354 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/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 diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index ac11013f8..a6eed743a 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/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 ( - HybridLinearKVPool, - NSATokenToKVPool, - SWAKVPool, -) +from sglang.srt.mem_cache.memory_pool import HybridLinearKVPool, NSATokenToKVPool +from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool from sglang.srt.tracing.trace import trace_event_batch, trace_slice, trace_slice_end if TYPE_CHECKING: diff --git a/python/sglang/srt/layers/attention/flashattention_backend.py b/python/sglang/srt/layers/attention/flashattention_backend.py index 53e3bfe9d..6b3e22e39 100644 --- a/python/sglang/srt/layers/attention/flashattention_backend.py +++ b/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 diff --git a/python/sglang/srt/layers/attention/flashinfer_backend.py b/python/sglang/srt/layers/attention/flashinfer_backend.py index 29a1ec83a..3a72264a8 100644 --- a/python/sglang/srt/layers/attention/flashinfer_backend.py +++ b/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 ( diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 468d8fb8a..4d7deb474 100644 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/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 ( - BaseTokenToKVPoolAllocator, - SWATokenToKVPoolAllocator, -) +from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator 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, diff --git a/python/sglang/srt/managers/schedule_policy.py b/python/sglang/srt/managers/schedule_policy.py index e03fd8a27..bb0197956 100644 --- a/python/sglang/srt/managers/schedule_policy.py +++ b/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 diff --git a/python/sglang/srt/mem_cache/allocator.py b/python/sglang/srt/mem_cache/allocator.py index dfc34ffe9..eaf29628b 100644 --- a/python/sglang/srt/mem_cache/allocator.py +++ b/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, diff --git a/python/sglang/srt/mem_cache/chunk_cache.py b/python/sglang/srt/mem_cache/chunk_cache.py index 09d2c951c..3eafb7592 100644 --- a/python/sglang/srt/mem_cache/chunk_cache.py +++ b/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 diff --git a/python/sglang/srt/mem_cache/common.py b/python/sglang/srt/mem_cache/common.py index b06e3a401..38a4eee8e 100644 --- a/python/sglang/srt/mem_cache/common.py +++ b/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 diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index 820a21db5..9fb9b939b 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/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, diff --git a/python/sglang/srt/mem_cache/swa_memory_pool.py b/python/sglang/srt/mem_cache/swa_memory_pool.py new file mode 100644 index 000000000..e323b0e31 --- /dev/null +++ b/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) diff --git a/python/sglang/srt/mem_cache/swa_radix_cache.py b/python/sglang/srt/mem_cache/swa_radix_cache.py index 901b67a99..625a8a166 100644 --- a/python/sglang/srt/mem_cache/swa_radix_cache.py +++ b/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: 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 d9645ebc8..7662dee1d 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 @@ -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, diff --git a/test/srt/test_swa_unittest.py b/test/srt/test_swa_unittest.py index ef8a69d68..95ebdb005 100644 --- a/test/srt/test_swa_unittest.py +++ b/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