diff --git a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py index bedf086a6..d071d7534 100644 --- a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py +++ b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py @@ -18,10 +18,10 @@ from sglang.srt.mem_cache.memory_pool_host import ( DSAIndexerPoolHost, HostPoolGroup, LogicalHostPool, - MambaPoolHost, PoolEntry, ) from sglang.srt.mem_cache.pool_host.common import get_allocator_type +from sglang.srt.mem_cache.pool_host.mamba import MambaPoolHost from sglang.srt.mem_cache.pool_host.mha import ( MHATokenToKOnlyPoolHost, get_mha_host_pool_cls, diff --git a/python/sglang/srt/mem_cache/memory_pool_host.py b/python/sglang/srt/mem_cache/memory_pool_host.py index 0024ebc6e..1ed1645b3 100644 --- a/python/sglang/srt/mem_cache/memory_pool_host.py +++ b/python/sglang/srt/mem_cache/memory_pool_host.py @@ -9,7 +9,6 @@ if TYPE_CHECKING: from sglang.srt.mem_cache.hicache_storage import PoolName from sglang.srt.mem_cache.pool_host.mla import MLATokenToKVPoolHost -import numpy as np import psutil import torch @@ -20,7 +19,9 @@ from sglang.kernels.ops.kvcache.hicache import ( transfer_hicache_all_layer_mla_staged_lf_pf as jit_transfer_hicache_all_layer_mla_staged_lf_pf, ) from sglang.kernels.ops.kvcache.hisparse import transfer_cache_dsv4_mla -from sglang.srt.mem_cache.memory_pool import DSATokenToKVPool, MambaPool +from sglang.srt.mem_cache.memory_pool import ( + DSATokenToKVPool, +) from sglang.srt.utils import is_cuda, is_hip, is_mps, is_npu, is_xpu _is_cuda = is_cuda() @@ -38,13 +39,6 @@ if _is_cuda or _is_hip: transfer_kv_per_layer_mla, transfer_kv_per_layer_mla_pf_lf, ) -if _is_cuda or _is_hip: - from sglang.kernels.ops.mamba.transfer_mamba import ( - transfer_kv_mamba_lf_pf, - transfer_kv_mamba_pf_lf, - ) -if _is_npu: - pass logger = logging.getLogger(__name__) @@ -53,7 +47,6 @@ from sglang.srt.mem_cache.pool_host import HostKVCache from sglang.srt.mem_cache.pool_host.base import ( _WRITE_BACK_STAGING_PAGE_CHUNK, HICACHE_HOST_MEMORY_RESERVE_BYTES, - sync_fixed_hicache_size, synchronized, ) from sglang.srt.mem_cache.pool_host.common import ( @@ -62,578 +55,6 @@ from sglang.srt.mem_cache.pool_host.common import ( ) from sglang.srt.mem_cache.pool_host.hisparse import HiSparseHostPoolMixin - -class MambaPoolHost(HostKVCache): - def __init__( - self, - device_pool: MambaPool, - host_to_device_ratio: float, - host_size: int, - pin_memory: bool = True, - device: str = "cpu", - allocator_type: str = "default", - layout: str = "layer_first", - ): - self.device_pool = device_pool - self.page_size = 1 - - assert layout in [ - "page_first", - "page_first_direct", - ], f"Unsupported layout: {layout}" - - self.layout = layout - self.pin_memory = pin_memory - self.device = device - self.allocator = get_allocator_from_storage(allocator_type) - self.num_mamba_layers = device_pool.num_mamba_layers - - self.conv_state_shapes = [ - conv_state.shape[2:] for conv_state in device_pool.mamba_cache.conv - ] - self.temporal_state_shape = device_pool.mamba_cache.temporal.shape[2:] - self.temporal_state_elem_size = int(np.prod(self.temporal_state_shape)) - self.conv_state_elem_sizes = [ - int(np.prod(conv_shape)) for conv_shape in self.conv_state_shapes - ] - self.conv_dtype = device_pool.mamba_cache.conv[0].dtype - self.temporal_dtype = device_pool.mamba_cache.temporal.dtype - self.dtype = self.conv_dtype - self.size_per_token = self.get_size_per_token() - - if host_size > 0: - self.size = sync_fixed_hicache_size( - int(host_size * 1e9 // self.size_per_token), host_size - ) - else: - self.size = int(device_pool.size * host_to_device_ratio) - - self.page_num = self.size // self.page_size + 1 - self.size = self.page_num * self.page_size - - if self.size <= device_pool.size: - logger.warning( - "HiCache host KV pool (%d tokens) is smaller than the device pool (%d tokens);" - "L2 cache effectiveness is reduced." - "Consider increasing --hicache-ratio (or --hicache-size) for higher L2 cache hit rate.", - self.size, - device_pool.size, - ) - - 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." - ) - logger.info( - "Allocating %.2f GB host memory for hierarchical Mamba cache (layout=%s).", - requested_bytes / 1e9, - self.layout, - ) - - self.temporal_device_ptrs = torch.tensor( - [ - device_pool.mamba_cache.temporal[i].data_ptr() - for i in range(self.num_mamba_layers) - ], - dtype=torch.uint64, - device=self.device_pool.device, - ) - self.conv_device_ptrs = [ - torch.tensor( - [conv_state[i].data_ptr() for i in range(self.num_mamba_layers)], - dtype=torch.uint64, - device=self.device_pool.device, - ) - for conv_state in device_pool.mamba_cache.conv - ] - - self.init_kv_buffer() - self._init_write_back_staging_buffers() - self.lock = threading.RLock() - self.clear() - - def init_kv_buffer(self): - _host_alloc = ALLOC_MEMORY_FUNCS[self.device_pool.device] - - def alloc_func(dims, *, dtype, device, pin_memory, allocator): - # conv-only linear attention has no ssm state: mmap can't map the - # 0-element temporal buffer, so hand back a plain empty tensor. - if np.prod(dims) == 0: - return torch.empty(dims, dtype=dtype, device=device) - return _host_alloc( - dims, - dtype=dtype, - device=device, - pin_memory=pin_memory, - allocator=allocator, - ) - - if self.layout in ["page_first", "page_first_direct"]: - # page-first: (page_num, num_layers, 1, *shape) — per-page data is contiguous - temporal_dims = ( - self.size, - self.num_mamba_layers, - 1, - ) + self.temporal_state_shape - self.temporal_buffer = alloc_func( - temporal_dims, - dtype=self.temporal_dtype, - device=self.device, - pin_memory=self.pin_memory, - allocator=self.allocator, - ) - self.conv_buffer = [] - for conv_shape in self.conv_state_shapes: - conv_dims = (self.size, self.num_mamba_layers, 1) + conv_shape - self.conv_buffer.append( - alloc_func( - conv_dims, - dtype=self.conv_dtype, - device=self.device, - pin_memory=self.pin_memory, - allocator=self.allocator, - ) - ) - else: - # layer-first: (num_layers, size, *shape) - temporal_dims = ( - self.num_mamba_layers, - self.size, - ) + self.temporal_state_shape - self.temporal_buffer = alloc_func( - temporal_dims, - dtype=self.temporal_dtype, - device=self.device, - pin_memory=self.pin_memory, - allocator=self.allocator, - ) - self.conv_buffer = [] - for conv_shape in self.conv_state_shapes: - conv_dims = (self.num_mamba_layers, self.size) + conv_shape - self.conv_buffer.append( - alloc_func( - conv_dims, - dtype=self.conv_dtype, - device=self.device, - pin_memory=self.pin_memory, - allocator=self.allocator, - ) - ) - - def _init_write_back_staging_buffers(self): - self.temporal_staging_buffer = None - self.conv_staging_buffers = [None] * len(self.conv_buffer) - # Must be True: HostPoolGroup computes can_use_write_back_jit as AND of - # all pools. When True, start_writing() keeps indices on CPU, which MLA's - # staged write-back kernel requires. MambaPoolHost's own backup path does - # not check this flag — it routes by layout + io_backend instead. - self.can_use_write_back_jit = True - self._temporal_can_use_jit = False - self._conv_can_use_jit = [False] * len(self.conv_buffer) - - def get_hybrid_pool_buffer(self): - # Expose all mamba host tensors that need Mooncake buffer registration. - return [self.temporal_buffer, *self.conv_buffer] - - def _iter_page_tensors(self, index: int): - if self.layout in ["page_first", "page_first_direct"]: - yield self.temporal_buffer[index] - for conv_buf in self.conv_buffer: - yield conv_buf[index] - else: - yield self.temporal_buffer[:, index : index + self.page_size] - for conv_buf in self.conv_buffer: - yield conv_buf[:, index : index + self.page_size] - - @staticmethod - def _flatten_tensor_bytes(tensor: torch.Tensor) -> torch.Tensor: - return tensor.contiguous().view(torch.uint8).reshape(-1) - - @synchronized - def clear(self): - self.mem_state = torch.zeros( - (self.size,), dtype=torch.uint8, device=self.device - ) - self.free_slots = torch.arange(self.size, dtype=torch.int64) - self.release_slots = [] - self.num_release_slots = 0 - - def available_size(self): - return len(self.free_slots) + self.num_release_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 - - if need_size > len(self.free_slots): - self._merge_release_slots() - - 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: - indices_cpu = indices.cpu() - if indices_cpu.numel() == 0: - return 0 - - self.release_slots.append(indices_cpu) - self.num_release_slots += len(indices_cpu) - return len(indices) - - def get_size_per_token(self): - conv_total_size = sum( - conv_elem_size * self.conv_dtype.itemsize - for conv_elem_size in self.conv_state_elem_sizes - ) - temporal_size = self.temporal_state_elem_size * self.temporal_dtype.itemsize - return (conv_total_size + temporal_size) * self.num_mamba_layers - - def get_ksize_per_token(self): - return self.get_size_per_token() - - @staticmethod - def _item_size_per_index(tensor: torch.Tensor) -> int: - if tensor.shape[0] == 0: - return 0 - return int(tensor[0].numel() * tensor.element_size()) - - @staticmethod - def _copy_tensor( - src: torch.Tensor, - dst: torch.Tensor, - src_indices: torch.Tensor, - dst_indices: torch.Tensor, - io_backend: str, - ) -> None: - if src_indices.numel() == 0: - return - if io_backend == "kernel": - # TODO: Rename the interface for clarity. - # Here, transfer_kv_per_layer_mla is reused to transfer the Mamba state. - # This has nothing to do with MLA; it's only reused because this interface happens to transfer a single Pool. - transfer_kv_per_layer_mla( - src=src, - dst=dst, - src_indices=src_indices, - dst_indices=dst_indices, - item_size=MambaPoolHost._item_size_per_index(src), - ) - elif io_backend == "direct": - transfer_kv_direct( - src_layers=[src], - dst_layers=[dst], - src_indices=src_indices, - dst_indices=dst_indices, - page_size=1, - ) - else: - raise ValueError(f"Unsupported io_backend: {io_backend}") - - @staticmethod - def _copy_tensor_pf_lf( - src: torch.Tensor, - dst: torch.Tensor, - src_indices: torch.Tensor, - dst_indices: torch.Tensor, - layer_id: int, - num_layers: int, - io_backend: str, - ) -> None: - if src_indices.numel() == 0: - return - if io_backend == "kernel": - item_size = MambaPoolHost._item_size_per_index(dst) - # Mamba JIT kernel expects all index tensors on CUDA. - # host_indices may be on CPU (kept there by start_writing when - # can_use_write_back_jit is True on the HostPoolGroup). - if src_indices.device.type != "cuda": - src_indices = src_indices.to(dst_indices.device, non_blocking=True) - transfer_kv_mamba_pf_lf( - src=src, - dst=dst, - src_indices=src_indices, - dst_indices=dst_indices, - layer_id=layer_id, - item_size=item_size, - src_layout_dim=item_size * num_layers, - ) - elif io_backend == "direct": - transfer_kv_per_layer_direct_pf_lf( - src_ptrs=[src], - dst_ptrs=[dst], - src_indices=src_indices, - dst_indices=dst_indices, - layer_id=layer_id, - page_size=1, - ) - else: - raise ValueError(f"Unsupported io_backend: {io_backend}") - - @staticmethod - def _copy_tensor_all_layers_lf_pf( - src_layers: torch.Tensor, - dst: torch.Tensor, - src_indices: torch.Tensor, - dst_indices: torch.Tensor, - num_layers: int, - io_backend: str, - src_ptrs: torch.Tensor, - staging: Optional[torch.Tensor] = None, - can_use_jit: bool = False, - ) -> None: - if src_indices.numel() == 0: - return - if io_backend == "kernel": - item_size = MambaPoolHost._item_size_per_index(src_layers[0]) - # Mamba JIT kernel expects all index tensors on CUDA. - # When can_use_write_back_jit is True on the HostPoolGroup, - # start_writing() keeps host_indices on CPU (for MLA staged kernel). - # Move dst_indices to CUDA here to satisfy the kernel's requirement. - if dst_indices.device.type != "cuda": - dst_indices = dst_indices.to(src_indices.device, non_blocking=True) - transfer_kv_mamba_lf_pf( - src_ptrs=src_ptrs, - dst=dst, - src_indices=src_indices, - dst_indices=dst_indices, - item_size=item_size, - dst_layout_dim=item_size * num_layers, - num_layers=num_layers, - ) - elif io_backend == "direct": - src_ptrs = [src_layers[i] for i in range(num_layers)] - transfer_kv_all_layer_direct_lf_pf( - src_ptrs=src_ptrs, - dst_ptrs=[dst], - src_indices=src_indices, - dst_indices=dst_indices, - page_size=1, - ) - else: - raise ValueError(f"Unsupported io_backend: {io_backend}") - - def load_to_device_per_layer( - self, - device_pool, - host_indices, - device_indices, - layer_id, - io_backend="kernel", - *, - is_draft: bool = False, - ): - if self.layout in ["page_first", "page_first_direct"]: - # no ssm state on conv-only models: nothing to transfer - if self.temporal_state_elem_size > 0: - self._copy_tensor_pf_lf( - src=self.temporal_buffer, - dst=device_pool.mamba_cache.temporal[layer_id], - src_indices=host_indices, - dst_indices=device_indices, - layer_id=layer_id, - num_layers=self.num_mamba_layers, - io_backend=io_backend, - ) - for conv_idx in range(len(self.conv_state_shapes)): - self._copy_tensor_pf_lf( - src=self.conv_buffer[conv_idx], - dst=device_pool.mamba_cache.conv[conv_idx][layer_id], - src_indices=host_indices, - dst_indices=device_indices, - layer_id=layer_id, - num_layers=self.num_mamba_layers, - io_backend=io_backend, - ) - else: - self._copy_tensor( - self.temporal_buffer[layer_id], - device_pool.mamba_cache.temporal[layer_id], - host_indices, - device_indices, - io_backend, - ) - for conv_idx in range(len(self.conv_state_shapes)): - self._copy_tensor( - self.conv_buffer[conv_idx][layer_id], - device_pool.mamba_cache.conv[conv_idx][layer_id], - host_indices, - device_indices, - io_backend, - ) - - def backup_from_device_all_layer( - self, device_pool, host_indices, device_indices, io_backend="kernel" - ): - if self.layout in ["page_first", "page_first_direct"]: - # no ssm state on conv-only models: a 0-size batched memcpy errors - if self.temporal_state_elem_size > 0: - self._copy_tensor_all_layers_lf_pf( - src_layers=device_pool.mamba_cache.temporal, - dst=self.temporal_buffer, - src_indices=device_indices, - dst_indices=host_indices, - num_layers=self.num_mamba_layers, - io_backend=io_backend, - staging=self.temporal_staging_buffer, - can_use_jit=self._temporal_can_use_jit, - src_ptrs=self.temporal_device_ptrs, - ) - for conv_idx in range(len(self.conv_state_shapes)): - self._copy_tensor_all_layers_lf_pf( - src_layers=device_pool.mamba_cache.conv[conv_idx], - dst=self.conv_buffer[conv_idx], - src_indices=device_indices, - dst_indices=host_indices, - num_layers=self.num_mamba_layers, - io_backend=io_backend, - staging=self.conv_staging_buffers[conv_idx], - can_use_jit=self._conv_can_use_jit[conv_idx], - src_ptrs=self.conv_device_ptrs[conv_idx], - ) - else: - for layer_id in range(self.num_mamba_layers): - self._copy_tensor( - device_pool.mamba_cache.temporal[layer_id], - self.temporal_buffer[layer_id], - device_indices, - host_indices, - io_backend, - ) - for conv_idx in range(len(self.conv_state_shapes)): - self._copy_tensor( - device_pool.mamba_cache.conv[conv_idx][layer_id], - self.conv_buffer[conv_idx][layer_id], - device_indices, - host_indices, - io_backend, - ) - - def get_data_page(self, index, flat: bool = True) -> torch.Tensor: - data_page = torch.cat( - [ - self._flatten_tensor_bytes(tensor) - for tensor in self._iter_page_tensors(index) - ] - ) - return data_page.flatten() if flat else data_page - - def get_dummy_flat_data_page(self) -> torch.Tensor: - return torch.zeros( - self.page_size * self.size_per_token, - dtype=torch.uint8, - device=self.device, - pin_memory=self.pin_memory, - ) - - def set_from_flat_data_page( - self, - index: int, - data_page: torch.Tensor, - ) -> None: - flat_bytes = data_page.contiguous().view(torch.uint8).reshape(-1) - start = 0 - for tensor in self._iter_page_tensors(index): - num_bytes = tensor.numel() * tensor.element_size() - tensor_bytes = flat_bytes[start : start + num_bytes] - start += num_bytes - restored = tensor_bytes.view(dtype=tensor.dtype).reshape(tensor.shape) - tensor.copy_(restored) - - def get_page_buffer_meta(self, indices): - """Meta data for zero-copy storage I/O. - - Only page-first layouts are supported for mamba storage zero-copy because - each page slot in temporal/conv buffers is directly addressable. - """ - assert len(indices) % self.page_size == 0 - if self.layout not in ["page_first", "page_first_direct"]: - raise ValueError( - f"Mamba storage zero-copy requires page_first layout, got {self.layout}" - ) - indices = indices.tolist() - ptr_list = [] - element_size_list = [] - - # Compute base pointers once; each page pointer is offset from these bases. - temporal_base_ptr = self.temporal_buffer.data_ptr() - conv_base_ptrs = [buf.data_ptr() for buf in self.conv_buffer] - # Component sizes are constant across pages, so precompute once as well. - temporal_element_size = ( - self.page_size - * self.num_mamba_layers - * self.temporal_dtype.itemsize - * self.temporal_state_elem_size - ) - conv_element_sizes = [ - ( - self.page_size - * self.num_mamba_layers - * self.conv_dtype.itemsize - * self.conv_state_elem_sizes[i] - ) - for i in range(len(self.conv_state_shapes)) - ] - - for i in range(0, len(indices), self.page_size): - # Emit component pointers in stable order: temporal first (dropped - # for conv-only models with no ssm state), then conv_0..conv_n. - # _get_hybrid_page_component_keys drops the temporal key under the - # same condition, keeping keys and buffers aligned. - if self.temporal_state_elem_size > 0: - temporal_ptr = ( - temporal_base_ptr - + indices[i] - * self.num_mamba_layers - * self.temporal_state_elem_size - * self.temporal_dtype.itemsize - ) - ptr_list.append(temporal_ptr) - element_size_list.append(temporal_element_size) - for j in range(len(self.conv_buffer)): - conv_ptr = ( - conv_base_ptrs[j] - + indices[i] - * self.num_mamba_layers - * self.conv_state_elem_sizes[j] - * self.conv_dtype.itemsize - ) - ptr_list.append(conv_ptr) - element_size_list.append(conv_element_sizes[j]) - return ptr_list, element_size_list - - def is_stride_page_aligned(self, page_size_bytes: int = 4096) -> bool: - if self.layout not in ["page_first", "page_first_direct"]: - return False - temporal_stride = ( - self.num_mamba_layers - * self.temporal_state_elem_size - * self.temporal_dtype.itemsize - ) - if self.temporal_buffer.data_ptr() % page_size_bytes != 0: - return False - if temporal_stride % page_size_bytes != 0: - return False - for buf, elem_size in zip(self.conv_buffer, self.conv_state_elem_sizes): - conv_stride = self.num_mamba_layers * elem_size * self.conv_dtype.itemsize - if buf.data_ptr() % page_size_bytes != 0: - return False - if conv_stride % page_size_bytes != 0: - return False - return True - - # ---- V4 Compressed KV Host Pools ---- diff --git a/python/sglang/srt/mem_cache/pool_host/mamba.py b/python/sglang/srt/mem_cache/pool_host/mamba.py new file mode 100644 index 000000000..22d1f8b17 --- /dev/null +++ b/python/sglang/srt/mem_cache/pool_host/mamba.py @@ -0,0 +1,610 @@ +from __future__ import annotations + +import logging +import threading +from typing import Optional + +import numpy as np +import psutil +import torch + +from sglang.srt.mem_cache.memory_pool import MambaPool +from sglang.srt.mem_cache.pool_host.base import ( + HICACHE_HOST_MEMORY_RESERVE_BYTES, + HostKVCache, + sync_fixed_hicache_size, + synchronized, +) +from sglang.srt.mem_cache.pool_host.common import ( + ALLOC_MEMORY_FUNCS, + get_allocator_from_storage, +) +from sglang.srt.utils import is_cuda, is_hip + +_is_cuda = is_cuda() +_is_hip = is_hip() +if _is_cuda or _is_hip: + from sgl_kernel.kvcacheio import ( + transfer_kv_all_layer_direct_lf_pf, + transfer_kv_direct, + transfer_kv_per_layer_direct_pf_lf, + transfer_kv_per_layer_mla, + ) +if _is_cuda or _is_hip: + from sglang.kernels.ops.mamba.transfer_mamba import ( + transfer_kv_mamba_lf_pf, + transfer_kv_mamba_pf_lf, + ) + +logger = logging.getLogger(__name__) + + +class MambaPoolHost(HostKVCache): + def __init__( + self, + device_pool: MambaPool, + host_to_device_ratio: float, + host_size: int, + pin_memory: bool = True, + device: str = "cpu", + allocator_type: str = "default", + layout: str = "layer_first", + ): + self.device_pool = device_pool + self.page_size = 1 + + assert layout in [ + "page_first", + "page_first_direct", + ], f"Unsupported layout: {layout}" + + self.layout = layout + self.pin_memory = pin_memory + self.device = device + self.allocator = get_allocator_from_storage(allocator_type) + self.num_mamba_layers = device_pool.num_mamba_layers + + self.conv_state_shapes = [ + conv_state.shape[2:] for conv_state in device_pool.mamba_cache.conv + ] + self.temporal_state_shape = device_pool.mamba_cache.temporal.shape[2:] + self.temporal_state_elem_size = int(np.prod(self.temporal_state_shape)) + self.conv_state_elem_sizes = [ + int(np.prod(conv_shape)) for conv_shape in self.conv_state_shapes + ] + self.conv_dtype = device_pool.mamba_cache.conv[0].dtype + self.temporal_dtype = device_pool.mamba_cache.temporal.dtype + self.dtype = self.conv_dtype + self.size_per_token = self.get_size_per_token() + + if host_size > 0: + self.size = sync_fixed_hicache_size( + int(host_size * 1e9 // self.size_per_token), host_size + ) + else: + self.size = int(device_pool.size * host_to_device_ratio) + + self.page_num = self.size // self.page_size + 1 + self.size = self.page_num * self.page_size + + if self.size <= device_pool.size: + logger.warning( + "HiCache host KV pool (%d tokens) is smaller than the device pool (%d tokens);" + "L2 cache effectiveness is reduced." + "Consider increasing --hicache-ratio (or --hicache-size) for higher L2 cache hit rate.", + self.size, + device_pool.size, + ) + + 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." + ) + logger.info( + "Allocating %.2f GB host memory for hierarchical Mamba cache (layout=%s).", + requested_bytes / 1e9, + self.layout, + ) + + self.temporal_device_ptrs = torch.tensor( + [ + device_pool.mamba_cache.temporal[i].data_ptr() + for i in range(self.num_mamba_layers) + ], + dtype=torch.uint64, + device=self.device_pool.device, + ) + self.conv_device_ptrs = [ + torch.tensor( + [conv_state[i].data_ptr() for i in range(self.num_mamba_layers)], + dtype=torch.uint64, + device=self.device_pool.device, + ) + for conv_state in device_pool.mamba_cache.conv + ] + + self.init_kv_buffer() + self._init_write_back_staging_buffers() + self.lock = threading.RLock() + self.clear() + + def init_kv_buffer(self): + _host_alloc = ALLOC_MEMORY_FUNCS[self.device_pool.device] + + def alloc_func(dims, *, dtype, device, pin_memory, allocator): + # conv-only linear attention has no ssm state: mmap can't map the + # 0-element temporal buffer, so hand back a plain empty tensor. + if np.prod(dims) == 0: + return torch.empty(dims, dtype=dtype, device=device) + return _host_alloc( + dims, + dtype=dtype, + device=device, + pin_memory=pin_memory, + allocator=allocator, + ) + + if self.layout in ["page_first", "page_first_direct"]: + # page-first: (page_num, num_layers, 1, *shape) — per-page data is contiguous + temporal_dims = ( + self.size, + self.num_mamba_layers, + 1, + ) + self.temporal_state_shape + self.temporal_buffer = alloc_func( + temporal_dims, + dtype=self.temporal_dtype, + device=self.device, + pin_memory=self.pin_memory, + allocator=self.allocator, + ) + self.conv_buffer = [] + for conv_shape in self.conv_state_shapes: + conv_dims = (self.size, self.num_mamba_layers, 1) + conv_shape + self.conv_buffer.append( + alloc_func( + conv_dims, + dtype=self.conv_dtype, + device=self.device, + pin_memory=self.pin_memory, + allocator=self.allocator, + ) + ) + else: + # layer-first: (num_layers, size, *shape) + temporal_dims = ( + self.num_mamba_layers, + self.size, + ) + self.temporal_state_shape + self.temporal_buffer = alloc_func( + temporal_dims, + dtype=self.temporal_dtype, + device=self.device, + pin_memory=self.pin_memory, + allocator=self.allocator, + ) + self.conv_buffer = [] + for conv_shape in self.conv_state_shapes: + conv_dims = (self.num_mamba_layers, self.size) + conv_shape + self.conv_buffer.append( + alloc_func( + conv_dims, + dtype=self.conv_dtype, + device=self.device, + pin_memory=self.pin_memory, + allocator=self.allocator, + ) + ) + + def _init_write_back_staging_buffers(self): + self.temporal_staging_buffer = None + self.conv_staging_buffers = [None] * len(self.conv_buffer) + # Must be True: HostPoolGroup computes can_use_write_back_jit as AND of + # all pools. When True, start_writing() keeps indices on CPU, which MLA's + # staged write-back kernel requires. MambaPoolHost's own backup path does + # not check this flag — it routes by layout + io_backend instead. + self.can_use_write_back_jit = True + self._temporal_can_use_jit = False + self._conv_can_use_jit = [False] * len(self.conv_buffer) + + def get_hybrid_pool_buffer(self): + # Expose all mamba host tensors that need Mooncake buffer registration. + return [self.temporal_buffer, *self.conv_buffer] + + def _iter_page_tensors(self, index: int): + if self.layout in ["page_first", "page_first_direct"]: + yield self.temporal_buffer[index] + for conv_buf in self.conv_buffer: + yield conv_buf[index] + else: + yield self.temporal_buffer[:, index : index + self.page_size] + for conv_buf in self.conv_buffer: + yield conv_buf[:, index : index + self.page_size] + + @staticmethod + def _flatten_tensor_bytes(tensor: torch.Tensor) -> torch.Tensor: + return tensor.contiguous().view(torch.uint8).reshape(-1) + + @synchronized + def clear(self): + self.mem_state = torch.zeros( + (self.size,), dtype=torch.uint8, device=self.device + ) + self.free_slots = torch.arange(self.size, dtype=torch.int64) + self.release_slots = [] + self.num_release_slots = 0 + + def available_size(self): + return len(self.free_slots) + self.num_release_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 + + if need_size > len(self.free_slots): + self._merge_release_slots() + + 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: + indices_cpu = indices.cpu() + if indices_cpu.numel() == 0: + return 0 + + self.release_slots.append(indices_cpu) + self.num_release_slots += len(indices_cpu) + return len(indices) + + def get_size_per_token(self): + conv_total_size = sum( + conv_elem_size * self.conv_dtype.itemsize + for conv_elem_size in self.conv_state_elem_sizes + ) + temporal_size = self.temporal_state_elem_size * self.temporal_dtype.itemsize + return (conv_total_size + temporal_size) * self.num_mamba_layers + + def get_ksize_per_token(self): + return self.get_size_per_token() + + @staticmethod + def _item_size_per_index(tensor: torch.Tensor) -> int: + if tensor.shape[0] == 0: + return 0 + return int(tensor[0].numel() * tensor.element_size()) + + @staticmethod + def _copy_tensor( + src: torch.Tensor, + dst: torch.Tensor, + src_indices: torch.Tensor, + dst_indices: torch.Tensor, + io_backend: str, + ) -> None: + if src_indices.numel() == 0: + return + if io_backend == "kernel": + # TODO: Rename the interface for clarity. + # Here, transfer_kv_per_layer_mla is reused to transfer the Mamba state. + # This has nothing to do with MLA; it's only reused because this interface happens to transfer a single Pool. + transfer_kv_per_layer_mla( + src=src, + dst=dst, + src_indices=src_indices, + dst_indices=dst_indices, + item_size=MambaPoolHost._item_size_per_index(src), + ) + elif io_backend == "direct": + transfer_kv_direct( + src_layers=[src], + dst_layers=[dst], + src_indices=src_indices, + dst_indices=dst_indices, + page_size=1, + ) + else: + raise ValueError(f"Unsupported io_backend: {io_backend}") + + @staticmethod + def _copy_tensor_pf_lf( + src: torch.Tensor, + dst: torch.Tensor, + src_indices: torch.Tensor, + dst_indices: torch.Tensor, + layer_id: int, + num_layers: int, + io_backend: str, + ) -> None: + if src_indices.numel() == 0: + return + if io_backend == "kernel": + item_size = MambaPoolHost._item_size_per_index(dst) + # Mamba JIT kernel expects all index tensors on CUDA. + # host_indices may be on CPU (kept there by start_writing when + # can_use_write_back_jit is True on the HostPoolGroup). + if src_indices.device.type != "cuda": + src_indices = src_indices.to(dst_indices.device, non_blocking=True) + transfer_kv_mamba_pf_lf( + src=src, + dst=dst, + src_indices=src_indices, + dst_indices=dst_indices, + layer_id=layer_id, + item_size=item_size, + src_layout_dim=item_size * num_layers, + ) + elif io_backend == "direct": + transfer_kv_per_layer_direct_pf_lf( + src_ptrs=[src], + dst_ptrs=[dst], + src_indices=src_indices, + dst_indices=dst_indices, + layer_id=layer_id, + page_size=1, + ) + else: + raise ValueError(f"Unsupported io_backend: {io_backend}") + + @staticmethod + def _copy_tensor_all_layers_lf_pf( + src_layers: torch.Tensor, + dst: torch.Tensor, + src_indices: torch.Tensor, + dst_indices: torch.Tensor, + num_layers: int, + io_backend: str, + src_ptrs: torch.Tensor, + staging: Optional[torch.Tensor] = None, + can_use_jit: bool = False, + ) -> None: + if src_indices.numel() == 0: + return + if io_backend == "kernel": + item_size = MambaPoolHost._item_size_per_index(src_layers[0]) + # Mamba JIT kernel expects all index tensors on CUDA. + # When can_use_write_back_jit is True on the HostPoolGroup, + # start_writing() keeps host_indices on CPU (for MLA staged kernel). + # Move dst_indices to CUDA here to satisfy the kernel's requirement. + if dst_indices.device.type != "cuda": + dst_indices = dst_indices.to(src_indices.device, non_blocking=True) + transfer_kv_mamba_lf_pf( + src_ptrs=src_ptrs, + dst=dst, + src_indices=src_indices, + dst_indices=dst_indices, + item_size=item_size, + dst_layout_dim=item_size * num_layers, + num_layers=num_layers, + ) + elif io_backend == "direct": + src_ptrs = [src_layers[i] for i in range(num_layers)] + transfer_kv_all_layer_direct_lf_pf( + src_ptrs=src_ptrs, + dst_ptrs=[dst], + src_indices=src_indices, + dst_indices=dst_indices, + page_size=1, + ) + else: + raise ValueError(f"Unsupported io_backend: {io_backend}") + + def load_to_device_per_layer( + self, + device_pool, + host_indices, + device_indices, + layer_id, + io_backend="kernel", + *, + is_draft: bool = False, + ): + if self.layout in ["page_first", "page_first_direct"]: + # no ssm state on conv-only models: nothing to transfer + if self.temporal_state_elem_size > 0: + self._copy_tensor_pf_lf( + src=self.temporal_buffer, + dst=device_pool.mamba_cache.temporal[layer_id], + src_indices=host_indices, + dst_indices=device_indices, + layer_id=layer_id, + num_layers=self.num_mamba_layers, + io_backend=io_backend, + ) + for conv_idx in range(len(self.conv_state_shapes)): + self._copy_tensor_pf_lf( + src=self.conv_buffer[conv_idx], + dst=device_pool.mamba_cache.conv[conv_idx][layer_id], + src_indices=host_indices, + dst_indices=device_indices, + layer_id=layer_id, + num_layers=self.num_mamba_layers, + io_backend=io_backend, + ) + else: + self._copy_tensor( + self.temporal_buffer[layer_id], + device_pool.mamba_cache.temporal[layer_id], + host_indices, + device_indices, + io_backend, + ) + for conv_idx in range(len(self.conv_state_shapes)): + self._copy_tensor( + self.conv_buffer[conv_idx][layer_id], + device_pool.mamba_cache.conv[conv_idx][layer_id], + host_indices, + device_indices, + io_backend, + ) + + def backup_from_device_all_layer( + self, device_pool, host_indices, device_indices, io_backend="kernel" + ): + if self.layout in ["page_first", "page_first_direct"]: + # no ssm state on conv-only models: a 0-size batched memcpy errors + if self.temporal_state_elem_size > 0: + self._copy_tensor_all_layers_lf_pf( + src_layers=device_pool.mamba_cache.temporal, + dst=self.temporal_buffer, + src_indices=device_indices, + dst_indices=host_indices, + num_layers=self.num_mamba_layers, + io_backend=io_backend, + staging=self.temporal_staging_buffer, + can_use_jit=self._temporal_can_use_jit, + src_ptrs=self.temporal_device_ptrs, + ) + for conv_idx in range(len(self.conv_state_shapes)): + self._copy_tensor_all_layers_lf_pf( + src_layers=device_pool.mamba_cache.conv[conv_idx], + dst=self.conv_buffer[conv_idx], + src_indices=device_indices, + dst_indices=host_indices, + num_layers=self.num_mamba_layers, + io_backend=io_backend, + staging=self.conv_staging_buffers[conv_idx], + can_use_jit=self._conv_can_use_jit[conv_idx], + src_ptrs=self.conv_device_ptrs[conv_idx], + ) + else: + for layer_id in range(self.num_mamba_layers): + self._copy_tensor( + device_pool.mamba_cache.temporal[layer_id], + self.temporal_buffer[layer_id], + device_indices, + host_indices, + io_backend, + ) + for conv_idx in range(len(self.conv_state_shapes)): + self._copy_tensor( + device_pool.mamba_cache.conv[conv_idx][layer_id], + self.conv_buffer[conv_idx][layer_id], + device_indices, + host_indices, + io_backend, + ) + + def get_data_page(self, index, flat: bool = True) -> torch.Tensor: + data_page = torch.cat( + [ + self._flatten_tensor_bytes(tensor) + for tensor in self._iter_page_tensors(index) + ] + ) + return data_page.flatten() if flat else data_page + + def get_dummy_flat_data_page(self) -> torch.Tensor: + return torch.zeros( + self.page_size * self.size_per_token, + dtype=torch.uint8, + device=self.device, + pin_memory=self.pin_memory, + ) + + def set_from_flat_data_page( + self, + index: int, + data_page: torch.Tensor, + ) -> None: + flat_bytes = data_page.contiguous().view(torch.uint8).reshape(-1) + start = 0 + for tensor in self._iter_page_tensors(index): + num_bytes = tensor.numel() * tensor.element_size() + tensor_bytes = flat_bytes[start : start + num_bytes] + start += num_bytes + restored = tensor_bytes.view(dtype=tensor.dtype).reshape(tensor.shape) + tensor.copy_(restored) + + def get_page_buffer_meta(self, indices): + """Meta data for zero-copy storage I/O. + + Only page-first layouts are supported for mamba storage zero-copy because + each page slot in temporal/conv buffers is directly addressable. + """ + assert len(indices) % self.page_size == 0 + if self.layout not in ["page_first", "page_first_direct"]: + raise ValueError( + f"Mamba storage zero-copy requires page_first layout, got {self.layout}" + ) + indices = indices.tolist() + ptr_list = [] + element_size_list = [] + + # Compute base pointers once; each page pointer is offset from these bases. + temporal_base_ptr = self.temporal_buffer.data_ptr() + conv_base_ptrs = [buf.data_ptr() for buf in self.conv_buffer] + # Component sizes are constant across pages, so precompute once as well. + temporal_element_size = ( + self.page_size + * self.num_mamba_layers + * self.temporal_dtype.itemsize + * self.temporal_state_elem_size + ) + conv_element_sizes = [ + ( + self.page_size + * self.num_mamba_layers + * self.conv_dtype.itemsize + * self.conv_state_elem_sizes[i] + ) + for i in range(len(self.conv_state_shapes)) + ] + + for i in range(0, len(indices), self.page_size): + # Emit component pointers in stable order: temporal first (dropped + # for conv-only models with no ssm state), then conv_0..conv_n. + # _get_hybrid_page_component_keys drops the temporal key under the + # same condition, keeping keys and buffers aligned. + if self.temporal_state_elem_size > 0: + temporal_ptr = ( + temporal_base_ptr + + indices[i] + * self.num_mamba_layers + * self.temporal_state_elem_size + * self.temporal_dtype.itemsize + ) + ptr_list.append(temporal_ptr) + element_size_list.append(temporal_element_size) + for j in range(len(self.conv_buffer)): + conv_ptr = ( + conv_base_ptrs[j] + + indices[i] + * self.num_mamba_layers + * self.conv_state_elem_sizes[j] + * self.conv_dtype.itemsize + ) + ptr_list.append(conv_ptr) + element_size_list.append(conv_element_sizes[j]) + return ptr_list, element_size_list + + def is_stride_page_aligned(self, page_size_bytes: int = 4096) -> bool: + if self.layout not in ["page_first", "page_first_direct"]: + return False + temporal_stride = ( + self.num_mamba_layers + * self.temporal_state_elem_size + * self.temporal_dtype.itemsize + ) + if self.temporal_buffer.data_ptr() % page_size_bytes != 0: + return False + if temporal_stride % page_size_bytes != 0: + return False + for buf, elem_size in zip(self.conv_buffer, self.conv_state_elem_sizes): + conv_stride = self.num_mamba_layers * elem_size * self.conv_dtype.itemsize + if buf.data_ptr() % page_size_bytes != 0: + return False + if conv_stride % page_size_bytes != 0: + return False + return True diff --git a/test/registered/kernels/ops/mamba/test_transfer_mamba.py b/test/registered/kernels/ops/mamba/test_transfer_mamba.py index 8f7644c23..ff552f30f 100644 --- a/test/registered/kernels/ops/mamba/test_transfer_mamba.py +++ b/test/registered/kernels/ops/mamba/test_transfer_mamba.py @@ -12,7 +12,7 @@ from types import SimpleNamespace import pytest import torch -from sglang.srt.mem_cache.memory_pool_host import MambaPoolHost +from sglang.srt.mem_cache.pool_host.mamba import MambaPoolHost from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci register_cuda_ci(est_time=10, stage="base-b-kernel-unit", runner_config="1-gpu-large") diff --git a/test/registered/unit/mem_cache/test_hicache_staged_write_back_dispatch.py b/test/registered/unit/mem_cache/test_hicache_staged_write_back_dispatch.py index fc07cfb81..8dc1542a3 100644 --- a/test/registered/unit/mem_cache/test_hicache_staged_write_back_dispatch.py +++ b/test/registered/unit/mem_cache/test_hicache_staged_write_back_dispatch.py @@ -24,9 +24,9 @@ from sglang.srt.mem_cache.memory_pool_host import ( DSAIndexerPoolHost, HostPoolGroup, LogicalHostPool, - MambaPoolHost, PoolEntry, ) +from sglang.srt.mem_cache.pool_host.mamba import MambaPoolHost from sglang.srt.mem_cache.pool_host.mha import MHATokenToKVPoolHost from sglang.srt.mem_cache.pool_host.mla import MLATokenToKVPoolHost from sglang.test.ci.ci_register import register_cpu_ci diff --git a/test/registered/unit/mem_cache/test_mem_pool_host.py b/test/registered/unit/mem_cache/test_mem_pool_host.py index dae2acb2f..02dc0e531 100644 --- a/test/registered/unit/mem_cache/test_mem_pool_host.py +++ b/test/registered/unit/mem_cache/test_mem_pool_host.py @@ -9,8 +9,8 @@ from sglang.srt.mem_cache.memory_pool import MHATokenToKVPool from sglang.srt.mem_cache.memory_pool_host import ( DeepSeekV4PagedHostPool, LogicalHostPool, - MambaPoolHost, ) +from sglang.srt.mem_cache.pool_host.mamba import MambaPoolHost from sglang.srt.mem_cache.pool_host.mha import MHATokenToKVPoolHost from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.test_utils import CustomTestCase