import logging from typing import Dict, List, Optional, Tuple import torch from sglang.srt.layers.radix_attention import RadixAttention from sglang.srt.mem_cache.base_swa_memory_pool import BaseSWAKVPool from sglang.srt.mem_cache.memory_pool import ( KVCache, MHATokenToKVPool, unwrap_write_loc, ) from sglang.srt.mem_cache.utils import maybe_init_custom_mem_pool logger = logging.getLogger(__name__) GB = 1024 * 1024 * 1024 class SWAKVPool(BaseSWAKVPool): """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], 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.layer_num = self.full_layer_nums + self.swa_layer_nums self.start_layer = 0 self.page_size = page_size self.layer_transfer_counter = None kwargs["page_size"] = page_size kwargs["enable_memory_saver"] = False kwargs["head_num"] = head_num kwargs["head_dim"] = head_dim kwargs["device"] = device # 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}" ) @property def post_capture_active(self) -> bool: """True iff the sub-pools took the post-capture VA-backed path (both share the flag).""" return self.full_kv_pool.post_capture_active @property def post_capture_backed_bytes(self) -> int: """Physically-backed KV bytes across both sub-pools (post-capture only).""" return ( self.full_kv_pool.post_capture_backed_bytes + self.swa_kv_pool.post_capture_backed_bytes ) def finalize_backing(self, config) -> None: """Back both sub-pools to their post-capture final sizes and record them.""" self.full_kv_pool._finalize_backing_tokens(config.full_max_total_num_tokens) self.swa_kv_pool._finalize_backing_tokens(config.swa_max_total_num_tokens) self.size = int(config.full_max_total_num_tokens) self.size_swa = int(config.swa_max_total_num_tokens) def register_mapping(self, full_to_swa_index_mapping: torch.Tensor): self.full_to_swa_index_mapping = full_to_swa_index_mapping def register_layer_transfer_counter(self, layer_transfer_counter): # Wait happens at this wrapper. Inner pools must not wait again. self.layer_transfer_counter = layer_transfer_counter self.full_kv_pool.register_layer_transfer_counter(None) self.swa_kv_pool.register_layer_transfer_counter(None) def _wait_for_layer(self, layer_id: int): if self.layer_transfer_counter is not None: self.layer_transfer_counter.wait_until(layer_id - self.start_layer) 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): self._wait_for_layer(layer_id) 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): self._wait_for_layer(layer_id) 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): self._wait_for_layer(layer_id) 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 translate_loc_from_full_to_swa(self, kv_indices: torch.Tensor) -> torch.Tensor: assert self.full_to_swa_index_mapping is not None # -1 in kv_indices maps to -1 via the sentinel appended to the mapping. return self.full_to_swa_index_mapping[kv_indices] def set_kv_buffer( self, layer: RadixAttention, loc_info, cache_k: torch.Tensor, cache_v: torch.Tensor, k_scale: float = 1.0, v_scale: float = 1.0, ): # loc_info bundles the full loc and the pre-translated SWA loc. loc, swa_loc, _ = unwrap_write_loc(loc_info) layer_id = layer.layer_id layer_id_pool, is_swa_layer = self.layers_mapping[layer_id] if is_swa_layer: # swa_loc is the full->SWA translation, computed once per forward by # the attention backend; set_kv_buffer never translates internally. assert swa_loc is not None self.swa_kv_pool.set_kv_buffer( None, swa_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 _filter_swa_cpu_copy(self, swa_kv_cpu, row_mask: torch.Tensor): if swa_kv_cpu is None: return None if row_mask is None or bool(torch.all(row_mask).item()): return swa_kv_cpu chunk_size = getattr( self.swa_kv_pool, "cpu_offloading_chunk_size", len(row_mask) ) filtered = [] for layer_chunks in swa_kv_cpu: if len(layer_chunks) == 0: filtered.append([]) continue k_cpu = torch.cat([chunk[0] for chunk in layer_chunks], dim=0) v_cpu = torch.cat([chunk[1] for chunk in layer_chunks], dim=0) k_cpu = k_cpu[row_mask] v_cpu = v_cpu[row_mask] filtered_layer = [] for i in range(0, len(k_cpu), chunk_size): filtered_layer.append( [k_cpu[i : i + chunk_size], v_cpu[i : i + chunk_size]] ) filtered.append(filtered_layer) return filtered def get_cpu_copy(self, indices, mamba_indices=None): # 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) swa_mask = None if self.full_to_swa_index_mapping is not None: swa_indices = self.full_to_swa_index_mapping[indices] # Slot 0 is reserved as a dummy slot. Tail-only SWA allocations leave # the out-of-window full KV indices unmapped, so only copy mapped SWA # tokens and keep their positions for load_cpu_copy(). swa_mask = swa_indices > 0 if torch.any(swa_mask): swa_kv_cpu = self.swa_kv_pool.get_cpu_copy(swa_indices[swa_mask]) swa_mask = swa_mask.cpu() else: swa_kv_cpu = None else: swa_kv_cpu = None return {"full": full_kv_cpu, "swa": swa_kv_cpu, "swa_mask": swa_mask} def load_cpu_copy(self, kv_cache_cpu, indices, mamba_indices=None): # 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] new_swa_mask = swa_indices > 0 old_swa_mask = kv_cache_cpu.get("swa_mask") if old_swa_mask is not None: old_swa_mask = old_swa_mask.to(indices.device) row_mask = new_swa_mask[old_swa_mask].cpu() swa_indices = swa_indices[old_swa_mask][row_mask.to(indices.device)] else: row_mask = new_swa_mask.cpu() swa_indices = swa_indices[new_swa_mask] if swa_indices.numel() == 0: return 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)