[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,
|
DSAIndexerPoolHost,
|
||||||
HostPoolGroup,
|
HostPoolGroup,
|
||||||
LogicalHostPool,
|
LogicalHostPool,
|
||||||
MambaPoolHost,
|
|
||||||
PoolEntry,
|
PoolEntry,
|
||||||
)
|
)
|
||||||
from sglang.srt.mem_cache.pool_host.common import get_allocator_type
|
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 (
|
from sglang.srt.mem_cache.pool_host.mha import (
|
||||||
MHATokenToKOnlyPoolHost,
|
MHATokenToKOnlyPoolHost,
|
||||||
get_mha_host_pool_cls,
|
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.hicache_storage import PoolName
|
||||||
from sglang.srt.mem_cache.pool_host.mla import MLATokenToKVPoolHost
|
from sglang.srt.mem_cache.pool_host.mla import MLATokenToKVPoolHost
|
||||||
|
|
||||||
import numpy as np
|
|
||||||
import psutil
|
import psutil
|
||||||
import torch
|
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,
|
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.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
|
from sglang.srt.utils import is_cuda, is_hip, is_mps, is_npu, is_xpu
|
||||||
|
|
||||||
_is_cuda = is_cuda()
|
_is_cuda = is_cuda()
|
||||||
@@ -38,13 +39,6 @@ if _is_cuda or _is_hip:
|
|||||||
transfer_kv_per_layer_mla,
|
transfer_kv_per_layer_mla,
|
||||||
transfer_kv_per_layer_mla_pf_lf,
|
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__)
|
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 (
|
from sglang.srt.mem_cache.pool_host.base import (
|
||||||
_WRITE_BACK_STAGING_PAGE_CHUNK,
|
_WRITE_BACK_STAGING_PAGE_CHUNK,
|
||||||
HICACHE_HOST_MEMORY_RESERVE_BYTES,
|
HICACHE_HOST_MEMORY_RESERVE_BYTES,
|
||||||
sync_fixed_hicache_size,
|
|
||||||
synchronized,
|
synchronized,
|
||||||
)
|
)
|
||||||
from sglang.srt.mem_cache.pool_host.common import (
|
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
|
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 ----
|
# ---- 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 pytest
|
||||||
import torch
|
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
|
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")
|
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,
|
DSAIndexerPoolHost,
|
||||||
HostPoolGroup,
|
HostPoolGroup,
|
||||||
LogicalHostPool,
|
LogicalHostPool,
|
||||||
MambaPoolHost,
|
|
||||||
PoolEntry,
|
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.mha import MHATokenToKVPoolHost
|
||||||
from sglang.srt.mem_cache.pool_host.mla import MLATokenToKVPoolHost
|
from sglang.srt.mem_cache.pool_host.mla import MLATokenToKVPoolHost
|
||||||
from sglang.test.ci.ci_register import register_cpu_ci
|
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 (
|
from sglang.srt.mem_cache.memory_pool_host import (
|
||||||
DeepSeekV4PagedHostPool,
|
DeepSeekV4PagedHostPool,
|
||||||
LogicalHostPool,
|
LogicalHostPool,
|
||||||
MambaPoolHost,
|
|
||||||
)
|
)
|
||||||
|
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.mha import MHATokenToKVPoolHost
|
||||||
from sglang.test.ci.ci_register import register_cpu_ci
|
from sglang.test.ci.ci_register import register_cpu_ci
|
||||||
from sglang.test.test_utils import CustomTestCase
|
from sglang.test.test_utils import CustomTestCase
|
||||||
|
|||||||
Reference in New Issue
Block a user