[mem_cache][5/N] refactor: extract host KV cache base layer into pool_host package

python/sglang/srt/managers/cache_controller.py

@@ -31,7 +31,7 @@
 
 if TYPE_CHECKING:
     from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator
-    from sglang.srt.mem_cache.memory_pool_host import HostKVCache
+    from sglang.srt.mem_cache.pool_host import HostKVCache
 
 from sglang.srt.distributed import (
     get_pipeline_model_parallel_rank,

python/sglang/srt/mem_cache/hicache_storage.py

@@ -14,7 +14,7 @@
 from sglang.srt.environ import envs
 
 if TYPE_CHECKING:
-    from sglang.srt.mem_cache.memory_pool_host import HostKVCache
+    from sglang.srt.mem_cache.pool_host import HostKVCache
 
 logger = logging.getLogger(__name__)
 

python/sglang/srt/mem_cache/memory_pool_host.py

@@ -1,11 +1,8 @@
 from __future__ import annotations
 
-import abc
 import logging
 import threading
-from collections import defaultdict
 from dataclasses import dataclass
-from functools import wraps
 from typing import TYPE_CHECKING, Any, Callable, Optional
 
 if TYPE_CHECKING:
@@ -40,12 +37,10 @@
 from sglang.jit_kernel.hisparse import transfer_cache_dsv4_mla
 from sglang.srt.mem_cache.memory_pool import (
     DSATokenToKVPool,
-    KVCache,
     MambaPool,
     MHATokenToKVPool,
     MLATokenToKVPool,
 )
-from sglang.srt.mem_cache.mmap_allocator import alloc_mmap
 from sglang.srt.utils import is_cuda, is_hip, is_mps, is_npu, is_xpu
 
 _is_cuda = is_cuda()
@@ -74,319 +69,19 @@
 
 logger = logging.getLogger(__name__)
 
-# Host RAM to leave free when sizing HiCache pools (OS, other processes).
-HICACHE_HOST_MEMORY_RESERVE_BYTES: int = 10 * (1024**3)
 
-_WRITE_BACK_STAGING_PAGE_CHUNK = 64
-
-
-def synchronized(func):
-    @wraps(func)
-    def wrapper(self, *args, **kwargs):
-        with self.lock:
-            return func(self, *args, **kwargs)
-
-    return wrapper
-
-
-class HostTensorAllocator:
-    def __init__(self):
-        """Initialize the HostTensorAllocator."""
-        self.dtype = None
-        self.dims = None
-
-    def allocate(self, dims: tuple, dtype: torch.dtype, device: str) -> torch.Tensor:
-        assert (
-            device == "cpu"
-        ), f"HostTensorAllocator only supports CPU allocations; got device={device!r}"
-        self.dtype = dtype
-        self.dims = dims
-        return alloc_mmap(dims, dtype)
-
-
-class HiSparseHostPoolMixin:
-    def _round_up_to_page_size(self, size: int) -> int:
-        return (size + self.page_size - 1) // self.page_size * self.page_size
-
-    def alloc_page(self, num_pages: int) -> Optional[torch.Tensor]:
-        return self.alloc(num_pages * self.page_size)
-
-    def alloc_paged_token_slots(
-        self,
-        req_to_host_pool: torch.Tensor,
-        req_to_host_pool_allocated_len: torch.Tensor,
-        req_pool_idx: int,
-        start_pos: int,
-        num_tokens: int,
-    ) -> torch.Tensor:
-        """Allocate request host slots by page and return token-granular slots."""
-        device = req_to_host_pool.device
-        if num_tokens <= 0:
-            return torch.empty((0,), dtype=torch.int64, device=device)
-
-        allocated_len = int(req_to_host_pool_allocated_len[req_pool_idx])
-        end_pos = start_pos + num_tokens
-        page_end = self._round_up_to_page_size(end_pos)
-        assert start_pos <= allocated_len
-
-        if page_end > allocated_len:
-            num_new_pages = (page_end - allocated_len) // self.page_size
-            host_locs = self.alloc_page(num_new_pages)
-            if host_locs is None:
-                logger.error(
-                    "HiSparse: host mem pool alloc failed for %d host pages "
-                    "(req_pool_idx=%d, start_pos=%d, num_tokens=%d)",
-                    num_new_pages,
-                    req_pool_idx,
-                    start_pos,
-                    num_tokens,
-                )
-                raise RuntimeError(
-                    f"HiSparse host mem pool alloc failed for {num_new_pages} pages"
-                )
-
-            req_to_host_pool[req_pool_idx, allocated_len:page_end] = host_locs.to(
-                device=device, non_blocking=True
-            )
-            req_to_host_pool_allocated_len[req_pool_idx] = page_end
-
-        return req_to_host_pool[req_pool_idx, start_pos:end_pos]
-
-    def allocated_host_indices(
-        self,
-        req_to_host_pool: torch.Tensor,
-        req_pool_idx: int,
-        allocated_len: int,
-    ) -> torch.Tensor:
-        allocated_len = int(allocated_len)
-        host_len = min(
-            self._round_up_to_page_size(allocated_len),
-            req_to_host_pool.shape[1],
-        )
-        host_indices = req_to_host_pool[req_pool_idx, :host_len]
-        return host_indices[host_indices >= 0]
-
-
-def get_allocator_from_storage(allocator_type):
-    if allocator_type == "mooncake":
-        try:
-            from sglang.srt.mem_cache.storage.mooncake_store.mooncake_store import (
-                MooncakeHostTensorAllocator,
-            )
-
-            return MooncakeHostTensorAllocator()
-        except ImportError:
-            logger.warning(
-                "Mooncake's tensor allocator requires mooncake >= 0.3.8.post1. "
-                "Please upgrade Mooncake by 'pip install mooncake-transfer-engine --upgrade'. "
-                "Fallback to use default allocator."
-            )
-            return HostTensorAllocator()
-    else:
-        return HostTensorAllocator()
-
-
-def _cuda_host_register(buffer: torch.Tensor) -> None:
-    cudart = torch.cuda.cudart()
-    n_bytes = buffer.numel() * buffer.element_size()
-    rc = cudart.cudaHostRegister(buffer.data_ptr(), n_bytes, 0)
-    if int(rc) != 0:
-        raise RuntimeError(
-            f"cudaHostRegister failed (rc={int(rc)}, "
-            f"{cudart.cudaGetErrorString(rc)}) for ptr={buffer.data_ptr():#x} "
-            f"size={n_bytes}; host buffer is not pinned and device transfers "
-            f"may silently return stale data."
-        )
-
-
-def alloc_with_host_register(
-    dims: tuple,
-    dtype: torch.dtype,
-    device: str,
-    pin_memory: bool,
-    allocator: HostTensorAllocator,
-) -> torch.Tensor:
-    """
-    Allocate tensor and register host memory with cudaHostRegister.
-    CudaHostRegister only applies when pin_memory=True.
-    """
-    buffer = allocator.allocate(dims, dtype=dtype, device=device)
-    if pin_memory:
-        _cuda_host_register(buffer)
-    return buffer
-
-
-def alloc_with_pin_memory(
-    dims: tuple,
-    dtype: torch.dtype,
-    device: str,
-    pin_memory: bool,
-    allocator: None,
-) -> torch.Tensor:
-    """
-    Allocate tensor using PyTorch's built-in pin_memory flag.
-    """
-    buffer = torch.empty(dims, dtype=dtype, device=device, pin_memory=pin_memory)
-    return buffer
-
-
-ALLOC_MEMORY_FUNCS = defaultdict(
-    lambda: alloc_with_host_register,
-    {
-        "npu": alloc_with_pin_memory,
-        "musa": alloc_with_pin_memory,
-    },
+from sglang.srt.mem_cache.pool_host import HostKVCache
+from sglang.srt.mem_cache.pool_host.base import (
+    HICACHE_HOST_MEMORY_RESERVE_BYTES,
+    synchronized,
 )
+from sglang.srt.mem_cache.pool_host.common import (
+    ALLOC_MEMORY_FUNCS,
+    get_allocator_from_storage,
+)
+from sglang.srt.mem_cache.pool_host.hisparse import HiSparseHostPoolMixin
 
-
-class HostKVCache(abc.ABC):
-
-    def __init__(
-        self,
-        device_pool: KVCache,
-        host_to_device_ratio: float,
-        host_size: int,
-        page_size: int,
-        layout: str,
-        pin_memory: bool,
-        device: str,
-        allocator_type: str = "default",
-    ):
-        self.device_pool = device_pool
-        self.page_size = page_size
-        self.layout = layout
-        self.pin_memory = pin_memory
-        self.device = device
-        self.allocator = get_allocator_from_storage(allocator_type)
-        self.can_use_write_back_jit = False
-
-        self.dtype = device_pool.store_dtype
-        self.size_per_token = self.get_size_per_token()
-        if host_size > 0:
-            self.size = int(host_size * 1e9 // self.size_per_token)
-        else:
-            self.size = int(device_pool.size * host_to_device_ratio)
-        # Align up the host memory pool size to the page size
-        self.page_num = self.size // self.page_size + 1
-        self.size = self.page_num * self.page_size
-        self.start_layer = device_pool.start_layer
-        self.end_layer = device_pool.end_layer
-
-        assert (
-            self.size > device_pool.size
-        ), "The host memory should be larger than the device memory with the current protocol"
-
-        # Verify there is enough available host memory.
-        host_mem = psutil.virtual_memory()
-        requested_bytes = self.size * self.size_per_token
-        available_bytes = host_mem.available - HICACHE_HOST_MEMORY_RESERVE_BYTES
-        if requested_bytes > available_bytes:
-            raise ValueError(
-                f"Not enough host memory available. Requesting "
-                f"{requested_bytes / 1e9:.2f} GB but only have "
-                f"{available_bytes / 1e9:.2f} GB free. Please reduce the "
-                f"size of the hierarchical cache."
-            )
-        else:
-            logger.info(
-                f"Allocating {requested_bytes / 1e9:.2f} GB host memory for hierarchical KV cache."
-            )
-
-        self.kv_buffer = self.init_kv_buffer()
-
-        # A lock for synchronized operations on memory allocation and state transitions.
-        self.lock = threading.RLock()
-        self.clear()
-
-    @abc.abstractmethod
-    def get_size_per_token(self):
-        raise NotImplementedError()
-
-    @abc.abstractmethod
-    def init_kv_buffer(self):
-        raise NotImplementedError()
-
-    @abc.abstractmethod
-    def load_to_device_per_layer(
-        self, device_pool, host_indices, device_indices, layer_id, io_backend
-    ) -> None:
-        """
-        Load KV data from the host memory pool to the device memory pool for a specific layer.
-        """
-        raise NotImplementedError()
-
-    @abc.abstractmethod
-    def backup_from_device_all_layer(
-        self, device_pool, host_indices, device_indices, io_backend
-    ) -> None:
-        """
-        Backup KV data from the device memory pool to the host memory pool for all layers.
-        """
-        raise NotImplementedError()
-
-    @abc.abstractmethod
-    def get_data_page(self, index, flat: bool = True) -> torch.Tensor:
-        """
-        Get a flat data page from the host memory pool.
-        """
-        raise NotImplementedError()
-
-    @abc.abstractmethod
-    def get_dummy_flat_data_page(self) -> torch.Tensor:
-        """
-        Get a dummy flat data page from the host memory pool.
-        This is used for prefetching or initializing empty pages.
-        """
-        raise NotImplementedError()
-
-    @abc.abstractmethod
-    def set_from_flat_data_page(self, index: int, data_page: torch.Tensor) -> None:
-        """
-        Set a flat data page to the host memory pool.
-        """
-        raise NotImplementedError()
-
-    def is_stride_page_aligned(self, page_size_bytes: int = 4096) -> bool:
-        """Return True if per-page strides are multiples of *page_size_bytes*.
-
-        Subclasses should override this with a layout-specific stride formula.
-        This base implementation logs a warning and returns False (safe default).
-        """
-        logger.warning(
-            "%s does not implement is_stride_page_aligned(); assuming not aligned. "
-            "O_DIRECT with a file-based NIXL backend will fall back to copy mode for this pool.",
-            type(self).__name__,
-        )
-        return False
-
-    @synchronized
-    def clear(self):
-        # Initialize memory states and tracking structures.
-        self.mem_state = torch.zeros(
-            (self.size,), dtype=torch.uint8, device=self.device
-        )
-        self.free_slots = torch.arange(self.size, dtype=torch.int64)
-
-    def available_size(self):
-        return len(self.free_slots)
-
-    @synchronized
-    def alloc(self, need_size: int) -> Optional[torch.Tensor]:
-        assert (
-            need_size % self.page_size == 0
-        ), "The requested size should be a multiple of the page size."
-        if need_size > self.available_size():
-            return None
-
-        select_index = self.free_slots[:need_size]
-        self.free_slots = self.free_slots[need_size:]
-
-        return select_index
-
-    @synchronized
-    def free(self, indices: torch.Tensor) -> int:
-        self.free_slots = torch.cat([self.free_slots, indices.cpu()])
-        return len(indices)
+_WRITE_BACK_STAGING_PAGE_CHUNK = 64
 
 
 class MHATokenToKVPoolHost(HostKVCache):

python/sglang/srt/mem_cache/pool_host/__init__.py

@@ -0,0 +1,7 @@
+from sglang.srt.mem_cache.pool_host.base import HostKVCache
+from sglang.srt.mem_cache.pool_host.common import HostTensorAllocator
+
+__all__ = [
+    "HostKVCache",
+    "HostTensorAllocator",
+]