diff --git a/python/sglang/srt/disaggregation/decode_kvcache_offload_manager.py b/python/sglang/srt/disaggregation/decode_kvcache_offload_manager.py index 9a1f7eaec..b02109f2a 100644 --- a/python/sglang/srt/disaggregation/decode_kvcache_offload_manager.py +++ b/python/sglang/srt/disaggregation/decode_kvcache_offload_manager.py @@ -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 diff --git a/python/sglang/srt/managers/hisparse_coordinator.py b/python/sglang/srt/managers/hisparse_coordinator.py index 5e8f1bf8c..0b9ac7190 100644 --- a/python/sglang/srt/managers/hisparse_coordinator.py +++ b/python/sglang/srt/managers/hisparse_coordinator.py @@ -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() diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index 733479550..2bf611f99 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/python/sglang/srt/mem_cache/hiradix_cache.py @@ -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, diff --git a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py index cb5ee0768..18b314bd9 100644 --- a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py +++ b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py @@ -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: diff --git a/python/sglang/srt/mem_cache/kv_cache_builder.py b/python/sglang/srt/mem_cache/kv_cache_builder.py index d4e4cda98..24bb32796 100644 --- a/python/sglang/srt/mem_cache/kv_cache_builder.py +++ b/python/sglang/srt/mem_cache/kv_cache_builder.py @@ -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): diff --git a/python/sglang/srt/mem_cache/memory_pool_host.py b/python/sglang/srt/mem_cache/memory_pool_host.py index e1bd76477..5d05fa666 100644 --- a/python/sglang/srt/mem_cache/memory_pool_host.py +++ b/python/sglang/srt/mem_cache/memory_pool_host.py @@ -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__( diff --git a/python/sglang/srt/mem_cache/pool_host/mla.py b/python/sglang/srt/mem_cache/pool_host/mla.py new file mode 100644 index 000000000..943395a50 --- /dev/null +++ b/python/sglang/srt/mem_cache/pool_host/mla.py @@ -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 diff --git a/python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py b/python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py index e12841165..d72794cbc 100644 --- a/python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py +++ b/python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py @@ -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 diff --git a/test/registered/jit/test_hicache.py b/test/registered/jit/test_hicache.py index b877e0b96..6cebf2ec4 100644 --- a/test/registered/jit/test_hicache.py +++ b/test/registered/jit/test_hicache.py @@ -5,12 +5,12 @@ import torch 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_host import MLATokenToKVPoolHost from sglang.srt.mem_cache.pool_host.common import ( ALLOC_MEMORY_FUNCS, alloc_with_pin_memory, ) 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.test.ci.ci_register import register_cuda_ci diff --git a/test/registered/jit/test_hicache_page_first_write_back.py b/test/registered/jit/test_hicache_page_first_write_back.py index 55dd6ca82..3d6897e10 100644 --- a/test/registered/jit/test_hicache_page_first_write_back.py +++ b/test/registered/jit/test_hicache_page_first_write_back.py @@ -15,12 +15,12 @@ import torch 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_host import MLATokenToKVPoolHost from sglang.srt.mem_cache.pool_host.common import ( ALLOC_MEMORY_FUNCS, alloc_with_pin_memory, ) 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.test.ci.ci_register import register_amd_ci, register_cuda_ci diff --git a/test/registered/unit/mem_cache/test_dsa_pool_host_unit.py b/test/registered/unit/mem_cache/test_dsa_pool_host_unit.py index 1c6a4825d..1e92cd4e5 100644 --- a/test/registered/unit/mem_cache/test_dsa_pool_host_unit.py +++ b/test/registered/unit/mem_cache/test_dsa_pool_host_unit.py @@ -4,14 +4,12 @@ import unittest import torch from sglang.srt.mem_cache.memory_pool import DSATokenToKVPool -from sglang.srt.mem_cache.memory_pool_host import ( - DSAIndexerPoolHost, - MLATokenToKVPoolHost, -) +from sglang.srt.mem_cache.memory_pool_host import DSAIndexerPoolHost from sglang.srt.mem_cache.pool_host.common import ( ALLOC_MEMORY_FUNCS, 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.test.ci.ci_register import register_amd_ci, register_cuda_ci @@ -148,7 +146,7 @@ class TestDSAHiCacheTransfer(unittest.TestCase): @unittest.skipIf( 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). " "The other 62 tests in this file pass on AMD.", ) diff --git a/test/registered/unit/mem_cache/test_hicache_staged_write_back_dispatch.py b/test/registered/unit/mem_cache/test_hicache_staged_write_back_dispatch.py index ac3a1b650..8c79c3565 100644 --- a/test/registered/unit/mem_cache/test_hicache_staged_write_back_dispatch.py +++ b/test/registered/unit/mem_cache/test_hicache_staged_write_back_dispatch.py @@ -25,16 +25,17 @@ from sglang.srt.mem_cache.memory_pool_host import ( HostPoolGroup, LogicalHostPool, MambaPoolHost, - MLATokenToKVPoolHost, PoolEntry, ) 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 register_cpu_ci(est_time=3, suite="base-a-test-cpu") MEMORY_POOL_HOST_MODULE = "sglang.srt.mem_cache.memory_pool_host" 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: @@ -155,7 +156,7 @@ class _FakeDeviceModule: 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 if src_registry is not None: staged_side_effect = lambda **kwargs: _cpu_staged_lf_pf_copy( @@ -163,15 +164,15 @@ class TestHiCacheStagedWriteBackDispatch(unittest.TestCase): ) return ( 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, ), 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, ), 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, create=True, ), @@ -329,16 +330,18 @@ class TestHiCacheStagedWriteBackDispatch(unittest.TestCase): ) 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 ( staged_patch as staged, fallback_patch as fallback, 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, ) as load, 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, ) as can_use_write_back_jit_kernel, ): diff --git a/test/registered/unit/mem_cache/test_mooncake_group_semantics.py b/test/registered/unit/mem_cache/test_mooncake_group_semantics.py index e41ae493b..456090c7a 100644 --- a/test/registered/unit/mem_cache/test_mooncake_group_semantics.py +++ b/test/registered/unit/mem_cache/test_mooncake_group_semantics.py @@ -42,14 +42,14 @@ def _fake_mooncake_modules(fake_store_cls, replicate_config_cls): } -def _fake_memory_pool_host_module(): - memory_pool_host = types.ModuleType("sglang.srt.mem_cache.memory_pool_host") +def _fake_pool_host_mla_module(): + pool_host_mla = types.ModuleType("sglang.srt.mem_cache.pool_host.mla") class MLATokenToKVPoolHost: pass - memory_pool_host.MLATokenToKVPoolHost = MLATokenToKVPoolHost - return memory_pool_host + pool_host_mla.MLATokenToKVPoolHost = MLATokenToKVPoolHost + return pool_host_mla def _fake_pool_host_module(): @@ -68,8 +68,8 @@ def _fake_pool_host_module(): def _fake_host_pool_modules(): 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.mla": _fake_pool_host_mla_module(), }