minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[mem_cache][7/N] refactor: move MLATokenToKVPoolHost to pool_host.mla (#30616)

Browse Source
This commit is contained in:
shuwenn
2026-07-13 14:23:38 +08:00
committed by GitHub
parent cbcbef6811
commit 9dd57ef8c4
13 changed files with 602 additions and 564 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/disaggregation/decode_kvcache_offload_manager.py
+1 -1
View File
@@ -18,8 +18,8 @@ from sglang.srt.mem_cache.memory_pool import (
MLATokenToKVPool,
ReqToTokenPool,
)
from sglang.srt.mem_cache.memory_pool_host import MLATokenToKVPoolHost
from sglang.srt.mem_cache.pool_host.mha import get_mha_host_pool_cls
from sglang.srt.mem_cache.pool_host.mla import MLATokenToKVPoolHost
from sglang.srt.server_args import ServerArgs
from sglang.srt.utils.common import ceil_align
python/sglang/srt/managers/hisparse_coordinator.py
+2 -4
View File
@@ -18,10 +18,8 @@ from sglang.srt.mem_cache.hisparse_memory_pool import (
HiSparseDSATokenToKVPool,
)
from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
from sglang.srt.mem_cache.memory_pool_host import (
DeepSeekV4PagedHostPool,
MLATokenToKVPoolHost,
)
from sglang.srt.mem_cache.memory_pool_host import DeepSeekV4PagedHostPool
from sglang.srt.mem_cache.pool_host.mla import MLATokenToKVPoolHost
from sglang.srt.utils import get_device_module, is_hip
device_module = get_device_module()
python/sglang/srt/mem_cache/hiradix_cache.py
+1 -1
View File
@@ -48,8 +48,8 @@ from sglang.srt.mem_cache.memory_pool import (
MiniMaxSparseKVPool,
MLATokenToKVPool,
)
from sglang.srt.mem_cache.memory_pool_host import MLATokenToKVPoolHost
from sglang.srt.mem_cache.pool_host.mha import get_mha_host_pool_cls
from sglang.srt.mem_cache.pool_host.mla import MLATokenToKVPoolHost
from sglang.srt.mem_cache.radix_cache import (
RadixCache,
RadixKey,
python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py
+1 -1
View File
@@ -19,13 +19,13 @@ from sglang.srt.mem_cache.memory_pool_host import (
HostPoolGroup,
LogicalHostPool,
MambaPoolHost,
MLATokenToKVPoolHost,
PoolEntry,
)
from sglang.srt.mem_cache.pool_host.mha import (
MHATokenToKOnlyPoolHost,
get_mha_host_pool_cls,
)
from sglang.srt.mem_cache.pool_host.mla import MLATokenToKVPoolHost
from sglang.srt.mem_cache.unified_cache_components import ComponentType
if TYPE_CHECKING:
python/sglang/srt/mem_cache/kv_cache_builder.py
+1 -1
View File
@@ -89,8 +89,8 @@ def maybe_register_hicache_draft(
MHATokenToKVPool,
MLATokenToKVPool,
)
from sglang.srt.mem_cache.memory_pool_host import MLATokenToKVPoolHost
from sglang.srt.mem_cache.pool_host.mha import get_mha_host_pool_cls
from sglang.srt.mem_cache.pool_host.mla import MLATokenToKVPoolHost
pool = draft_kv_pool
if isinstance(pool, HybridLinearKVPool):
python/sglang/srt/mem_cache/memory_pool_host.py
+1 -535
View File
@@ -7,29 +7,22 @@ from typing import TYPE_CHECKING, Any, Callable, Optional
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
from sglang.jit_kernel.hicache import (
can_use_hicache_jit_kernel,
can_use_write_back_jit_kernel,
)
from sglang.jit_kernel.hicache import (
transfer_hicache_all_layer_mla as jit_transfer_hicache_all_layer_mla,
)
from sglang.jit_kernel.hicache import (
transfer_hicache_all_layer_mla_staged_lf_pf as jit_transfer_hicache_all_layer_mla_staged_lf_pf,
)
from sglang.jit_kernel.hicache import (
transfer_hicache_one_layer_mla as jit_transfer_hicache_one_layer_mla,
)
from sglang.jit_kernel.hisparse import transfer_cache_dsv4_mla
from sglang.srt.mem_cache.memory_pool import (
DSATokenToKVPool,
MambaPool,
MLATokenToKVPool,
)
from sglang.srt.utils import is_cuda, is_hip, is_mps, is_npu, is_xpu
@@ -48,8 +41,6 @@ if _is_cuda or _is_hip:
transfer_kv_per_layer_mla,
transfer_kv_per_layer_mla_pf_lf,
)
if _is_npu:
from sgl_kernel_npu.kvcacheio import TransferDirection, transfer_kv_dim_exchange
logger = logging.getLogger(__name__)
@@ -68,531 +59,6 @@ from sglang.srt.mem_cache.pool_host.common import (
from sglang.srt.mem_cache.pool_host.hisparse import HiSparseHostPoolMixin
class MLATokenToKVPoolHost(HiSparseHostPoolMixin, HostKVCache):
device_pool: MLATokenToKVPool
def __init__(
self,
device_pool: MLATokenToKVPool,
host_to_device_ratio: float,
host_size: int,
page_size: int,
layout: str,
pin_memory: bool = True,
device: str = "cpu",
allocator_type: str = "default",
override_kv_cache_dim: Optional[int] = None,
):
self.override_kv_cache_dim = override_kv_cache_dim
super().__init__(
device_pool,
host_to_device_ratio,
host_size,
page_size,
layout,
pin_memory,
device,
allocator_type,
)
# The JIT HiCache kernels also build with hipcc (ROCm): the PTX-only
# helpers in hicache.cuh are guarded by USE_ROCM and the staged
# write-back kernel has a ROCm path, so enable them on HIP too. This
# keeps the ROCm write-back path consistent with CUDA.
self.can_use_jit = (_is_cuda or _is_hip) and can_use_hicache_jit_kernel(
element_size=self.kv_cache_dim * self.dtype.itemsize
)
if self.layout == "page_first":
# Transpose [page, layer, ...] -> [layer, page, ...] to get per-layer views
# This swaps strides without copying data
transposed = self.kv_buffer.transpose(0, 1)
self.data_refs = [transposed[i] for i in range(self.layer_num)]
else:
self.data_refs = [self.kv_buffer[i] for i in range(self.layer_num)]
self.data_ptrs = torch.tensor(
[x.data_ptr() for x in self.data_refs],
dtype=torch.uint64,
device=self.device_pool.device,
)
self._init_write_back_staging_buffers()
def get_contiguous_buf_infos(self):
"""Return (data_ptrs, data_lens, item_lens) in the same format as device pool,
for registering host memory with the disaggregation transfer engine."""
data_ptrs = [int(self.data_ptrs[i].item()) for i in range(self.layer_num)]
data_lens = [self.kv_buffer[i].nbytes for i in range(self.layer_num)]
item_lens = [self.token_stride_size * self.page_size] * self.layer_num
return data_ptrs, data_lens, item_lens
def get_size_per_token(self):
self.kv_lora_rank = self.device_pool.kv_lora_rank
self.qk_rope_head_dim = self.device_pool.qk_rope_head_dim
self.layer_num = self._effective_host_layer_num()
self.kv_cache_dim = self.override_kv_cache_dim or (
self.kv_lora_rank + self.qk_rope_head_dim
)
return self.kv_cache_dim * self.dtype.itemsize * self.layer_num
def get_ksize_per_token(self):
return self.get_size_per_token()
def init_kv_buffer(self):
if self.layout == "layer_first":
dims = (
self.layer_num,
self.size,
1,
self.kv_cache_dim,
)
elif self.layout == "page_first":
dims = (
self.size,
self.layer_num,
1,
self.kv_cache_dim,
)
elif self.layout == "page_first_direct":
dims = (
self.page_num,
self.layer_num,
self.page_size,
1,
self.kv_cache_dim,
)
# Ascend-specific: Aligns with NPUMLATokenToKVPool layout
# Separately allocate k_buffer and v_buffer for easier data transfer.
elif self.layout == "page_first_kv_split":
base_dims = (
self.page_num,
self.layer_num,
self.page_size,
1,
)
alloc_func = ALLOC_MEMORY_FUNCS[self.device_pool.device]
self.k_buffer = alloc_func(
(*base_dims, self.kv_lora_rank),
dtype=self.dtype,
device=self.device,
pin_memory=self.pin_memory,
allocator=self.allocator,
)
self.v_buffer = alloc_func(
(*base_dims, self.qk_rope_head_dim),
dtype=self.dtype,
device=self.device,
pin_memory=self.pin_memory,
allocator=self.allocator,
)
self.index_k_buffer = None
if self.device_pool.index_head_dim is not None:
self.index_k_buffer = alloc_func(
(*base_dims, self.device_pool.index_head_dim),
dtype=self.dtype,
device=self.device,
pin_memory=self.pin_memory,
allocator=self.allocator,
)
# Return k_buffer to preserve original kv_buffer and data_refs init logic,
# though Ascend doesn't use these parameters.
return self.k_buffer
else:
raise ValueError(f"Unsupported layout: {self.layout}")
self.token_stride_size = self.kv_cache_dim * self.dtype.itemsize
self.layout_dim = self.token_stride_size * self.layer_num
alloc_func = ALLOC_MEMORY_FUNCS[self.device_pool.device]
buffer = alloc_func(
dims,
dtype=self.dtype,
device=self.device,
pin_memory=self.pin_memory,
allocator=self.allocator,
)
return buffer
def _init_write_back_staging_buffers(self):
self.staging_page_capacity = 0
self.staging_token_capacity = 0
self.staging_buffer = None
self.can_use_write_back_jit = False
if self.layout != "page_first" or (_is_npu or _is_xpu or _is_mps):
return
# The staged write-back JIT kernel builds with hipcc and has a ROCm
# path, so enable it on HIP too (consistent with the CUDA path).
self.can_use_write_back_jit = (
_is_cuda or _is_hip
) and can_use_write_back_jit_kernel(
element_size=self.kv_cache_dim * self.dtype.itemsize,
)
if not self.can_use_write_back_jit:
return
self.staging_page_capacity = min(self.page_num, _WRITE_BACK_STAGING_PAGE_CHUNK)
self.staging_token_capacity = self.staging_page_capacity * self.page_size
self.staging_buffer = torch.empty(
(
self.staging_token_capacity,
self.layer_num,
1,
self.kv_cache_dim,
),
dtype=self.dtype,
device=self.device_pool.device,
)
def load_to_device_per_layer(
self, device_pool, host_indices, device_indices, layer_id, io_backend
):
if not self._is_device_layer_owned(device_pool, layer_id):
return
host_layer = self._host_layer_index(layer_id)
if io_backend == "kernel":
if self.layout == "layer_first":
if self.can_use_jit:
jit_transfer_hicache_one_layer_mla(
cache_dst=device_pool.kv_buffer[layer_id],
cache_src=self.kv_buffer[host_layer],
indices_dst=device_indices,
indices_src=host_indices,
element_dim=self.kv_cache_dim,
)
else:
transfer_kv_per_layer_mla(
src=self.kv_buffer[host_layer],
dst=device_pool.kv_buffer[layer_id],
src_indices=host_indices,
dst_indices=device_indices,
item_size=self.token_stride_size,
)
elif self.layout == "page_first":
if self.can_use_jit:
jit_transfer_hicache_one_layer_mla(
cache_dst=device_pool.kv_buffer[layer_id],
cache_src=self.data_refs[host_layer],
indices_dst=device_indices,
indices_src=host_indices,
element_dim=self.kv_cache_dim,
)
else:
transfer_kv_per_layer_mla_pf_lf(
src=self.kv_buffer,
dst=device_pool.kv_buffer[layer_id],
src_indices=host_indices,
dst_indices=device_indices,
layer_id=host_layer,
item_size=self.token_stride_size,
src_layout_dim=self.layout_dim,
)
else:
raise ValueError(f"Unsupported layout: {self.layout}")
elif io_backend == "direct":
if self.layout == "layer_first":
transfer_kv_direct(
src_layers=[self.kv_buffer[host_layer]],
dst_layers=[device_pool.kv_buffer[layer_id]],
src_indices=host_indices,
dst_indices=device_indices,
page_size=self.page_size,
)
elif self.layout == "page_first_direct":
transfer_kv_per_layer_direct_pf_lf(
src_ptrs=[self.kv_buffer],
dst_ptrs=[device_pool.kv_buffer[layer_id]],
src_indices=host_indices,
dst_indices=device_indices,
layer_id=host_layer,
page_size=self.page_size,
)
else:
raise ValueError(f"Unsupported layout: {self.layout}")
elif io_backend == "kernel_ascend":
if self.layout == "page_first_kv_split":
# Ascend-specific: transfer KV data for all layers when layer_id == 0
if layer_id == 0:
transfer_kv_dim_exchange(
device_indices=device_indices,
host_indices=host_indices,
device_k=device_pool.k_buffer,
host_k=self.k_buffer,
device_v=device_pool.v_buffer,
host_v=self.v_buffer,
device_index_k=device_pool.index_k_buffer,
host_index_k=self.index_k_buffer,
page_size=self.page_size,
direction=TransferDirection.H2D,
)
else:
raise ValueError(f"Unsupported layout: {self.layout}")
else:
raise ValueError(f"Unsupported IO backend: {io_backend}")
def _backup_from_device_per_layer(
self, device_pool, host_indices, device_indices, layer_id, io_backend
):
host_layer = self._host_layer_index(layer_id)
if io_backend == "kernel":
if self.layout == "layer_first":
if self.can_use_jit:
jit_transfer_hicache_one_layer_mla(
cache_dst=self.kv_buffer[host_layer],
cache_src=device_pool.kv_buffer[layer_id],
indices_dst=host_indices,
indices_src=device_indices,
element_dim=self.kv_cache_dim,
)
else:
transfer_kv_per_layer_mla(
src=device_pool.kv_buffer[layer_id],
dst=self.kv_buffer[host_layer],
src_indices=device_indices,
dst_indices=host_indices,
item_size=self.token_stride_size,
)
elif self.layout == "page_first":
if self.can_use_jit:
jit_transfer_hicache_one_layer_mla(
cache_dst=self.data_refs[host_layer],
cache_src=device_pool.kv_buffer[layer_id],
indices_dst=host_indices,
indices_src=device_indices,
element_dim=self.kv_cache_dim,
)
else:
raise ValueError(
"Layer-sharded MLA HiCache backup with page_first layout "
"requires the JIT one-layer kernel."
)
else:
raise ValueError(
f"Layer-sharded HiCache backup does not support layout: {self.layout}"
)
elif io_backend == "direct":
if self.layout == "layer_first":
transfer_kv_direct(
src_layers=[device_pool.kv_buffer[layer_id]],
dst_layers=[self.kv_buffer[host_layer]],
src_indices=device_indices,
dst_indices=host_indices,
page_size=self.page_size,
)
else:
raise ValueError(
"Layer-sharded direct HiCache backup only supports "
f"layer_first layout, got {self.layout}"
)
else:
raise ValueError(
f"Layer-sharded HiCache backup does not support IO backend: {io_backend}"
)
def backup_from_device_all_layer(
self, device_pool, host_indices, device_indices, io_backend
):
if self._is_device_layer_sharded(device_pool):
for layer_id in self._owned_device_layer_ids(device_pool):
self._backup_from_device_per_layer(
device_pool, host_indices, device_indices, layer_id, io_backend
)
return
if io_backend == "kernel":
if self.layout == "layer_first":
if self.can_use_jit:
jit_transfer_hicache_all_layer_mla(
ptr_dst=self.data_ptrs,
indices_dst=host_indices,
ptr_src=device_pool.data_ptrs,
indices_src=device_indices,
cache_dst_stride_bytes=self.token_stride_size,
cache_src_stride_bytes=self.token_stride_size,
element_size=self.kv_cache_dim * self.dtype.itemsize,
)
else:
transfer_kv_all_layer_mla(
src_layers=device_pool.data_ptrs,
dst_layers=self.data_ptrs,
src_indices=device_indices,
dst_indices=host_indices,
item_size=self.token_stride_size,
num_layers=self.layer_num,
)
elif self.layout == "page_first":
if self.can_use_write_back_jit:
jit_transfer_hicache_all_layer_mla_staged_lf_pf(
ptr_src=device_pool.data_ptrs,
src_indices=device_indices,
dst_indices=host_indices,
staging=self.staging_buffer,
dst=self.kv_buffer,
page_size=self.page_size,
)
else:
transfer_kv_all_layer_mla_lf_pf(
src_layers=device_pool.data_ptrs,
dst=self.kv_buffer,
src_indices=device_indices,
dst_indices=host_indices,
item_size=self.token_stride_size,
dst_layout_dim=self.layout_dim,
num_layers=self.layer_num,
)
else:
raise ValueError(f"Unsupported layout: {self.layout}")
elif io_backend == "direct":
if self.layout == "layer_first":
transfer_kv_direct(
src_layers=device_pool.kv_buffer,
dst_layers=self.data_refs,
src_indices=device_indices,
dst_indices=host_indices,
page_size=self.page_size,
)
elif self.layout == "page_first_direct":
transfer_kv_all_layer_direct_lf_pf(
src_ptrs=device_pool.kv_buffer,
dst_ptrs=[self.kv_buffer],
src_indices=device_indices,
dst_indices=host_indices,
page_size=self.page_size,
)
else:
raise ValueError(f"Unsupported layout: {self.layout}")
elif io_backend == "kernel_ascend":
if self.layout == "page_first_kv_split":
transfer_kv_dim_exchange(
device_indices=device_indices,
host_indices=host_indices,
device_k=device_pool.k_buffer,
host_k=self.k_buffer,
device_v=device_pool.v_buffer,
host_v=self.v_buffer,
device_index_k=device_pool.index_k_buffer,
host_index_k=self.index_k_buffer,
page_size=self.page_size,
direction=TransferDirection.D2H,
)
else:
raise ValueError(f"Unsupported layout: {self.layout}")
else:
raise ValueError(f"Unsupported IO backend: {io_backend}")
def get_data_page(self, index, flat: bool = True) -> torch.Tensor:
if self.layout == "layer_first":
data_page = self.kv_buffer[:, index : index + self.page_size, :, :]
elif self.layout == "page_first":
data_page = self.kv_buffer[index : index + self.page_size, :, :, :]
elif self.layout == "page_first_direct":
real_index = index // self.page_size
data_page = self.kv_buffer[real_index : real_index + 1, :, :, :, :]
else:
raise ValueError(f"Unsupported layout: {self.layout}")
if flat:
data_page = data_page.flatten()
return data_page
def get_dummy_flat_data_page(self) -> torch.Tensor:
return torch.zeros(
(
self.layer_num,
self.page_size,
1,
self.kv_cache_dim,
),
dtype=self.dtype,
device=self.device,
pin_memory=self.pin_memory,
).flatten()
def set_from_flat_data_page(self, index: int, data_page: torch.Tensor) -> None:
if self.layout == "layer_first":
self.kv_buffer[:, index : index + self.page_size, :, :] = data_page.reshape(
self.layer_num,
self.page_size,
1,
self.kv_cache_dim,
)
elif self.layout == "page_first":
self.kv_buffer[index : index + self.page_size, :, :, :] = data_page.reshape(
self.page_size,
self.layer_num,
1,
self.kv_cache_dim,
)
elif self.layout == "page_first_direct":
real_index = index // self.page_size
self.kv_buffer[real_index : real_index + 1, :, :, :, :] = data_page.reshape(
1,
self.layer_num,
self.page_size,
1,
self.kv_cache_dim,
)
else:
raise ValueError(f"Unsupported layout: {self.layout}")
def get_page_buffer_meta(self, indices):
"""
meta data for zero copy
"""
assert len(indices) % self.page_size == 0
ptr_list = []
kv_buffer_data_ptr = self.kv_buffer.data_ptr()
indices = indices.tolist()
if self.layout == "layer_first":
for index in range(0, len(indices), self.page_size):
for layer_id in range(self.layer_num):
k_ptr = (
kv_buffer_data_ptr
+ indices[index] * self.kv_cache_dim * self.dtype.itemsize
+ layer_id * self.size * self.kv_cache_dim * self.dtype.itemsize
)
ptr_list.append(k_ptr)
element_size = self.dtype.itemsize * self.page_size * self.kv_cache_dim
element_size_list = [element_size] * len(ptr_list)
elif self.layout in ["page_first", "page_first_direct"]:
for index in range(0, len(indices), self.page_size):
k_ptr = (
kv_buffer_data_ptr
+ indices[index]
* self.layer_num
* self.kv_cache_dim
* self.dtype.itemsize
)
ptr_list.append(k_ptr)
element_size = (
self.layer_num
* self.dtype.itemsize
* self.page_size
* self.kv_cache_dim
)
element_size_list = [element_size] * len(ptr_list)
else:
raise ValueError(f"Unsupported layout: {self.layout}")
return ptr_list, element_size_list
def is_stride_page_aligned(self, page_size_bytes: int = 4096) -> bool:
"""Return True if per-page strides are multiples of *page_size_bytes*.
When O_DIRECT is used with any file-based NIXL backend, every data pointer
passed to the kernel must be page-aligned. In zero-copy mode the
pointer for KV page ``p`` is:
base_ptr + p * page_size * layer_num * kv_cache_dim * itemsize
For this to be page-aligned (given a page-aligned ``base_ptr``) the per-page
stride must itself be a multiple of the OS page size.
"""
if self.layout not in ("page_first", "page_first_direct"):
return False
stride = (
self.page_size * self.layer_num * self.kv_cache_dim * self.dtype.itemsize
)
base_aligned = self.kv_buffer.data_ptr() % page_size_bytes == 0
return base_aligned and stride % page_size_bytes == 0
class MambaPoolHost(HostKVCache):
def __init__(
python/sglang/srt/mem_cache/pool_host/mla.py
+573
View File
@@ -0,0 +1,573 @@
from __future__ import annotations
import logging
from typing import Optional
import torch
from sglang.jit_kernel.hicache import (
can_use_hicache_jit_kernel,
can_use_write_back_jit_kernel,
)
from sglang.jit_kernel.hicache import (
transfer_hicache_all_layer_mla as jit_transfer_hicache_all_layer_mla,
)
from sglang.jit_kernel.hicache import (
transfer_hicache_all_layer_mla_staged_lf_pf as jit_transfer_hicache_all_layer_mla_staged_lf_pf,
)
from sglang.jit_kernel.hicache import (
transfer_hicache_one_layer_mla as jit_transfer_hicache_one_layer_mla,
)
from sglang.srt.mem_cache.memory_pool import MLATokenToKVPool
from sglang.srt.mem_cache.pool_host.base import (
_WRITE_BACK_STAGING_PAGE_CHUNK,
HostKVCache,
)
from sglang.srt.mem_cache.pool_host.common import ALLOC_MEMORY_FUNCS
from sglang.srt.mem_cache.pool_host.hisparse import HiSparseHostPoolMixin
from sglang.srt.utils import is_cuda, is_hip, is_mps, is_npu, is_xpu
_is_cuda = is_cuda()
_is_hip = is_hip()
_is_npu = is_npu()
_is_xpu = is_xpu()
_is_mps = is_mps()
if _is_cuda or _is_hip:
from sgl_kernel.kvcacheio import (
transfer_kv_all_layer_direct_lf_pf,
transfer_kv_all_layer_mla,
transfer_kv_all_layer_mla_lf_pf,
transfer_kv_direct,
transfer_kv_per_layer_direct_pf_lf,
transfer_kv_per_layer_mla,
transfer_kv_per_layer_mla_pf_lf,
)
if _is_npu:
from sgl_kernel_npu.kvcacheio import TransferDirection, transfer_kv_dim_exchange
logger = logging.getLogger(__name__)
class MLATokenToKVPoolHost(HiSparseHostPoolMixin, HostKVCache):
device_pool: MLATokenToKVPool
def __init__(
self,
device_pool: MLATokenToKVPool,
host_to_device_ratio: float,
host_size: int,
page_size: int,
layout: str,
pin_memory: bool = True,
device: str = "cpu",
allocator_type: str = "default",
override_kv_cache_dim: Optional[int] = None,
):
self.override_kv_cache_dim = override_kv_cache_dim
super().__init__(
device_pool,
host_to_device_ratio,
host_size,
page_size,
layout,
pin_memory,
device,
allocator_type,
)
# The JIT HiCache kernels also build with hipcc (ROCm): the PTX-only
# helpers in hicache.cuh are guarded by USE_ROCM and the staged
# write-back kernel has a ROCm path, so enable them on HIP too. This
# keeps the ROCm write-back path consistent with CUDA.
self.can_use_jit = (_is_cuda or _is_hip) and can_use_hicache_jit_kernel(
element_size=self.kv_cache_dim * self.dtype.itemsize
)
if self.layout == "page_first":
# Transpose [page, layer, ...] -> [layer, page, ...] to get per-layer views
# This swaps strides without copying data
transposed = self.kv_buffer.transpose(0, 1)
self.data_refs = [transposed[i] for i in range(self.layer_num)]
else:
self.data_refs = [self.kv_buffer[i] for i in range(self.layer_num)]
self.data_ptrs = torch.tensor(
[x.data_ptr() for x in self.data_refs],
dtype=torch.uint64,
device=self.device_pool.device,
)
self._init_write_back_staging_buffers()
def get_contiguous_buf_infos(self):
"""Return (data_ptrs, data_lens, item_lens) in the same format as device pool,
for registering host memory with the disaggregation transfer engine."""
data_ptrs = [int(self.data_ptrs[i].item()) for i in range(self.layer_num)]
data_lens = [self.kv_buffer[i].nbytes for i in range(self.layer_num)]
item_lens = [self.token_stride_size * self.page_size] * self.layer_num
return data_ptrs, data_lens, item_lens
def get_size_per_token(self):
self.kv_lora_rank = self.device_pool.kv_lora_rank
self.qk_rope_head_dim = self.device_pool.qk_rope_head_dim
self.layer_num = self._effective_host_layer_num()
self.kv_cache_dim = self.override_kv_cache_dim or (
self.kv_lora_rank + self.qk_rope_head_dim
)
return self.kv_cache_dim * self.dtype.itemsize * self.layer_num
def get_ksize_per_token(self):
return self.get_size_per_token()
def init_kv_buffer(self):
if self.layout == "layer_first":
dims = (
self.layer_num,
self.size,
1,
self.kv_cache_dim,
)
elif self.layout == "page_first":
dims = (
self.size,
self.layer_num,
1,
self.kv_cache_dim,
)
elif self.layout == "page_first_direct":
dims = (
self.page_num,
self.layer_num,
self.page_size,
1,
self.kv_cache_dim,
)
# Ascend-specific: Aligns with NPUMLATokenToKVPool layout
# Separately allocate k_buffer and v_buffer for easier data transfer.
elif self.layout == "page_first_kv_split":
base_dims = (
self.page_num,
self.layer_num,
self.page_size,
1,
)
alloc_func = ALLOC_MEMORY_FUNCS[self.device_pool.device]
self.k_buffer = alloc_func(
(*base_dims, self.kv_lora_rank),
dtype=self.dtype,
device=self.device,
pin_memory=self.pin_memory,
allocator=self.allocator,
)
self.v_buffer = alloc_func(
(*base_dims, self.qk_rope_head_dim),
dtype=self.dtype,
device=self.device,
pin_memory=self.pin_memory,
allocator=self.allocator,
)
self.index_k_buffer = None
if self.device_pool.index_head_dim is not None:
self.index_k_buffer = alloc_func(
(*base_dims, self.device_pool.index_head_dim),
dtype=self.dtype,
device=self.device,
pin_memory=self.pin_memory,
allocator=self.allocator,
)
# Return k_buffer to preserve original kv_buffer and data_refs init logic,
# though Ascend doesn't use these parameters.
return self.k_buffer
else:
raise ValueError(f"Unsupported layout: {self.layout}")
self.token_stride_size = self.kv_cache_dim * self.dtype.itemsize
self.layout_dim = self.token_stride_size * self.layer_num
alloc_func = ALLOC_MEMORY_FUNCS[self.device_pool.device]
buffer = alloc_func(
dims,
dtype=self.dtype,
device=self.device,
pin_memory=self.pin_memory,
allocator=self.allocator,
)
return buffer
def _init_write_back_staging_buffers(self):
self.staging_page_capacity = 0
self.staging_token_capacity = 0
self.staging_buffer = None
self.can_use_write_back_jit = False
if self.layout != "page_first" or (_is_npu or _is_xpu or _is_mps):
return
# The staged write-back JIT kernel builds with hipcc and has a ROCm
# path, so enable it on HIP too (consistent with the CUDA path).
self.can_use_write_back_jit = (
_is_cuda or _is_hip
) and can_use_write_back_jit_kernel(
element_size=self.kv_cache_dim * self.dtype.itemsize,
)
if not self.can_use_write_back_jit:
return
self.staging_page_capacity = min(self.page_num, _WRITE_BACK_STAGING_PAGE_CHUNK)
self.staging_token_capacity = self.staging_page_capacity * self.page_size
self.staging_buffer = torch.empty(
(
self.staging_token_capacity,
self.layer_num,
1,
self.kv_cache_dim,
),
dtype=self.dtype,
device=self.device_pool.device,
)
def load_to_device_per_layer(
self, device_pool, host_indices, device_indices, layer_id, io_backend
):
if not self._is_device_layer_owned(device_pool, layer_id):
return
host_layer = self._host_layer_index(layer_id)
if io_backend == "kernel":
if self.layout == "layer_first":
if self.can_use_jit:
jit_transfer_hicache_one_layer_mla(
cache_dst=device_pool.kv_buffer[layer_id],
cache_src=self.kv_buffer[host_layer],
indices_dst=device_indices,
indices_src=host_indices,
element_dim=self.kv_cache_dim,
)
else:
transfer_kv_per_layer_mla(
src=self.kv_buffer[host_layer],
dst=device_pool.kv_buffer[layer_id],
src_indices=host_indices,
dst_indices=device_indices,
item_size=self.token_stride_size,
)
elif self.layout == "page_first":
if self.can_use_jit:
jit_transfer_hicache_one_layer_mla(
cache_dst=device_pool.kv_buffer[layer_id],
cache_src=self.data_refs[host_layer],
indices_dst=device_indices,
indices_src=host_indices,
element_dim=self.kv_cache_dim,
)
else:
transfer_kv_per_layer_mla_pf_lf(
src=self.kv_buffer,
dst=device_pool.kv_buffer[layer_id],
src_indices=host_indices,
dst_indices=device_indices,
layer_id=host_layer,
item_size=self.token_stride_size,
src_layout_dim=self.layout_dim,
)
else:
raise ValueError(f"Unsupported layout: {self.layout}")
elif io_backend == "direct":
if self.layout == "layer_first":
transfer_kv_direct(
src_layers=[self.kv_buffer[host_layer]],
dst_layers=[device_pool.kv_buffer[layer_id]],
src_indices=host_indices,
dst_indices=device_indices,
page_size=self.page_size,
)
elif self.layout == "page_first_direct":
transfer_kv_per_layer_direct_pf_lf(
src_ptrs=[self.kv_buffer],
dst_ptrs=[device_pool.kv_buffer[layer_id]],
src_indices=host_indices,
dst_indices=device_indices,
layer_id=host_layer,
page_size=self.page_size,
)
else:
raise ValueError(f"Unsupported layout: {self.layout}")
elif io_backend == "kernel_ascend":
if self.layout == "page_first_kv_split":
# Ascend-specific: transfer KV data for all layers when layer_id == 0
if layer_id == 0:
transfer_kv_dim_exchange(
device_indices=device_indices,
host_indices=host_indices,
device_k=device_pool.k_buffer,
host_k=self.k_buffer,
device_v=device_pool.v_buffer,
host_v=self.v_buffer,
device_index_k=device_pool.index_k_buffer,
host_index_k=self.index_k_buffer,
page_size=self.page_size,
direction=TransferDirection.H2D,
)
else:
raise ValueError(f"Unsupported layout: {self.layout}")
else:
raise ValueError(f"Unsupported IO backend: {io_backend}")
def _backup_from_device_per_layer(
self, device_pool, host_indices, device_indices, layer_id, io_backend
):
host_layer = self._host_layer_index(layer_id)
if io_backend == "kernel":
if self.layout == "layer_first":
if self.can_use_jit:
jit_transfer_hicache_one_layer_mla(
cache_dst=self.kv_buffer[host_layer],
cache_src=device_pool.kv_buffer[layer_id],
indices_dst=host_indices,
indices_src=device_indices,
element_dim=self.kv_cache_dim,
)
else:
transfer_kv_per_layer_mla(
src=device_pool.kv_buffer[layer_id],
dst=self.kv_buffer[host_layer],
src_indices=device_indices,
dst_indices=host_indices,
item_size=self.token_stride_size,
)
elif self.layout == "page_first":
if self.can_use_jit:
jit_transfer_hicache_one_layer_mla(
cache_dst=self.data_refs[host_layer],
cache_src=device_pool.kv_buffer[layer_id],
indices_dst=host_indices,
indices_src=device_indices,
element_dim=self.kv_cache_dim,
)
else:
raise ValueError(
"Layer-sharded MLA HiCache backup with page_first layout "
"requires the JIT one-layer kernel."
)
else:
raise ValueError(
f"Layer-sharded HiCache backup does not support layout: {self.layout}"
)
elif io_backend == "direct":
if self.layout == "layer_first":
transfer_kv_direct(
src_layers=[device_pool.kv_buffer[layer_id]],
dst_layers=[self.kv_buffer[host_layer]],
src_indices=device_indices,
dst_indices=host_indices,
page_size=self.page_size,
)
else:
raise ValueError(
"Layer-sharded direct HiCache backup only supports "
f"layer_first layout, got {self.layout}"
)
else:
raise ValueError(
f"Layer-sharded HiCache backup does not support IO backend: {io_backend}"
)
def backup_from_device_all_layer(
self, device_pool, host_indices, device_indices, io_backend
):
if self._is_device_layer_sharded(device_pool):
for layer_id in self._owned_device_layer_ids(device_pool):
self._backup_from_device_per_layer(
device_pool, host_indices, device_indices, layer_id, io_backend
)
return
if io_backend == "kernel":
if self.layout == "layer_first":
if self.can_use_jit:
jit_transfer_hicache_all_layer_mla(
ptr_dst=self.data_ptrs,
indices_dst=host_indices,
ptr_src=device_pool.data_ptrs,
indices_src=device_indices,
cache_dst_stride_bytes=self.token_stride_size,
cache_src_stride_bytes=self.token_stride_size,
element_size=self.kv_cache_dim * self.dtype.itemsize,
)
else:
transfer_kv_all_layer_mla(
src_layers=device_pool.data_ptrs,
dst_layers=self.data_ptrs,
src_indices=device_indices,
dst_indices=host_indices,
item_size=self.token_stride_size,
num_layers=self.layer_num,
)
elif self.layout == "page_first":
if self.can_use_write_back_jit:
jit_transfer_hicache_all_layer_mla_staged_lf_pf(
ptr_src=device_pool.data_ptrs,
src_indices=device_indices,
dst_indices=host_indices,
staging=self.staging_buffer,
dst=self.kv_buffer,
page_size=self.page_size,
)
else:
transfer_kv_all_layer_mla_lf_pf(
src_layers=device_pool.data_ptrs,
dst=self.kv_buffer,
src_indices=device_indices,
dst_indices=host_indices,
item_size=self.token_stride_size,
dst_layout_dim=self.layout_dim,
num_layers=self.layer_num,
)
else:
raise ValueError(f"Unsupported layout: {self.layout}")
elif io_backend == "direct":
if self.layout == "layer_first":
transfer_kv_direct(
src_layers=device_pool.kv_buffer,
dst_layers=self.data_refs,
src_indices=device_indices,
dst_indices=host_indices,
page_size=self.page_size,
)
elif self.layout == "page_first_direct":
transfer_kv_all_layer_direct_lf_pf(
src_ptrs=device_pool.kv_buffer,
dst_ptrs=[self.kv_buffer],
src_indices=device_indices,
dst_indices=host_indices,
page_size=self.page_size,
)
else:
raise ValueError(f"Unsupported layout: {self.layout}")
elif io_backend == "kernel_ascend":
if self.layout == "page_first_kv_split":
transfer_kv_dim_exchange(
device_indices=device_indices,
host_indices=host_indices,
device_k=device_pool.k_buffer,
host_k=self.k_buffer,
device_v=device_pool.v_buffer,
host_v=self.v_buffer,
device_index_k=device_pool.index_k_buffer,
host_index_k=self.index_k_buffer,
page_size=self.page_size,
direction=TransferDirection.D2H,
)
else:
raise ValueError(f"Unsupported layout: {self.layout}")
else:
raise ValueError(f"Unsupported IO backend: {io_backend}")
def get_data_page(self, index, flat: bool = True) -> torch.Tensor:
if self.layout == "layer_first":
data_page = self.kv_buffer[:, index : index + self.page_size, :, :]
elif self.layout == "page_first":
data_page = self.kv_buffer[index : index + self.page_size, :, :, :]
elif self.layout == "page_first_direct":
real_index = index // self.page_size
data_page = self.kv_buffer[real_index : real_index + 1, :, :, :, :]
else:
raise ValueError(f"Unsupported layout: {self.layout}")
if flat:
data_page = data_page.flatten()
return data_page
def get_dummy_flat_data_page(self) -> torch.Tensor:
return torch.zeros(
(
self.layer_num,
self.page_size,
1,
self.kv_cache_dim,
),
dtype=self.dtype,
device=self.device,
pin_memory=self.pin_memory,
).flatten()
def set_from_flat_data_page(self, index: int, data_page: torch.Tensor) -> None:
if self.layout == "layer_first":
self.kv_buffer[:, index : index + self.page_size, :, :] = data_page.reshape(
self.layer_num,
self.page_size,
1,
self.kv_cache_dim,
)
elif self.layout == "page_first":
self.kv_buffer[index : index + self.page_size, :, :, :] = data_page.reshape(
self.page_size,
self.layer_num,
1,
self.kv_cache_dim,
)
elif self.layout == "page_first_direct":
real_index = index // self.page_size
self.kv_buffer[real_index : real_index + 1, :, :, :, :] = data_page.reshape(
1,
self.layer_num,
self.page_size,
1,
self.kv_cache_dim,
)
else:
raise ValueError(f"Unsupported layout: {self.layout}")
def get_page_buffer_meta(self, indices):
"""
meta data for zero copy
"""
assert len(indices) % self.page_size == 0
ptr_list = []
kv_buffer_data_ptr = self.kv_buffer.data_ptr()
indices = indices.tolist()
if self.layout == "layer_first":
for index in range(0, len(indices), self.page_size):
for layer_id in range(self.layer_num):
k_ptr = (
kv_buffer_data_ptr
+ indices[index] * self.kv_cache_dim * self.dtype.itemsize
+ layer_id * self.size * self.kv_cache_dim * self.dtype.itemsize
)
ptr_list.append(k_ptr)
element_size = self.dtype.itemsize * self.page_size * self.kv_cache_dim
element_size_list = [element_size] * len(ptr_list)
elif self.layout in ["page_first", "page_first_direct"]:
for index in range(0, len(indices), self.page_size):
k_ptr = (
kv_buffer_data_ptr
+ indices[index]
* self.layer_num
* self.kv_cache_dim
* self.dtype.itemsize
)
ptr_list.append(k_ptr)
element_size = (
self.layer_num
* self.dtype.itemsize
* self.page_size
* self.kv_cache_dim
)
element_size_list = [element_size] * len(ptr_list)
else:
raise ValueError(f"Unsupported layout: {self.layout}")
return ptr_list, element_size_list
def is_stride_page_aligned(self, page_size_bytes: int = 4096) -> bool:
"""Return True if per-page strides are multiples of *page_size_bytes*.
When O_DIRECT is used with any file-based NIXL backend, every data pointer
passed to the kernel must be page-aligned. In zero-copy mode the
pointer for KV page ``p`` is:
base_ptr + p * page_size * layer_num * kv_cache_dim * itemsize
For this to be page-aligned (given a page-aligned ``base_ptr``) the per-page
stride must itself be a multiple of the OS page size.
"""
if self.layout not in ("page_first", "page_first_direct"):
return False
stride = (
self.page_size * self.layer_num * self.kv_cache_dim * self.dtype.itemsize
)
base_aligned = self.kv_buffer.data_ptr() % page_size_bytes == 0
return base_aligned and stride % page_size_bytes == 0
python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py
+1 -1
View File
@@ -21,8 +21,8 @@ from sglang.srt.mem_cache.hicache_storage import (
PoolTransfer,
PoolTransferResult,
)
from sglang.srt.mem_cache.memory_pool_host import MLATokenToKVPoolHost
from sglang.srt.mem_cache.pool_host import HostKVCache, HostTensorAllocator
from sglang.srt.mem_cache.pool_host.mla import MLATokenToKVPoolHost
from sglang.srt.observability.metrics_collector import StorageMetrics
DEFAULT_LOCAL_BUFFER_SIZE = 16 * 1024 * 1024 # 16 MB