From 66076f2409030c19626a68834fc085bb208bf344 Mon Sep 17 00:00:00 2001 From: shuwenn <47200617+alphabetc1@users.noreply.github.com> Date: Thu, 11 Jun 2026 19:50:24 +0800 Subject: [PATCH] [mem_cache][3/N] refactor: move HiSparse allocators to allocator/hisparse.py (#26678) Co-authored-by: Claude Opus 4.8 (1M context) --- .../srt/managers/hisparse_coordinator.py | 6 +- python/sglang/srt/managers/schedule_policy.py | 6 +- .../srt/mem_cache/allocator/hisparse.py | 565 ++++++++++++++++++ python/sglang/srt/mem_cache/chunk_cache.py | 6 +- .../srt/mem_cache/hisparse_memory_pool.py | 563 ----------------- .../model_runner_kv_cache_mixin.py | 10 +- .../unit/managers/test_hisparse_unit.py | 4 +- 7 files changed, 582 insertions(+), 578 deletions(-) create mode 100644 python/sglang/srt/mem_cache/allocator/hisparse.py diff --git a/python/sglang/srt/managers/hisparse_coordinator.py b/python/sglang/srt/managers/hisparse_coordinator.py index 810b2f368..a67c1aa84 100644 --- a/python/sglang/srt/managers/hisparse_coordinator.py +++ b/python/sglang/srt/managers/hisparse_coordinator.py @@ -6,11 +6,13 @@ from typing import List, NamedTuple, Union import torch from sglang.srt.managers.schedule_batch import Req -from sglang.srt.mem_cache.hisparse_memory_pool import ( +from sglang.srt.mem_cache.allocator.hisparse import ( DeepSeekV4HiSparseTokenToKVPoolAllocator, - HiSparseDSATokenToKVPool, HiSparseTokenToKVPoolAllocator, ) +from sglang.srt.mem_cache.hisparse_memory_pool import ( + HiSparseDSATokenToKVPool, +) from sglang.srt.mem_cache.memory_pool_host import ( DeepSeekV4PagedHostPool, MLATokenToKVPoolHost, diff --git a/python/sglang/srt/managers/schedule_policy.py b/python/sglang/srt/managers/schedule_policy.py index 3f3a31e57..1b34f0f54 100644 --- a/python/sglang/srt/managers/schedule_policy.py +++ b/python/sglang/srt/managers/schedule_policy.py @@ -38,6 +38,9 @@ from sglang.srt.dllm.config import DllmConfig from sglang.srt.layers.attention.dsa.utils import is_dsa_prefill_cp_in_seq_split from sglang.srt.layers.utils.cp_utils import is_prefill_context_parallel_enabled from sglang.srt.managers.schedule_batch import Req, ScheduleBatch +from sglang.srt.mem_cache.allocator.hisparse import ( + DeepSeekV4HiSparseTokenToKVPoolAllocator, +) from sglang.srt.mem_cache.allocator.swa import SWATokenToKVPoolAllocator from sglang.srt.mem_cache.base_prefix_cache import ( BasePrefixCache, @@ -46,9 +49,6 @@ from sglang.srt.mem_cache.base_prefix_cache import ( MatchPrefixParams, zero_match_result, ) -from sglang.srt.mem_cache.hisparse_memory_pool import ( - DeepSeekV4HiSparseTokenToKVPoolAllocator, -) from sglang.srt.mem_cache.radix_cache import RadixCache, RadixKey, TreeNode from sglang.srt.server_args import ServerArgs, get_global_server_args diff --git a/python/sglang/srt/mem_cache/allocator/hisparse.py b/python/sglang/srt/mem_cache/allocator/hisparse.py new file mode 100644 index 000000000..5e23ce596 --- /dev/null +++ b/python/sglang/srt/mem_cache/allocator/hisparse.py @@ -0,0 +1,565 @@ +import weakref + +import torch + +from sglang.srt.mem_cache.allocator.base import BaseTokenToKVPoolAllocator +from sglang.srt.mem_cache.allocator.paged import PagedTokenToKVPoolAllocator +from sglang.srt.mem_cache.deepseek_v4_memory_pool import ( + DeepSeekV4TokenToKVPool, + HiSparseC4DevicePool, +) +from sglang.srt.mem_cache.hisparse_memory_pool import HiSparseDSATokenToKVPool +from sglang.srt.utils.common import get_num_new_pages + + +class HiSparseTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator): + def __init__( + self, + size: int, + page_size: int, + dtype: torch.dtype, + device: torch.device, + kvcache: HiSparseDSATokenToKVPool, + need_sort: bool, + host_to_device_ratio: int = 2, + ): + self._kvcache = kvcache + self._size_full = size * host_to_device_ratio + self._size_hisparse = size + self.compress_ratio = 1 + self.dtype = dtype + self.device = device + self.page_size = page_size + self.need_sort = need_sort + + self.logical_attn_allocator = PagedTokenToKVPoolAllocator( + self._size_full, + self.page_size, + self.dtype, + self.device, + kvcache, + need_sort, + ) + self.hisparse_attn_allocator = PagedTokenToKVPoolAllocator( + self._size_hisparse, + self.page_size, + self.dtype, + self.device, + kvcache, + need_sort, + ) + self.full_to_hisparse_device_index_mapping = torch.cat( + [ + torch.zeros( + self._size_full + self.page_size, + dtype=torch.int64, + device=self.device, + ), + torch.tensor([-1], dtype=torch.int64, device=self.device), + ] + ) + + self.free_pages = None + self.release_pages = None + self.is_not_in_free_group = True + self.free_group = [] + self.clear() + self._kvcache.register_mapping( + weakref.proxy(self.full_to_hisparse_device_index_mapping) + ) + + @property + def size_full(self) -> int: + return self._size_full + + @property + def size(self) -> int: + return self._size_full + + def available_size(self) -> int: + return min( + self.logical_attn_allocator.available_size(), + self.hisparse_attn_allocator.available_size(), + ) + + def get_kvcache(self): + return self._kvcache + + def alloc(self, need_size: int): + raise NotImplementedError( + "HiSparse allocator does not support direct token allocation; " + "use alloc_extend or alloc_decode instead." + ) + + def alloc_logical_only( + self, + prefix_lens: torch.Tensor, + prefix_lens_cpu: torch.Tensor, + seq_lens: torch.Tensor, + seq_lens_cpu: torch.Tensor, + last_loc: torch.Tensor, + extend_num_tokens: int, + ): + """Allocate only logical indices without hisparse device indices. + + Used in the direct-to-host transfer path where KV data is written + directly to host memory by the prefill node, skipping GPU staging. + """ + return self.logical_attn_allocator.alloc_extend( + prefix_lens, + prefix_lens_cpu, + seq_lens, + seq_lens_cpu, + last_loc, + extend_num_tokens, + ) + + def alloc_device_buffer(self, allocated_indices, need_size: int): + assert need_size % self.page_size == 0 + # clear original reference and isolate the buffer from outside addressing, allocate new buffer if needed + hisparse_indices = self.full_to_hisparse_device_index_mapping[allocated_indices] + self.full_to_hisparse_device_index_mapping[allocated_indices] = 0 + # Filter valid (non-zero) hisparse indices. + # In the direct-to-host path, mapping is all zeros since no hisparse + # device indices were pre-allocated. + hisparse_indices = hisparse_indices[hisparse_indices > 0] + if len(hisparse_indices) >= need_size: + buffer_indices = hisparse_indices[:need_size] + self.free_hisparse_indices(hisparse_indices[need_size:]) + else: + # page alignment, claiming the residual space for an incomplete page + page_residual_length = len(hisparse_indices) % self.page_size + if page_residual_length != 0: + hisparse_indices = torch.cat( + [ + hisparse_indices, + torch.arange( + hisparse_indices[-1] + 1, + hisparse_indices[-1] + + self.page_size + - page_residual_length + + 1, + device=self.device, + ), + ] + ) + extra_indices = self.hisparse_attn_allocator.alloc( + need_size - len(hisparse_indices) + ) + assert ( + extra_indices is not None + ), "Hisparse allocation failed in alloc_device_buffer" + buffer_indices = torch.cat([hisparse_indices, extra_indices]) + return buffer_indices + + def free_hisparse_indices(self, buffer_indices: torch.Tensor): + # disable free group mechanism for device buffer free + self.hisparse_attn_allocator.is_not_in_free_group = True + self.hisparse_attn_allocator.free(buffer_indices[buffer_indices > 0]) + + def get_last_loc_compressed(self, last_locs: torch.Tensor): + return last_locs + + def get_last_loc_hisparse_device(self, last_locs: torch.Tensor): + return self._kvcache._translate_loc_to_hisparse_device(last_locs) + + def alloc_extend( + self, + prefix_lens: torch.Tensor, + prefix_lens_cpu: torch.Tensor, + seq_lens: torch.Tensor, + seq_lens_cpu: torch.Tensor, + last_loc: torch.Tensor, # last_loc for full layers + extend_num_tokens: int, + ): + assert self.page_size > 1 + + num_new_pages = get_num_new_pages( + seq_lens=seq_lens_cpu, page_size=self.page_size, prefix_lens=prefix_lens_cpu + ) + if ( + num_new_pages + > self.logical_attn_allocator.available_size() // self.page_size + ): + return None + if ( + num_new_pages + > self.hisparse_attn_allocator.available_size() // self.page_size + ): + return None + + logical_indices = self.logical_attn_allocator.alloc_extend( + prefix_lens, + prefix_lens_cpu, + seq_lens, + seq_lens_cpu, + last_loc, + extend_num_tokens, + ) + assert logical_indices is not None, "Logical allocation failed in alloc_extend" + + hisparse_last_loc = self.get_last_loc_hisparse_device(last_loc) + hisparse_indices = self.hisparse_attn_allocator.alloc_extend( + prefix_lens, + prefix_lens_cpu, + seq_lens, + seq_lens_cpu, + hisparse_last_loc, + len(logical_indices), + num_new_pages=num_new_pages, + ) + assert ( + hisparse_indices is not None + ), "Hisparse allocation failed in alloc_extend" + self.full_to_hisparse_device_index_mapping[logical_indices] = hisparse_indices + return logical_indices + + def alloc_decode( + self, + seq_lens: torch.Tensor, + seq_lens_cpu: torch.Tensor, + last_loc: torch.Tensor, # last_loc for full layers + ): + return self.logical_attn_allocator.alloc_decode( + seq_lens, seq_lens_cpu, last_loc + ) + + def free_hisparse(self, free_indices: torch.Tensor): + hisparse_indices = self._kvcache._translate_loc_to_hisparse_device(free_indices) + hisparse_indices = hisparse_indices[hisparse_indices > 0] + self.free_hisparse_indices(hisparse_indices) + self.full_to_hisparse_device_index_mapping[free_indices] = 0 + + def clear(self): + self.logical_attn_allocator.clear() + self.hisparse_attn_allocator.clear() + # Note: the last item is -1, we don't clear it, see the comment in __init__ + self.full_to_hisparse_device_index_mapping[:-1].fill_(0) + self.is_not_in_free_group = True + self.free_group = [] + + def free_group_begin(self): + return + + def free_group_end(self): + return + + def free(self, free_index: torch.Tensor): + if free_index.numel() == 0: + return + if self.is_not_in_free_group: + self.logical_attn_allocator.free(free_index) + self.free_hisparse(free_index) + else: + self.free_group.append(free_index) + assert ( + self.logical_attn_allocator.available_size() + <= self.logical_attn_allocator.size + ) + assert ( + self.hisparse_attn_allocator.available_size() + <= self.hisparse_attn_allocator.size + ) + + +class DeepSeekV4HiSparseTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator): + + def __init__( + self, + logical_attn_allocator: BaseTokenToKVPoolAllocator, + ): + assert isinstance(logical_attn_allocator._kvcache, DeepSeekV4TokenToKVPool) + assert isinstance( + logical_attn_allocator._kvcache.c4_kv_pool, HiSparseC4DevicePool + ) + self.compress_ratio = 4 + + self.hisparse_kvcache = logical_attn_allocator._kvcache.c4_kv_pool + self._size_full = logical_attn_allocator.size_full + self._size_hisparse = self.hisparse_kvcache.size + + self.dtype = self.hisparse_kvcache.dtype + self.device = self.hisparse_kvcache.device + # Keep the public page_size as the logical DSV4 full/SWA page size. + # C4 HiSparse allocation/device-buffer code must use the compressed page size. + self.page_size = logical_attn_allocator.page_size + self.hisparse_page_size = self.hisparse_kvcache.page_size + + self.logical_attn_allocator = logical_attn_allocator + self._kvcache = logical_attn_allocator._kvcache + self.hisparse_attn_allocator = PagedTokenToKVPoolAllocator( + self._size_hisparse, + self.hisparse_page_size, + self.dtype, + self.device, + self.hisparse_kvcache, + logical_attn_allocator.need_sort, + ) + + self.full_to_hisparse_device_index_mapping = torch.cat( + [ + torch.zeros( + self._kvcache.c4_logical_size + self.hisparse_page_size, + dtype=torch.int64, + device=self.device, + ), + torch.tensor([-1], dtype=torch.int64, device=self.device), + ] + ) + + self.need_sort = logical_attn_allocator.need_sort + self.free_pages = None + self.release_pages = None + self.is_not_in_free_group = True + self.free_group = [] + self.clear() + + self.hisparse_kvcache.register_mapping( + weakref.proxy(self.full_to_hisparse_device_index_mapping) + ) + + @property + def size_full(self) -> int: + return self._size_full + + @property + def size(self) -> int: + return self.logical_attn_allocator.size + + @property + def size_swa(self) -> int: + return self.logical_attn_allocator.size_swa + + @property + def full_to_swa_index_mapping(self): + return self.logical_attn_allocator.full_to_swa_index_mapping + + def debug_print(self) -> str: + msg = self.logical_attn_allocator.debug_print() + msg += ( + f"#hisparse-available-size: " + f"{self.hisparse_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): + return self.logical_attn_allocator.translate_loc_from_full_to_swa(kv_indices) + + def full_available_size(self): + return min( + self.logical_attn_allocator.full_available_size(), + self.hisparse_attn_allocator.available_size() * self.compress_ratio, + ) + + def swa_available_size(self): + return self.logical_attn_allocator.swa_available_size() + + def free_swa(self, free_indices: torch.Tensor): + self.logical_attn_allocator.free_swa(free_indices) + + def available_size(self) -> int: + return min( + self.logical_attn_allocator.available_size(), + self.hisparse_attn_allocator.available_size() * self.compress_ratio, + ) + + def alloc(self, need_size: int): + raise NotImplementedError( + "DeepSeek V4 HiSparse allocator does not support direct token allocation; " + "use alloc_extend or alloc_decode instead." + ) + + def alloc_logical_only( + self, + prefix_lens: torch.Tensor, + prefix_lens_cpu: torch.Tensor, + seq_lens: torch.Tensor, + seq_lens_cpu: torch.Tensor, + last_loc: torch.Tensor, + extend_num_tokens: int, + ): + """Allocate decode logical indices without allocating C4 hisparse device pages.""" + return self.logical_attn_allocator.alloc_extend( + prefix_lens, + prefix_lens_cpu, + seq_lens, + seq_lens_cpu, + last_loc, + extend_num_tokens, + ) + + def alloc_device_buffer(self, allocated_indices, need_size: int): + assert need_size % self.hisparse_page_size == 0 + hisparse_indices = self.full_to_hisparse_device_index_mapping[allocated_indices] + self.full_to_hisparse_device_index_mapping[allocated_indices] = 0 + hisparse_indices = hisparse_indices[hisparse_indices > 0] + + device_buffer_size = need_size - self.hisparse_page_size + P = len(hisparse_indices) + if P > device_buffer_size + 1: + newest_src = hisparse_indices[P - 1].clone() + old_at_dbs = hisparse_indices[device_buffer_size].clone() + hisparse_indices[device_buffer_size] = newest_src + hisparse_indices[P - 1] = old_at_dbs + + if len(hisparse_indices) >= need_size: + buffer_indices = hisparse_indices[:need_size] + surplus = hisparse_indices[need_size:] + if surplus.numel() > 0: + buffer_pages = torch.unique(buffer_indices // self.hisparse_page_size) + surplus_pages = torch.unique(surplus // self.hisparse_page_size) + pure_surplus = surplus_pages[~torch.isin(surplus_pages, buffer_pages)] + if pure_surplus.numel() > 0: + self.hisparse_attn_allocator.is_not_in_free_group = True + self.hisparse_attn_allocator.free( + pure_surplus * self.hisparse_page_size + ) + else: + page_residual_length = len(hisparse_indices) % self.hisparse_page_size + if page_residual_length != 0: + hisparse_indices = torch.cat( + [ + hisparse_indices, + torch.arange( + hisparse_indices[-1] + 1, + hisparse_indices[-1] + + self.hisparse_page_size + - page_residual_length + + 1, + device=self.device, + ), + ] + ) + extra_indices = self.hisparse_attn_allocator.alloc( + need_size - len(hisparse_indices) + ) + assert ( + extra_indices is not None + ), "Hisparse allocation failed in alloc_device_buffer" + buffer_indices = torch.cat([hisparse_indices, extra_indices]) + return buffer_indices + + def free_hisparse_indices(self, buffer_indices: torch.Tensor): + self.hisparse_attn_allocator.is_not_in_free_group = True + self.hisparse_attn_allocator.free(buffer_indices[buffer_indices > 0]) + + def get_last_loc_compressed(self, last_locs: torch.Tensor): + return (last_locs - 3) // self.compress_ratio + + def get_last_loc_hisparse_device(self, last_locs: torch.Tensor): + return self.hisparse_kvcache._translate_loc_to_hisparse_device( + self.get_last_loc_compressed(last_locs) + ) + + 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, + extend_num_tokens: int, + ): + assert self.page_size > 1 + + num_new_pages_logical = get_num_new_pages( + seq_lens=seq_lens_cpu, page_size=self.page_size, prefix_lens=prefix_lens_cpu + ) + num_new_pages_hisparse = get_num_new_pages( + seq_lens=seq_lens_cpu // self.compress_ratio, + page_size=self.hisparse_page_size, + prefix_lens=prefix_lens_cpu // self.compress_ratio, + ) + if ( + num_new_pages_logical + > self.logical_attn_allocator.available_size() // self.page_size + ): + return None + if ( + num_new_pages_hisparse + > self.hisparse_attn_allocator.available_size() // self.hisparse_page_size + ): + return None + + logical_indices = self.logical_attn_allocator.alloc_extend( + prefix_lens, + prefix_lens_cpu, + seq_lens, + seq_lens_cpu, + last_loc, + extend_num_tokens, + ) + assert logical_indices is not None, "Logical allocation failed in alloc_extend" + + compressed_logical_indices = ( + self.hisparse_kvcache.translate_loc_from_full_to_compressed(logical_indices) + ) + hisparse_last_loc = self.get_last_loc_hisparse_device(last_loc) + hisparse_indices = self.hisparse_attn_allocator.alloc_extend( + prefix_lens // self.compress_ratio, + prefix_lens_cpu // self.compress_ratio, + seq_lens // self.compress_ratio, + seq_lens_cpu // self.compress_ratio, + hisparse_last_loc, + len(compressed_logical_indices), + ) + assert ( + hisparse_indices is not None + ), "Hisparse allocation failed in alloc_extend" + + self.full_to_hisparse_device_index_mapping[compressed_logical_indices] = ( + hisparse_indices.to(torch.int64) + ) + return logical_indices + + def alloc_decode( + self, + seq_lens: torch.Tensor, + seq_lens_cpu: torch.Tensor, + last_loc: torch.Tensor, + ): + return self.logical_attn_allocator.alloc_decode( + seq_lens, seq_lens_cpu, last_loc + ) + + def free_compressed(self, compressed_indices: torch.Tensor): + hisparse_indices = self.hisparse_kvcache.translate_loc_to_hisparse_device( + compressed_indices + ) + hisparse_indices = hisparse_indices[hisparse_indices > 0] + self.free_hisparse_indices(hisparse_indices) + self.full_to_hisparse_device_index_mapping[compressed_indices] = 0 + + def free_hisparse(self, free_indices: torch.Tensor): + compressed_indices = ( + self.hisparse_kvcache.translate_loc_from_full_to_compressed(free_indices) + ) + self.free_compressed(compressed_indices) + + def clear(self): + self.logical_attn_allocator.clear() + self.hisparse_attn_allocator.clear() + + self.full_to_hisparse_device_index_mapping[:-1].fill_(0) + self.is_not_in_free_group = True + self.free_group = [] + + def free(self, free_index: torch.Tensor): + if free_index.numel() == 0: + return + + if self.is_not_in_free_group: + self.logical_attn_allocator.free(free_index) + else: + self.free_group.append(free_index) + assert ( + self.logical_attn_allocator.available_size() + <= self.logical_attn_allocator.size + ) + assert ( + self.hisparse_attn_allocator.available_size() + <= self.hisparse_attn_allocator.size + ) diff --git a/python/sglang/srt/mem_cache/chunk_cache.py b/python/sglang/srt/mem_cache/chunk_cache.py index cfccdc815..d68eb7217 100644 --- a/python/sglang/srt/mem_cache/chunk_cache.py +++ b/python/sglang/srt/mem_cache/chunk_cache.py @@ -7,6 +7,9 @@ from typing import TYPE_CHECKING, Any, Optional import torch +from sglang.srt.mem_cache.allocator.hisparse import ( + DeepSeekV4HiSparseTokenToKVPoolAllocator, +) from sglang.srt.mem_cache.allocator.swa import SWATokenToKVPoolAllocator from sglang.srt.mem_cache.base_prefix_cache import ( BasePrefixCache, @@ -20,9 +23,6 @@ from sglang.srt.mem_cache.base_prefix_cache import ( MatchPrefixParams, MatchResult, ) -from sglang.srt.mem_cache.hisparse_memory_pool import ( - DeepSeekV4HiSparseTokenToKVPoolAllocator, -) if TYPE_CHECKING: from sglang.srt.managers.schedule_batch import Req diff --git a/python/sglang/srt/mem_cache/hisparse_memory_pool.py b/python/sglang/srt/mem_cache/hisparse_memory_pool.py index c0dcd02b8..1bf865ca2 100644 --- a/python/sglang/srt/mem_cache/hisparse_memory_pool.py +++ b/python/sglang/srt/mem_cache/hisparse_memory_pool.py @@ -1,23 +1,13 @@ # mapping on device memory, host memory and memory allocator import logging -import weakref from typing import Optional import torch from sglang.srt.layers.radix_attention import RadixAttention -from sglang.srt.mem_cache.allocator import ( - BaseTokenToKVPoolAllocator, - PagedTokenToKVPoolAllocator, -) -from sglang.srt.mem_cache.deepseek_v4_memory_pool import ( - DeepSeekV4TokenToKVPool, - HiSparseC4DevicePool, -) from sglang.srt.mem_cache.memory_pool import DSATokenToKVPool from sglang.srt.utils import is_cuda, is_hip -from sglang.srt.utils.common import get_num_new_pages logger = logging.getLogger(__name__) @@ -130,556 +120,3 @@ class HiSparseDSATokenToKVPool(DSATokenToKVPool): def load_cpu_copy(self, kv_cache_cpu, indices, mamba_indices=None): raise NotImplementedError("HiSparseDevicePool does not support load_cpu_copy") - - -class HiSparseTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator): - def __init__( - self, - size: int, - page_size: int, - dtype: torch.dtype, - device: torch.device, - kvcache: HiSparseDSATokenToKVPool, - need_sort: bool, - host_to_device_ratio: int = 2, - ): - self._kvcache = kvcache - self._size_full = size * host_to_device_ratio - self._size_hisparse = size - self.compress_ratio = 1 - self.dtype = dtype - self.device = device - self.page_size = page_size - self.need_sort = need_sort - - self.logical_attn_allocator = PagedTokenToKVPoolAllocator( - self._size_full, - self.page_size, - self.dtype, - self.device, - kvcache, - need_sort, - ) - self.hisparse_attn_allocator = PagedTokenToKVPoolAllocator( - self._size_hisparse, - self.page_size, - self.dtype, - self.device, - kvcache, - need_sort, - ) - self.full_to_hisparse_device_index_mapping = torch.cat( - [ - torch.zeros( - self._size_full + self.page_size, - dtype=torch.int64, - device=self.device, - ), - torch.tensor([-1], dtype=torch.int64, device=self.device), - ] - ) - - self.free_pages = None - self.release_pages = None - self.is_not_in_free_group = True - self.free_group = [] - self.clear() - self._kvcache.register_mapping( - weakref.proxy(self.full_to_hisparse_device_index_mapping) - ) - - @property - def size_full(self) -> int: - return self._size_full - - @property - def size(self) -> int: - return self._size_full - - def available_size(self) -> int: - return min( - self.logical_attn_allocator.available_size(), - self.hisparse_attn_allocator.available_size(), - ) - - def get_kvcache(self): - return self._kvcache - - def alloc(self, need_size: int): - raise NotImplementedError( - "HiSparse allocator does not support direct token allocation; " - "use alloc_extend or alloc_decode instead." - ) - - def alloc_logical_only( - self, - prefix_lens: torch.Tensor, - prefix_lens_cpu: torch.Tensor, - seq_lens: torch.Tensor, - seq_lens_cpu: torch.Tensor, - last_loc: torch.Tensor, - extend_num_tokens: int, - ): - """Allocate only logical indices without hisparse device indices. - - Used in the direct-to-host transfer path where KV data is written - directly to host memory by the prefill node, skipping GPU staging. - """ - return self.logical_attn_allocator.alloc_extend( - prefix_lens, - prefix_lens_cpu, - seq_lens, - seq_lens_cpu, - last_loc, - extend_num_tokens, - ) - - def alloc_device_buffer(self, allocated_indices, need_size: int): - assert need_size % self.page_size == 0 - # clear original reference and isolate the buffer from outside addressing, allocate new buffer if needed - hisparse_indices = self.full_to_hisparse_device_index_mapping[allocated_indices] - self.full_to_hisparse_device_index_mapping[allocated_indices] = 0 - # Filter valid (non-zero) hisparse indices. - # In the direct-to-host path, mapping is all zeros since no hisparse - # device indices were pre-allocated. - hisparse_indices = hisparse_indices[hisparse_indices > 0] - if len(hisparse_indices) >= need_size: - buffer_indices = hisparse_indices[:need_size] - self.free_hisparse_indices(hisparse_indices[need_size:]) - else: - # page alignment, claiming the residual space for an incomplete page - page_residual_length = len(hisparse_indices) % self.page_size - if page_residual_length != 0: - hisparse_indices = torch.cat( - [ - hisparse_indices, - torch.arange( - hisparse_indices[-1] + 1, - hisparse_indices[-1] - + self.page_size - - page_residual_length - + 1, - device=self.device, - ), - ] - ) - extra_indices = self.hisparse_attn_allocator.alloc( - need_size - len(hisparse_indices) - ) - assert ( - extra_indices is not None - ), "Hisparse allocation failed in alloc_device_buffer" - buffer_indices = torch.cat([hisparse_indices, extra_indices]) - return buffer_indices - - def free_hisparse_indices(self, buffer_indices: torch.Tensor): - # disable free group mechanism for device buffer free - self.hisparse_attn_allocator.is_not_in_free_group = True - self.hisparse_attn_allocator.free(buffer_indices[buffer_indices > 0]) - - def get_last_loc_compressed(self, last_locs: torch.Tensor): - return last_locs - - def get_last_loc_hisparse_device(self, last_locs: torch.Tensor): - return self._kvcache._translate_loc_to_hisparse_device(last_locs) - - def alloc_extend( - self, - prefix_lens: torch.Tensor, - prefix_lens_cpu: torch.Tensor, - seq_lens: torch.Tensor, - seq_lens_cpu: torch.Tensor, - last_loc: torch.Tensor, # last_loc for full layers - extend_num_tokens: int, - ): - assert self.page_size > 1 - - num_new_pages = get_num_new_pages( - seq_lens=seq_lens_cpu, page_size=self.page_size, prefix_lens=prefix_lens_cpu - ) - if ( - num_new_pages - > self.logical_attn_allocator.available_size() // self.page_size - ): - return None - if ( - num_new_pages - > self.hisparse_attn_allocator.available_size() // self.page_size - ): - return None - - logical_indices = self.logical_attn_allocator.alloc_extend( - prefix_lens, - prefix_lens_cpu, - seq_lens, - seq_lens_cpu, - last_loc, - extend_num_tokens, - ) - assert logical_indices is not None, "Logical allocation failed in alloc_extend" - - hisparse_last_loc = self.get_last_loc_hisparse_device(last_loc) - hisparse_indices = self.hisparse_attn_allocator.alloc_extend( - prefix_lens, - prefix_lens_cpu, - seq_lens, - seq_lens_cpu, - hisparse_last_loc, - len(logical_indices), - num_new_pages=num_new_pages, - ) - assert ( - hisparse_indices is not None - ), "Hisparse allocation failed in alloc_extend" - self.full_to_hisparse_device_index_mapping[logical_indices] = hisparse_indices - return logical_indices - - def alloc_decode( - self, - seq_lens: torch.Tensor, - seq_lens_cpu: torch.Tensor, - last_loc: torch.Tensor, # last_loc for full layers - ): - return self.logical_attn_allocator.alloc_decode( - seq_lens, seq_lens_cpu, last_loc - ) - - def free_hisparse(self, free_indices: torch.Tensor): - hisparse_indices = self._kvcache._translate_loc_to_hisparse_device(free_indices) - hisparse_indices = hisparse_indices[hisparse_indices > 0] - self.free_hisparse_indices(hisparse_indices) - self.full_to_hisparse_device_index_mapping[free_indices] = 0 - - def clear(self): - self.logical_attn_allocator.clear() - self.hisparse_attn_allocator.clear() - # Note: the last item is -1, we don't clear it, see the comment in __init__ - self.full_to_hisparse_device_index_mapping[:-1].fill_(0) - self.is_not_in_free_group = True - self.free_group = [] - - def free_group_begin(self): - return - - def free_group_end(self): - return - - def free(self, free_index: torch.Tensor): - if free_index.numel() == 0: - return - if self.is_not_in_free_group: - self.logical_attn_allocator.free(free_index) - self.free_hisparse(free_index) - else: - self.free_group.append(free_index) - assert ( - self.logical_attn_allocator.available_size() - <= self.logical_attn_allocator.size - ) - assert ( - self.hisparse_attn_allocator.available_size() - <= self.hisparse_attn_allocator.size - ) - - -class DeepSeekV4HiSparseTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator): - - def __init__( - self, - logical_attn_allocator: BaseTokenToKVPoolAllocator, - ): - assert isinstance(logical_attn_allocator._kvcache, DeepSeekV4TokenToKVPool) - assert isinstance( - logical_attn_allocator._kvcache.c4_kv_pool, HiSparseC4DevicePool - ) - self.compress_ratio = 4 - - self.hisparse_kvcache = logical_attn_allocator._kvcache.c4_kv_pool - self._size_full = logical_attn_allocator.size_full - self._size_hisparse = self.hisparse_kvcache.size - - self.dtype = self.hisparse_kvcache.dtype - self.device = self.hisparse_kvcache.device - # Keep the public page_size as the logical DSV4 full/SWA page size. - # C4 HiSparse allocation/device-buffer code must use the compressed page size. - self.page_size = logical_attn_allocator.page_size - self.hisparse_page_size = self.hisparse_kvcache.page_size - - self.logical_attn_allocator = logical_attn_allocator - self._kvcache = logical_attn_allocator._kvcache - self.hisparse_attn_allocator = PagedTokenToKVPoolAllocator( - self._size_hisparse, - self.hisparse_page_size, - self.dtype, - self.device, - self.hisparse_kvcache, - logical_attn_allocator.need_sort, - ) - - self.full_to_hisparse_device_index_mapping = torch.cat( - [ - torch.zeros( - self._kvcache.c4_logical_size + self.hisparse_page_size, - dtype=torch.int64, - device=self.device, - ), - torch.tensor([-1], dtype=torch.int64, device=self.device), - ] - ) - - self.need_sort = logical_attn_allocator.need_sort - self.free_pages = None - self.release_pages = None - self.is_not_in_free_group = True - self.free_group = [] - self.clear() - - self.hisparse_kvcache.register_mapping( - weakref.proxy(self.full_to_hisparse_device_index_mapping) - ) - - @property - def size_full(self) -> int: - return self._size_full - - @property - def size(self) -> int: - return self.logical_attn_allocator.size - - @property - def size_swa(self) -> int: - return self.logical_attn_allocator.size_swa - - @property - def full_to_swa_index_mapping(self): - return self.logical_attn_allocator.full_to_swa_index_mapping - - def debug_print(self) -> str: - msg = self.logical_attn_allocator.debug_print() - msg += ( - f"#hisparse-available-size: " - f"{self.hisparse_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): - return self.logical_attn_allocator.translate_loc_from_full_to_swa(kv_indices) - - def full_available_size(self): - return min( - self.logical_attn_allocator.full_available_size(), - self.hisparse_attn_allocator.available_size() * self.compress_ratio, - ) - - def swa_available_size(self): - return self.logical_attn_allocator.swa_available_size() - - def free_swa(self, free_indices: torch.Tensor): - self.logical_attn_allocator.free_swa(free_indices) - - def available_size(self) -> int: - return min( - self.logical_attn_allocator.available_size(), - self.hisparse_attn_allocator.available_size() * self.compress_ratio, - ) - - def alloc(self, need_size: int): - raise NotImplementedError( - "DeepSeek V4 HiSparse allocator does not support direct token allocation; " - "use alloc_extend or alloc_decode instead." - ) - - def alloc_logical_only( - self, - prefix_lens: torch.Tensor, - prefix_lens_cpu: torch.Tensor, - seq_lens: torch.Tensor, - seq_lens_cpu: torch.Tensor, - last_loc: torch.Tensor, - extend_num_tokens: int, - ): - """Allocate decode logical indices without allocating C4 hisparse device pages.""" - return self.logical_attn_allocator.alloc_extend( - prefix_lens, - prefix_lens_cpu, - seq_lens, - seq_lens_cpu, - last_loc, - extend_num_tokens, - ) - - def alloc_device_buffer(self, allocated_indices, need_size: int): - assert need_size % self.hisparse_page_size == 0 - hisparse_indices = self.full_to_hisparse_device_index_mapping[allocated_indices] - self.full_to_hisparse_device_index_mapping[allocated_indices] = 0 - hisparse_indices = hisparse_indices[hisparse_indices > 0] - - device_buffer_size = need_size - self.hisparse_page_size - P = len(hisparse_indices) - if P > device_buffer_size + 1: - newest_src = hisparse_indices[P - 1].clone() - old_at_dbs = hisparse_indices[device_buffer_size].clone() - hisparse_indices[device_buffer_size] = newest_src - hisparse_indices[P - 1] = old_at_dbs - - if len(hisparse_indices) >= need_size: - buffer_indices = hisparse_indices[:need_size] - surplus = hisparse_indices[need_size:] - if surplus.numel() > 0: - buffer_pages = torch.unique(buffer_indices // self.hisparse_page_size) - surplus_pages = torch.unique(surplus // self.hisparse_page_size) - pure_surplus = surplus_pages[~torch.isin(surplus_pages, buffer_pages)] - if pure_surplus.numel() > 0: - self.hisparse_attn_allocator.is_not_in_free_group = True - self.hisparse_attn_allocator.free( - pure_surplus * self.hisparse_page_size - ) - else: - page_residual_length = len(hisparse_indices) % self.hisparse_page_size - if page_residual_length != 0: - hisparse_indices = torch.cat( - [ - hisparse_indices, - torch.arange( - hisparse_indices[-1] + 1, - hisparse_indices[-1] - + self.hisparse_page_size - - page_residual_length - + 1, - device=self.device, - ), - ] - ) - extra_indices = self.hisparse_attn_allocator.alloc( - need_size - len(hisparse_indices) - ) - assert ( - extra_indices is not None - ), "Hisparse allocation failed in alloc_device_buffer" - buffer_indices = torch.cat([hisparse_indices, extra_indices]) - return buffer_indices - - def free_hisparse_indices(self, buffer_indices: torch.Tensor): - self.hisparse_attn_allocator.is_not_in_free_group = True - self.hisparse_attn_allocator.free(buffer_indices[buffer_indices > 0]) - - def get_last_loc_compressed(self, last_locs: torch.Tensor): - return (last_locs - 3) // self.compress_ratio - - def get_last_loc_hisparse_device(self, last_locs: torch.Tensor): - return self.hisparse_kvcache._translate_loc_to_hisparse_device( - self.get_last_loc_compressed(last_locs) - ) - - 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, - extend_num_tokens: int, - ): - assert self.page_size > 1 - - num_new_pages_logical = get_num_new_pages( - seq_lens=seq_lens_cpu, page_size=self.page_size, prefix_lens=prefix_lens_cpu - ) - num_new_pages_hisparse = get_num_new_pages( - seq_lens=seq_lens_cpu // self.compress_ratio, - page_size=self.hisparse_page_size, - prefix_lens=prefix_lens_cpu // self.compress_ratio, - ) - if ( - num_new_pages_logical - > self.logical_attn_allocator.available_size() // self.page_size - ): - return None - if ( - num_new_pages_hisparse - > self.hisparse_attn_allocator.available_size() // self.hisparse_page_size - ): - return None - - logical_indices = self.logical_attn_allocator.alloc_extend( - prefix_lens, - prefix_lens_cpu, - seq_lens, - seq_lens_cpu, - last_loc, - extend_num_tokens, - ) - assert logical_indices is not None, "Logical allocation failed in alloc_extend" - - compressed_logical_indices = ( - self.hisparse_kvcache.translate_loc_from_full_to_compressed(logical_indices) - ) - hisparse_last_loc = self.get_last_loc_hisparse_device(last_loc) - hisparse_indices = self.hisparse_attn_allocator.alloc_extend( - prefix_lens // self.compress_ratio, - prefix_lens_cpu // self.compress_ratio, - seq_lens // self.compress_ratio, - seq_lens_cpu // self.compress_ratio, - hisparse_last_loc, - len(compressed_logical_indices), - ) - assert ( - hisparse_indices is not None - ), "Hisparse allocation failed in alloc_extend" - - self.full_to_hisparse_device_index_mapping[compressed_logical_indices] = ( - hisparse_indices.to(torch.int64) - ) - return logical_indices - - def alloc_decode( - self, - seq_lens: torch.Tensor, - seq_lens_cpu: torch.Tensor, - last_loc: torch.Tensor, - ): - return self.logical_attn_allocator.alloc_decode( - seq_lens, seq_lens_cpu, last_loc - ) - - def free_compressed(self, compressed_indices: torch.Tensor): - hisparse_indices = self.hisparse_kvcache.translate_loc_to_hisparse_device( - compressed_indices - ) - hisparse_indices = hisparse_indices[hisparse_indices > 0] - self.free_hisparse_indices(hisparse_indices) - self.full_to_hisparse_device_index_mapping[compressed_indices] = 0 - - def free_hisparse(self, free_indices: torch.Tensor): - compressed_indices = ( - self.hisparse_kvcache.translate_loc_from_full_to_compressed(free_indices) - ) - self.free_compressed(compressed_indices) - - def clear(self): - self.logical_attn_allocator.clear() - self.hisparse_attn_allocator.clear() - - self.full_to_hisparse_device_index_mapping[:-1].fill_(0) - self.is_not_in_free_group = True - self.free_group = [] - - def free(self, free_index: torch.Tensor): - if free_index.numel() == 0: - return - - if self.is_not_in_free_group: - self.logical_attn_allocator.free(free_index) - else: - self.free_group.append(free_index) - assert ( - self.logical_attn_allocator.available_size() - <= self.logical_attn_allocator.size - ) - assert ( - self.hisparse_attn_allocator.available_size() - <= self.hisparse_attn_allocator.size - ) 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 163acb873..e81b649f9 100644 --- a/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py +++ b/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py @@ -17,14 +17,14 @@ from sglang.srt.mem_cache.allocator import ( PagedTokenToKVPoolAllocator, TokenToKVPoolAllocator, ) +from sglang.srt.mem_cache.allocator.hisparse import ( + DeepSeekV4HiSparseTokenToKVPoolAllocator, + HiSparseTokenToKVPoolAllocator, +) from sglang.srt.mem_cache.allocator.swa import SWATokenToKVPoolAllocator from sglang.srt.mem_cache.common import get_req_to_token_extra_context_len from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool -from sglang.srt.mem_cache.hisparse_memory_pool import ( - DeepSeekV4HiSparseTokenToKVPoolAllocator, - HiSparseDSATokenToKVPool, - HiSparseTokenToKVPoolAllocator, -) +from sglang.srt.mem_cache.hisparse_memory_pool import HiSparseDSATokenToKVPool from sglang.srt.mem_cache.memory_pool import ( DSATokenToKVPool, HybridLinearKVPool, diff --git a/test/registered/unit/managers/test_hisparse_unit.py b/test/registered/unit/managers/test_hisparse_unit.py index a31c39a72..e1f245ba6 100644 --- a/test/registered/unit/managers/test_hisparse_unit.py +++ b/test/registered/unit/managers/test_hisparse_unit.py @@ -94,10 +94,10 @@ class TestHiSparseUnit(unittest.TestCase): global_page_size = 1 if is_hip() else PAGE_SIZE - from sglang.srt.mem_cache.hisparse_memory_pool import ( - HiSparseDSATokenToKVPool, + from sglang.srt.mem_cache.allocator.hisparse import ( HiSparseTokenToKVPoolAllocator, ) + from sglang.srt.mem_cache.hisparse_memory_pool import HiSparseDSATokenToKVPool cls.device_pool = HiSparseDSATokenToKVPool( size=SIZE,