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, MLATokenToKVPool,
ReqToTokenPool, 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.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.server_args import ServerArgs
from sglang.srt.utils.common import ceil_align 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, HiSparseDSATokenToKVPool,
) )
from sglang.srt.mem_cache.memory_pool import ReqToTokenPool from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
from sglang.srt.mem_cache.memory_pool_host import ( from sglang.srt.mem_cache.memory_pool_host import DeepSeekV4PagedHostPool
DeepSeekV4PagedHostPool, from sglang.srt.mem_cache.pool_host.mla import MLATokenToKVPoolHost
MLATokenToKVPoolHost,
)
from sglang.srt.utils import get_device_module, is_hip from sglang.srt.utils import get_device_module, is_hip
device_module = get_device_module() 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, MiniMaxSparseKVPool,
MLATokenToKVPool, 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.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 ( from sglang.srt.mem_cache.radix_cache import (
RadixCache, RadixCache,
RadixKey, 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, HostPoolGroup,
LogicalHostPool, LogicalHostPool,
MambaPoolHost, MambaPoolHost,
MLATokenToKVPoolHost,
PoolEntry, PoolEntry,
) )
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,
) )
from sglang.srt.mem_cache.pool_host.mla import MLATokenToKVPoolHost
from sglang.srt.mem_cache.unified_cache_components import ComponentType from sglang.srt.mem_cache.unified_cache_components import ComponentType
if TYPE_CHECKING: 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, MHATokenToKVPool,
MLATokenToKVPool, 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.mha import get_mha_host_pool_cls
from sglang.srt.mem_cache.pool_host.mla import MLATokenToKVPoolHost
pool = draft_kv_pool pool = draft_kv_pool
if isinstance(pool, HybridLinearKVPool): 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: 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
import numpy as np import numpy as np
import psutil import psutil
import torch import torch
from sglang.jit_kernel.hicache import ( from sglang.jit_kernel.hicache import (
can_use_hicache_jit_kernel,
can_use_write_back_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 ( from sglang.jit_kernel.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.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.jit_kernel.hisparse import transfer_cache_dsv4_mla
from sglang.srt.mem_cache.memory_pool import ( from sglang.srt.mem_cache.memory_pool import (
DSATokenToKVPool, DSATokenToKVPool,
MambaPool, MambaPool,
MLATokenToKVPool,
) )
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
@@ -48,8 +41,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_npu:
from sgl_kernel_npu.kvcacheio import TransferDirection, transfer_kv_dim_exchange
logger = logging.getLogger(__name__) 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 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): class MambaPoolHost(HostKVCache):
def __init__( 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, PoolTransfer,
PoolTransferResult, 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 import HostKVCache, HostTensorAllocator
from sglang.srt.mem_cache.pool_host.mla import MLATokenToKVPoolHost
from sglang.srt.observability.metrics_collector import StorageMetrics from sglang.srt.observability.metrics_collector import StorageMetrics
DEFAULT_LOCAL_BUFFER_SIZE = 16 * 1024 * 1024 # 16 MB DEFAULT_LOCAL_BUFFER_SIZE = 16 * 1024 * 1024 # 16 MB
test/registered/jit/test_hicache.py
+1 -1
View File
@@ -5,12 +5,12 @@ import torch
from sglang.jit_kernel.hicache import can_use_write_back_jit_kernel from sglang.jit_kernel.hicache import can_use_write_back_jit_kernel
from sglang.srt.mem_cache.memory_pool import MHATokenToKVPool, MLATokenToKVPool from sglang.srt.mem_cache.memory_pool import MHATokenToKVPool, MLATokenToKVPool
from sglang.srt.mem_cache.memory_pool_host import MLATokenToKVPoolHost
from sglang.srt.mem_cache.pool_host.common import ( from sglang.srt.mem_cache.pool_host.common import (
ALLOC_MEMORY_FUNCS, ALLOC_MEMORY_FUNCS,
alloc_with_pin_memory, alloc_with_pin_memory,
) )
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.utils import is_cuda, is_hip, is_npu, is_xpu from sglang.srt.utils import is_cuda, is_hip, is_npu, is_xpu
from sglang.test.ci.ci_register import register_cuda_ci from sglang.test.ci.ci_register import register_cuda_ci
test/registered/jit/test_hicache_page_first_write_back.py
+1 -1
View File
@@ -15,12 +15,12 @@ import torch
from sglang.jit_kernel.hicache import can_use_write_back_jit_kernel from sglang.jit_kernel.hicache import can_use_write_back_jit_kernel
from sglang.srt.mem_cache.memory_pool import MHATokenToKVPool, MLATokenToKVPool from sglang.srt.mem_cache.memory_pool import MHATokenToKVPool, MLATokenToKVPool
from sglang.srt.mem_cache.memory_pool_host import MLATokenToKVPoolHost
from sglang.srt.mem_cache.pool_host.common import ( from sglang.srt.mem_cache.pool_host.common import (
ALLOC_MEMORY_FUNCS, ALLOC_MEMORY_FUNCS,
alloc_with_pin_memory, alloc_with_pin_memory,
) )
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.utils import is_cuda, is_hip, is_npu, is_xpu from sglang.srt.utils import is_cuda, is_hip, is_npu, is_xpu
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
test/registered/unit/mem_cache/test_dsa_pool_host_unit.py
+3 -5
View File
@@ -4,14 +4,12 @@ import unittest
import torch import torch
from sglang.srt.mem_cache.memory_pool import DSATokenToKVPool from sglang.srt.mem_cache.memory_pool import DSATokenToKVPool
from sglang.srt.mem_cache.memory_pool_host import ( from sglang.srt.mem_cache.memory_pool_host import DSAIndexerPoolHost
DSAIndexerPoolHost,
MLATokenToKVPoolHost,
)
from sglang.srt.mem_cache.pool_host.common import ( from sglang.srt.mem_cache.pool_host.common import (
ALLOC_MEMORY_FUNCS, ALLOC_MEMORY_FUNCS,
alloc_with_pin_memory, alloc_with_pin_memory,
) )
from sglang.srt.mem_cache.pool_host.mla import MLATokenToKVPoolHost
from sglang.srt.utils import is_cuda, is_hip, is_npu, is_xpu from sglang.srt.utils import is_cuda, is_hip, is_npu, is_xpu
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
@@ -148,7 +146,7 @@ class TestDSAHiCacheTransfer(unittest.TestCase):
@unittest.skipIf( @unittest.skipIf(
is_hip(), is_hip(),
'`io_backend="kernel"` path in memory_pool_host.backup_from_device_all_layer ' '`io_backend="kernel"` path in MLATokenToKVPoolHost.backup_from_device_all_layer '
"raises ValueError on AMD (only the `direct` IO backend is wired for ROCm). " "raises ValueError on AMD (only the `direct` IO backend is wired for ROCm). "
"The other 62 tests in this file pass on AMD.", "The other 62 tests in this file pass on AMD.",
) )
test/registered/unit/mem_cache/test_hicache_staged_write_back_dispatch.py
+11 -8
View File
@@ -25,16 +25,17 @@ from sglang.srt.mem_cache.memory_pool_host import (
HostPoolGroup, HostPoolGroup,
LogicalHostPool, LogicalHostPool,
MambaPoolHost, MambaPoolHost,
MLATokenToKVPoolHost,
PoolEntry, PoolEntry,
) )
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.test.ci.ci_register import register_cpu_ci from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(est_time=3, suite="base-a-test-cpu") register_cpu_ci(est_time=3, suite="base-a-test-cpu")
MEMORY_POOL_HOST_MODULE = "sglang.srt.mem_cache.memory_pool_host" MEMORY_POOL_HOST_MODULE = "sglang.srt.mem_cache.memory_pool_host"
MHA_POOL_HOST_MODULE = "sglang.srt.mem_cache.pool_host.mha" MHA_POOL_HOST_MODULE = "sglang.srt.mem_cache.pool_host.mha"
MLA_POOL_HOST_MODULE = "sglang.srt.mem_cache.pool_host.mla"
def _indices(start: int, end: int) -> torch.Tensor: def _indices(start: int, end: int) -> torch.Tensor:
@@ -155,7 +156,7 @@ class _FakeDeviceModule:
class TestHiCacheStagedWriteBackDispatch(unittest.TestCase): class TestHiCacheStagedWriteBackDispatch(unittest.TestCase):
def _patched_transfers(self, src_registry=None): def _patched_transfers(self, src_registry=None, module=MEMORY_POOL_HOST_MODULE):
staged_side_effect = None staged_side_effect = None
if src_registry is not None: if src_registry is not None:
staged_side_effect = lambda **kwargs: _cpu_staged_lf_pf_copy( staged_side_effect = lambda **kwargs: _cpu_staged_lf_pf_copy(
@@ -163,15 +164,15 @@ class TestHiCacheStagedWriteBackDispatch(unittest.TestCase):
) )
return ( return (
mock.patch( mock.patch(
f"{MEMORY_POOL_HOST_MODULE}.jit_transfer_hicache_all_layer_mla_staged_lf_pf", f"{module}.jit_transfer_hicache_all_layer_mla_staged_lf_pf",
side_effect=staged_side_effect, side_effect=staged_side_effect,
), ),
mock.patch( mock.patch(
f"{MEMORY_POOL_HOST_MODULE}.transfer_kv_all_layer_mla_lf_pf", f"{module}.transfer_kv_all_layer_mla_lf_pf",
create=True, create=True,
), ),
mock.patch( mock.patch(
f"{MEMORY_POOL_HOST_MODULE}.transfer_kv_per_layer_mla_pf_lf", f"{module}.transfer_kv_per_layer_mla_pf_lf",
side_effect=_cpu_per_layer_pf_lf_copy, side_effect=_cpu_per_layer_pf_lf_copy,
create=True, create=True,
), ),
@@ -329,16 +330,18 @@ class TestHiCacheStagedWriteBackDispatch(unittest.TestCase):
) )
src_registry = {_ptr_key_from_layers(device_layers): device_layers} src_registry = {_ptr_key_from_layers(device_layers): device_layers}
staged_patch, fallback_patch, _ = self._patched_transfers(src_registry) staged_patch, fallback_patch, _ = self._patched_transfers(
src_registry, module=MLA_POOL_HOST_MODULE
)
with ( with (
staged_patch as staged, staged_patch as staged,
fallback_patch as fallback, fallback_patch as fallback,
mock.patch( mock.patch(
f"{MEMORY_POOL_HOST_MODULE}.jit_transfer_hicache_one_layer_mla", f"{MLA_POOL_HOST_MODULE}.jit_transfer_hicache_one_layer_mla",
side_effect=_cpu_jit_one_layer_mla_copy, side_effect=_cpu_jit_one_layer_mla_copy,
) as load, ) as load,
mock.patch( mock.patch(
f"{MEMORY_POOL_HOST_MODULE}.can_use_write_back_jit_kernel", f"{MLA_POOL_HOST_MODULE}.can_use_write_back_jit_kernel",
return_value=True, return_value=True,
) as can_use_write_back_jit_kernel, ) as can_use_write_back_jit_kernel,
): ):
test/registered/unit/mem_cache/test_mooncake_group_semantics.py
+5 -5
View File
@@ -42,14 +42,14 @@ def _fake_mooncake_modules(fake_store_cls, replicate_config_cls):
} }
def _fake_memory_pool_host_module(): def _fake_pool_host_mla_module():
memory_pool_host = types.ModuleType("sglang.srt.mem_cache.memory_pool_host") pool_host_mla = types.ModuleType("sglang.srt.mem_cache.pool_host.mla")
class MLATokenToKVPoolHost: class MLATokenToKVPoolHost:
pass pass
memory_pool_host.MLATokenToKVPoolHost = MLATokenToKVPoolHost pool_host_mla.MLATokenToKVPoolHost = MLATokenToKVPoolHost
return memory_pool_host return pool_host_mla
def _fake_pool_host_module(): def _fake_pool_host_module():
@@ -68,8 +68,8 @@ def _fake_pool_host_module():
def _fake_host_pool_modules(): def _fake_host_pool_modules():
return { return {
"sglang.srt.mem_cache.memory_pool_host": _fake_memory_pool_host_module(),
"sglang.srt.mem_cache.pool_host": _fake_pool_host_module(), "sglang.srt.mem_cache.pool_host": _fake_pool_host_module(),
"sglang.srt.mem_cache.pool_host.mla": _fake_pool_host_mla_module(),
} }