from __future__ import annotations import dataclasses from contextlib import nullcontext from math import gcd import torch from sglang.srt.constants import GPU_MEMORY_TYPE_KV_CACHE from sglang.srt.mem_cache.utils import maybe_init_custom_mem_pool from sglang.srt.utils import is_hip, is_npu from sglang.srt.utils.torch_memory_saver_adapter import TorchMemorySaverAdapter _is_hip = is_hip() _is_npu = is_npu() def _lcm(a: int, b: int) -> int: return a // gcd(a, b) * b @dataclasses.dataclass class KVAndScore: kv_score: torch.Tensor @property def kv(self) -> torch.Tensor: return self.kv_score[..., : self._item_size] @property def score(self) -> torch.Tensor: return self.kv_score[..., self._item_size :] @property def shape(self): return self.kv_score.shape def __post_init__(self): self._item_size = self.kv_score.shape[-1] // 2 @staticmethod def from_kv_score(*, kv: torch.Tensor, score: torch.Tensor) -> KVAndScore: assert kv.shape == score.shape return KVAndScore(torch.cat([kv, score], dim=-1)) def new_empty(self, new_shape) -> KVAndScore: assert new_shape[-1] == self._item_size new_shape = list(new_shape) new_shape[-1] = 2 * self._item_size return KVAndScore(self.kv_score.new_empty(new_shape, requires_grad=False)) def __getitem__(self, index) -> KVAndScore: return KVAndScore(self.kv_score[index]) def __setitem__(self, index, value: KVAndScore): self.kv_score[index] = value.kv_score def clear(self): self.kv.zero_() self.score.fill_(float("-inf")) def view(self, *args): args = list(args) if isinstance(args[-1], int) and args[-1] != -1: args[-1] = 2 * self._item_size return KVAndScore(self.kv_score.view(*args)) def clone(self) -> KVAndScore: return KVAndScore(self.kv_score.clone()) @staticmethod def cat(tensors: list[KVAndScore], dim: int) -> KVAndScore: assert dim != -1, "Concatenation along last dim is not supported." assert len(tensors) > 0, "At least one tensor is required for concatenation." item_size = tensors[0]._item_size for v in tensors: assert ( v._item_size == item_size ), "All tensors must have the same item size." return KVAndScore(torch.cat([v.kv_score for v in tensors], dim=dim)) class CompressStatePool: def __init__( self, size: int, ring_size: int, overlap: bool, head_dim: int, dtype: torch.dtype, device: str, enable_memory_saver: bool, ratio: int, online: bool = False, swa_page_size: int = 0, online_mtp_max_draft_tokens: int = 0, ): self.ratio = ratio self.ring_size = ring_size self.swa_page_size = swa_page_size self.enable_memory_saver = enable_memory_saver self.online_mtp_state_slot_offset = 0 self.online_mtp_max_draft_tokens = 0 if online: assert ring_size == 1, "online compress requires ring_size=1" self._logical_size = size + self.ring_size + 1 if online_mtp_max_draft_tokens > 0: # Bank 0 is the committed state. Banks 1..N cache per-draft # prefix states for lazy commit after target verify. self.online_mtp_max_draft_tokens = online_mtp_max_draft_tokens self.online_mtp_state_slot_offset = self._logical_size self._size = self._logical_size * (1 + self.online_mtp_max_draft_tokens) last_dim = 3 * head_dim else: self._size = size + self.ring_size + 1 # Pad to lcm(ratio, page_size) so the flat buffer reshapes cleanly into # [block_num, page_size, last_dim] for the fused compressor op; page_size=1 falls back to ratio-only padding. pad_to = ( _lcm(ratio, swa_page_size) if (swa_page_size > 1 and _is_npu) else ratio ) self._size = (self._size + pad_to - 1) // pad_to * pad_to self._logical_size = self._size last_dim = 2 * (1 + overlap) * head_dim self.last_dim = last_dim self._alloc_kv_score_buffer( dtype=dtype, device=device, enable_memory_saver=enable_memory_saver ) if not online: if _is_hip and ratio == 128: # Request-scoped C128 state is addressed by req_pool_idx (or a # per-request ring). The pool is allocated with torch.empty(), # so a cold server can otherwise read uninitialized partial # states before a request slot has been written for the first # time. Initialize all C128 rows to the empty-state sentinel; # C4 keeps the historical last-row sentinel behavior. self.kv_score_buffer.clear() else: self.kv_score_buffer[-1].clear() def _alloc_kv_score_buffer( self, *, dtype: torch.dtype, device: str, enable_memory_saver: bool ) -> None: """Allocate the flat ``(self._size, self.last_dim)`` kv+score buffer under the memory-saver / custom-mem-pool context and wrap it in :class:`KVAndScore`. Sets ``self.memory_saver_adapter``, ``self.custom_mem_pool`` and ``self.kv_score_buffer``. Subclasses (e.g. :class:`NPUCompressStatePool`) that compute a different ``self._size`` reuse this instead of duplicating the allocation boilerplate. Requires ``self._size`` and ``self.last_dim`` to be set already. """ if _is_hip: self.kv_score_buffer = KVAndScore( torch.empty((self._size, self.last_dim), dtype=dtype, device=device) ) else: self.memory_saver_adapter = TorchMemorySaverAdapter.create( enable=enable_memory_saver ) self.enable_custom_mem_pool, self.custom_mem_pool, _ = ( maybe_init_custom_mem_pool(device=device) ) with self.memory_saver_adapter.region(GPU_MEMORY_TYPE_KV_CACHE): with ( torch.cuda.use_mem_pool(self.custom_mem_pool) if self.custom_mem_pool else nullcontext() ): self.kv_score_buffer = KVAndScore( torch.empty( (self._size, self.last_dim), dtype=dtype, device=device, ) ) @property def state_cache_3d(self) -> torch.Tensor: """``[block_num, page_size, last_dim]`` view of the flat kv+score buffer. ``last_dim = 2*(1+overlap)*head_dim`` — exactly the ``2*coff*D`` layout the fused compressor op wants for its ``state_cache`` argument (kv at ``[:, :, :coff*D]``, score at ``[:, :, coff*D:]``). Only valid for the non-online buffer; the online layout has ``last_dim = 3*head_dim`` which the fused path doesn't use. """ assert not self.online, ( "state_cache_3d is for the fused compressor path; " "online (3*head_dim) buffer is indexer-only." ) assert self.page_size > 1, ( "state_cache_3d requires page_size>1; pool was constructed " "with the default page_size=1 (flat 2D layout)." ) return self.kv_score_buffer.kv_score.view(-1, self.page_size, self.last_dim) def translate_from_swa_loc_to_state_loc( self, swa_loc: torch.Tensor ) -> torch.Tensor: swa_pages = swa_loc // self.swa_page_size state_loc = swa_pages * self.ring_size + (swa_loc % self.ring_size) state_loc = torch.where(swa_loc < 0, -1, state_loc) return state_loc def translate_from_req_position_to_state_loc( self, req_pool_indices: torch.Tensor, positions: torch.Tensor ) -> torch.Tensor: state_loc = req_pool_indices * self.ring_size + positions % self.ring_size state_loc = torch.where(positions < 0, -1, state_loc) return state_loc def get_state_by_state_loc(self, state_loc: torch.Tensor) -> KVAndScore: return self.kv_score_buffer[state_loc] def set_state_by_state_loc(self, state_loc: torch.Tensor, value: KVAndScore): self.kv_score_buffer[state_loc] = value self.kv_score_buffer[-1].clear()