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

Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

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,

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
 

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
+        )

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

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
-        )

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 (
+from sglang.srt.mem_cache.hisparse_memory_pool import HiSparseDSATokenToKVPool
-    DeepSeekV4HiSparseTokenToKVPoolAllocator,
-    HiSparseDSATokenToKVPool,
-    HiSparseTokenToKVPoolAllocator,
-)
 from sglang.srt.mem_cache.memory_pool import (
     DSATokenToKVPool,
     HybridLinearKVPool,

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 (
+        from sglang.srt.mem_cache.allocator.hisparse import (
-            HiSparseDSATokenToKVPool,
             HiSparseTokenToKVPoolAllocator,
         )
+        from sglang.srt.mem_cache.hisparse_memory_pool import HiSparseDSATokenToKVPool
 
         cls.device_pool = HiSparseDSATokenToKVPool(
             size=SIZE,