[mem_cache][8/N] refactor: move MambaPoolHost to pool_host.mamba (#31180)
This commit is contained in:
Shuwen Wang
@@ -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,
|
||||
|
||||
@@ -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 ----
|
||||
|
||||
|
||||
|
||||
@@ -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
|
||||
@@ -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")
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user