minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[mem_cache][3/N] refactor: move HiSparse allocators to allocator/hisparse.py (#26678)

Browse Source 
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
shuwenn
2026-06-11 19:50:24 +08:00
committed by GitHub
co-authored by Claude Opus 4.8
parent be45745f38
commit 66076f2409
7 changed files with 582 additions and 578 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/managers/hisparse_coordinator.py
+4 -2
View File
@@ -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,
python/sglang/srt/managers/schedule_policy.py
+3 -3
View File
@@ -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
python/sglang/srt/mem_cache/allocator/hisparse.py
+565
View File
@@ -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
)
python/sglang/srt/mem_cache/chunk_cache.py
+3 -3
View File
@@ -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
python/sglang/srt/mem_cache/hisparse_memory_pool.py
-563
View File
@@ -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
)
python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py
+5 -5
View File
@@ -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,