diff --git a/python/sglang/jit_kernel/hicache.py b/python/sglang/jit_kernel/hicache.py index b4b80bdb7..268eabe18 100644 --- a/python/sglang/jit_kernel/hicache.py +++ b/python/sglang/jit_kernel/hicache.py @@ -26,14 +26,33 @@ def _jit_hicache_module(*, element_size: int, unroll: int, block_quota: int) -> *args, cuda_files=[ "kvcacheio/hicache.cuh", - "kvcacheio/relayout.cuh", - "kvcacheio/staged_write_back.cuh", ], cuda_wrappers=[ ("launch_one", f"&HiCacheKernel<{args}>::run_one"), ("launch_all", f"&HiCacheKernel<{args}>::run_all"), ("launch_one_mla", f"&HiCacheKernel<{args}>::run_one_mla"), ("launch_all_mla", f"&HiCacheKernel<{args}>::run_all_mla"), + ], + ) + + +@cache_once +def _jit_hicache_staged_module( + *, element_size: int, unroll: int, block_quota: int +) -> Module: + args = make_cpp_args( + element_size, + unroll, + block_quota, + 1024, # num_threads, kept for template compatibility + ) + return load_jit( + "hicache_staged", + *args, + cuda_files=[ + "kvcacheio/staged_write_back.cuh", + ], + cuda_wrappers=[ ( "launch_all_lf_pf_staged", f"&HiCacheStagedWriteBackKernel<{args}>::run_all_lf_pf_staged", @@ -70,6 +89,30 @@ def can_use_hicache_jit_kernel( return False +def can_use_write_back_jit_kernel( + *, + element_size: int, + unroll: int | None = None, # can be tuned for performance + block_quota: int | None = None, # can be tuned for less interference +) -> bool: + logger = logging.getLogger(__name__) + if element_size % 16 != 0: + logger.warning(f"Unsupported {element_size = } for staged JIT HiCache kernel") + return False + try: + unroll = unroll or _default_unroll(element_size) + block_quota = block_quota or DEFAULT_BLOCK_QUOTA + _jit_hicache_staged_module( + element_size=element_size, + unroll=unroll, + block_quota=block_quota, + ) + return True + except Exception as e: + logger.warning(f"Failed to load staged JIT HiCache kernel: {e}") + return False + + def _default_unroll(element_size: int) -> int: if element_size <= 512: return 4 @@ -238,7 +281,7 @@ def transfer_hicache_all_layer_staged_lf_pf( block_quota = block_quota or DEFAULT_BLOCK_QUOTA unroll = unroll or _default_unroll(element_size) src_page_indices = src_indices[::page_size].contiguous() - module = _jit_hicache_module( + module = _jit_hicache_staged_module( element_size=element_size, unroll=unroll, block_quota=block_quota, @@ -284,7 +327,7 @@ def transfer_hicache_all_layer_mla_staged_lf_pf( block_quota = block_quota or DEFAULT_BLOCK_QUOTA unroll = unroll or _default_unroll(element_size) src_page_indices = src_indices[::page_size].contiguous() - module = _jit_hicache_module( + module = _jit_hicache_staged_module( element_size=element_size, unroll=unroll, block_quota=block_quota, diff --git a/python/sglang/srt/managers/cache_controller.py b/python/sglang/srt/managers/cache_controller.py index 0a7f2f297..d5e782aa1 100644 --- a/python/sglang/srt/managers/cache_controller.py +++ b/python/sglang/srt/managers/cache_controller.py @@ -726,10 +726,12 @@ class HiCacheController: return op = CacheOperation.merge_ops(self.write_queue) - # For now, kernel write-back keeps host indices on CPU only for page_first. - # More layouts can use this path once their write-back kernels accept CPU - # destination indices. - if self.io_backend == "kernel" and self.mem_pool_host.layout == "page_first": + # Page-first write-back JIT kernels can keep destination host indices on CPU. + if ( + self.io_backend == "kernel" + and self.mem_pool_host.layout == "page_first" + and getattr(self.mem_pool_host, "can_use_write_back_jit", False) + ): host_indices, device_indices = op.host_indices, op.device_indices else: host_indices, device_indices = self.move_indices( diff --git a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py index 6a3c617e2..011a877e7 100644 --- a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py +++ b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py @@ -394,10 +394,12 @@ class HybridCacheController(BaseHiCacheController): if not self.write_queue: return op = CacheOperation.merge_ops(self.write_queue) - # For now, kernel write-back keeps host indices on CPU only for page_first. - # More layouts can use this path once their write-back kernels accept CPU - # destination indices. - if self.io_backend == "kernel" and self.mem_pool_host.layout == "page_first": + # Page-first write-back JIT kernels can keep destination host indices on CPU. + if ( + self.io_backend == "kernel" + and self.mem_pool_host.layout == "page_first" + and getattr(self.mem_pool_host, "can_use_write_back_jit", False) + ): host_indices = op.host_indices device_indices = op.device_indices resolved_pool_transfers = op.pool_transfers 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 8c94c53b8..daaf05e63 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 @@ -321,7 +321,9 @@ def build_deepseek_v4_hicache_stack( swa_page_size=kvcache.swa_page_size, ) - logical_host_pool = LogicalHostPool(num_host_pages * page_size, page_size) + logical_host_pool = LogicalHostPool( + num_host_pages * page_size, page_size, layout=server_args.hicache_mem_layout + ) swa_host_pool = DeepSeekV4PagedHostPool( pool_name=str(PoolName.SWA), device_buffers=kvcache.swa_kv_pool.kv_buffer, diff --git a/python/sglang/srt/mem_cache/memory_pool_host.py b/python/sglang/srt/mem_cache/memory_pool_host.py index 67e1eda1d..ed1718f86 100644 --- a/python/sglang/srt/mem_cache/memory_pool_host.py +++ b/python/sglang/srt/mem_cache/memory_pool_host.py @@ -17,6 +17,7 @@ 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 as jit_transfer_hicache_all_layer, @@ -257,6 +258,7 @@ class HostKVCache(abc.ABC): self.pin_memory = pin_memory self.device = device self.allocator = get_allocator_from_storage(allocator_type) + self.can_use_write_back_jit = False self.dtype = device_pool.store_dtype self.size_per_token = self.get_size_per_token() @@ -490,9 +492,16 @@ class MHATokenToKVPoolHost(HostKVCache): self.staging_token_capacity = 0 self.staging_k_buffer = None self.staging_v_buffer = None + self.can_use_write_back_jit = False if self.layout != "page_first" or (_is_npu or _is_xpu or _is_mps): return + self.can_use_write_back_jit = _is_cuda and can_use_write_back_jit_kernel( + element_size=self.element_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_k_buffer = torch.empty( @@ -661,7 +670,7 @@ class MHATokenToKVPoolHost(HostKVCache): num_layers=self.layer_num, ) elif self.layout == "page_first": - if self.can_use_jit: + if self.can_use_write_back_jit: jit_transfer_hicache_all_layer_staged_lf_pf( k_ptr_src=device_pool.k_data_ptrs, v_ptr_src=device_pool.v_data_ptrs, @@ -1270,7 +1279,7 @@ class MLATokenToKVPoolHost(HiSparseHostPoolMixin, HostKVCache): element_size=self.kv_cache_dim * self.dtype.itemsize ) - if self.layout == "page_first" and self.can_use_jit: + 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) @@ -1382,9 +1391,16 @@ class MLATokenToKVPoolHost(HiSparseHostPoolMixin, HostKVCache): 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 + self.can_use_write_back_jit = _is_cuda 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( @@ -1506,7 +1522,7 @@ class MLATokenToKVPoolHost(HiSparseHostPoolMixin, HostKVCache): num_layers=self.layer_num, ) elif self.layout == "page_first": - if self.can_use_jit: + 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, @@ -1747,7 +1763,25 @@ class MambaPoolHost(HostKVCache): self.layout, ) + self.temporal_device_ptrs = torch.tensor( + [ + device_pool.mamba_cache.temporal[i].data_ptr() + for i in range(self.num_mamba_layers) + ], + dtype=torch.uint64, + device=self.device_pool.device, + ) + self.conv_device_ptrs = [ + torch.tensor( + [conv_state[i].data_ptr() for i in range(self.num_mamba_layers)], + dtype=torch.uint64, + device=self.device_pool.device, + ) + for conv_state in device_pool.mamba_cache.conv + ] + self.init_kv_buffer() + self._init_write_back_staging_buffers() self.lock = threading.RLock() self.clear() @@ -1806,6 +1840,54 @@ class MambaPoolHost(HostKVCache): ) ) + def _init_write_back_staging_buffers(self): + self.temporal_staging_buffer = None + self.conv_staging_buffers = [None] * len(self.conv_buffer) + self.can_use_write_back_jit = False + self._temporal_can_use_jit = False + self._conv_can_use_jit = [False] * len(self.conv_buffer) + if self.layout != "page_first" or (_is_npu or _is_xpu or _is_mps): + return + + self._temporal_can_use_jit = _is_cuda and can_use_write_back_jit_kernel( + element_size=self._item_size_per_index(self.temporal_buffer[0]), + ) + self._conv_can_use_jit = [ + _is_cuda + and can_use_write_back_jit_kernel( + element_size=self._item_size_per_index(buf[0]), + ) + for buf in self.conv_buffer + ] + self.can_use_write_back_jit = self._temporal_can_use_jit and all( + self._conv_can_use_jit + ) + 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.temporal_staging_buffer = torch.empty( + ( + self.staging_token_capacity, + self.num_mamba_layers, + 1, + *self.temporal_state_shape, + ), + dtype=self.temporal_dtype, + device=self.device_pool.device, + ) + self.conv_staging_buffers = [ + torch.empty( + ( + self.staging_token_capacity, + self.num_mamba_layers, + 1, + *conv_shape, + ), + dtype=self.conv_dtype, + device=self.device_pool.device, + ) + for conv_shape in self.conv_state_shapes + ] + def get_hybrid_pool_buffer(self): # Expose all mamba host tensors that need Mooncake buffer registration. return [self.temporal_buffer, *self.conv_buffer] @@ -1941,27 +2023,35 @@ class MambaPoolHost(HostKVCache): src_indices: torch.Tensor, dst_indices: torch.Tensor, num_layers: int, - device: str, io_backend: str, + src_ptrs: torch.Tensor, + staging: Optional[torch.Tensor] = None, + can_use_jit: bool = False, ) -> None: if src_indices.numel() == 0: return if io_backend == "kernel": item_size = MambaPoolHost._item_size_per_index(src_layers[0]) - src_ptrs = torch.tensor( - [src_layers[i].data_ptr() for i in range(num_layers)], - dtype=torch.uint64, - device=device, - ) - transfer_kv_all_layer_mla_lf_pf( - src_layers=src_ptrs, - dst=dst, - src_indices=src_indices, - dst_indices=dst_indices, - item_size=item_size, - dst_layout_dim=item_size * num_layers, - num_layers=num_layers, - ) + if can_use_jit: + jit_transfer_hicache_all_layer_mla_staged_lf_pf( + ptr_src=src_ptrs, + src_indices=src_indices, + dst_indices=dst_indices, + staging=staging, + dst=dst, + page_size=1, + element_size=item_size, + ) + else: + transfer_kv_all_layer_mla_lf_pf( + src_layers=src_ptrs, + dst=dst, + src_indices=src_indices, + dst_indices=dst_indices, + item_size=item_size, + dst_layout_dim=item_size * num_layers, + num_layers=num_layers, + ) elif io_backend == "direct": src_ptrs = [src_layers[i] for i in range(num_layers)] transfer_kv_all_layer_direct_lf_pf( @@ -2029,8 +2119,10 @@ class MambaPoolHost(HostKVCache): src_indices=device_indices, dst_indices=host_indices, num_layers=self.num_mamba_layers, - device=self.device_pool.device, io_backend=io_backend, + staging=self.temporal_staging_buffer, + can_use_jit=self._temporal_can_use_jit, + src_ptrs=self.temporal_device_ptrs, ) for conv_idx in range(len(self.conv_state_shapes)): self._copy_tensor_all_layers_lf_pf( @@ -2039,8 +2131,10 @@ class MambaPoolHost(HostKVCache): src_indices=device_indices, dst_indices=host_indices, num_layers=self.num_mamba_layers, - device=self.device_pool.device, io_backend=io_backend, + staging=self.conv_staging_buffers[conv_idx], + can_use_jit=self._conv_can_use_jit[conv_idx], + src_ptrs=self.conv_device_ptrs[conv_idx], ) else: for layer_id in range(self.num_mamba_layers): @@ -2161,7 +2255,7 @@ class LogicalHostPool: compressed side pools use these logical FULL indices as stable page anchors. """ - def __init__(self, size: int, page_size: int): + def __init__(self, size: int, page_size: int, layout: str = "layer_first"): if size % page_size != 0: raise ValueError( "LogicalHostPool size must be page-aligned, " @@ -2170,7 +2264,7 @@ class LogicalHostPool: self.size = size self.page_size = page_size self.device = "cpu" - self.layout = "layer_first" + self.layout = layout self.dtype = torch.uint8 self.layer_num = 0 self.start_layer = 0 @@ -2178,6 +2272,7 @@ class LogicalHostPool: self.kv_buffer = None self.size_per_token = 0 self.allocator = None + self.can_use_write_back_jit = True self.lock = threading.RLock() self.clear() @@ -2342,8 +2437,26 @@ class DeepSeekV4PagedHostPool(HiSparseHostPoolMixin, HostKVCache): if self.data_refs else None ) + self.can_use_jit = False + self.can_use_write_back_jit = False + self._init_write_back_staging_buffers() self.clear() + def _init_write_back_staging_buffers(self): + self.staging_buffer = None + if self.layout != "page_first" or (_is_npu or _is_xpu or _is_mps): + return + + self.can_use_write_back_jit = _is_cuda and can_use_write_back_jit_kernel( + element_size=self.item_bytes * self.dtype.itemsize, + ) + staging_page_capacity = min(self.num_host_pages, _WRITE_BACK_STAGING_PAGE_CHUNK) + self.staging_buffer = torch.empty( + (staging_page_capacity, self.layer_num, self.item_bytes), + dtype=self.dtype, + device=self.gpu_device, + ) + def get_contiguous_buf_infos(self): """Return per-layer page-row buffers for PD direct-to-host transfer.""" data_ptrs = [int(self.data_ptrs[i].item()) for i in range(self.layer_num)] @@ -2434,15 +2547,26 @@ class DeepSeekV4PagedHostPool(HiSparseHostPoolMixin, HostKVCache): num_layers=self.layer_num, ) elif io_backend == "kernel" and self.layout == "page_first": - transfer_kv_all_layer_mla_lf_pf( - src_layers=self.device_ptrs, - dst=self.kv_buffer, - src_indices=device_rows, - dst_indices=host_rows, - item_size=self.item_bytes, - dst_layout_dim=self.layer_num * self.item_bytes, - num_layers=self.layer_num, - ) + if self.can_use_write_back_jit: + jit_transfer_hicache_all_layer_mla_staged_lf_pf( + ptr_src=self.device_ptrs, + src_indices=device_rows, + dst_indices=host_rows, + staging=self.staging_buffer, + dst=self.kv_buffer, + page_size=1, + element_size=self.item_bytes, + ) + else: + transfer_kv_all_layer_mla_lf_pf( + src_layers=self.device_ptrs, + dst=self.kv_buffer, + src_indices=device_rows, + dst_indices=host_rows, + item_size=self.item_bytes, + dst_layout_dim=self.layer_num * self.item_bytes, + num_layers=self.layer_num, + ) elif io_backend == "direct" and self.layout == "layer_first": transfer_kv_direct( src_layers=self.device_buffers, @@ -2690,6 +2814,9 @@ class DeepSeekV4StateHostPool(HostKVCache): if self.data_refs else None ) + self.can_use_jit = False + self.can_use_write_back_jit = False + self._init_write_back_staging_buffers() def _init_device_page_views(self) -> None: expected_ring_size = None @@ -2724,6 +2851,21 @@ class DeepSeekV4StateHostPool(HostKVCache): self.ring_size = expected_ring_size or 0 self.state_page_bytes = expected_state_page_bytes or 0 + def _init_write_back_staging_buffers(self): + self.staging_buffer = None + if self.layout != "page_first" or (_is_npu or _is_xpu or _is_mps): + return + + self.can_use_write_back_jit = _is_cuda and can_use_write_back_jit_kernel( + element_size=self.state_page_bytes * self.dtype.itemsize, + ) + staging_page_capacity = min(self.num_host_pages, _WRITE_BACK_STAGING_PAGE_CHUNK) + self.staging_buffer = torch.empty( + (staging_page_capacity, self.layer_num, self.state_page_bytes), + dtype=self.dtype, + device=self.gpu_device, + ) + def _to_page_indices(self, indices: torch.Tensor) -> torch.Tensor: if indices.numel() % self.swa_page_size != 0: raise ValueError( @@ -2782,15 +2924,26 @@ class DeepSeekV4StateHostPool(HostKVCache): num_layers=self.layer_num, ) elif io_backend == "kernel" and self.layout == "page_first": - transfer_kv_all_layer_mla_lf_pf( - src_layers=self.device_ptrs, - dst=self.kv_buffer, - src_indices=device_rows, - dst_indices=host_rows, - item_size=self.state_page_bytes, - dst_layout_dim=self.layer_num * self.state_page_bytes, - num_layers=self.layer_num, - ) + if self.can_use_write_back_jit: + jit_transfer_hicache_all_layer_mla_staged_lf_pf( + ptr_src=self.device_ptrs, + src_indices=device_rows, + dst_indices=host_rows, + staging=self.staging_buffer, + dst=self.kv_buffer, + page_size=1, + element_size=self.state_page_bytes, + ) + else: + transfer_kv_all_layer_mla_lf_pf( + src_layers=self.device_ptrs, + dst=self.kv_buffer, + src_indices=device_rows, + dst_indices=host_rows, + item_size=self.state_page_bytes, + dst_layout_dim=self.layer_num * self.state_page_bytes, + num_layers=self.layer_num, + ) elif io_backend == "direct" and self.layout == "layer_first": transfer_kv_direct( src_layers=self.device_page_views, @@ -2960,6 +3113,10 @@ class HostPoolGroup: self.page_size = self.anchor_entry.host_pool.page_size self.device = self.anchor_entry.host_pool.device self.size = self.anchor_entry.host_pool.size + self.can_use_write_back_jit = all( + getattr(entry.host_pool, "can_use_write_back_jit", False) + for entry in entries + ) @property def kv_buffer(self): @@ -3141,6 +3298,9 @@ class DSAIndexerPoolHost(HostKVCache): layout, ) self.init_kv_buffer() + self.can_use_jit = False + self.can_use_write_back_jit = False + self._init_write_back_staging_buffers() self.lock = threading.RLock() self.clear() @@ -3191,6 +3351,28 @@ class DSAIndexerPoolHost(HostKVCache): else: raise ValueError(f"Unsupported layout: {self.layout}") + def _init_write_back_staging_buffers(self): + self.staging_buffer = None + if self.layout != "page_first" or (_is_npu or _is_xpu or _is_mps): + return + + self.can_use_write_back_jit = _is_cuda and can_use_write_back_jit_kernel( + element_size=self.indexer_page_stride_size * self.indexer_dtype.itemsize, + ) + staging_page_capacity = min( + self.indexer_page_num, _WRITE_BACK_STAGING_PAGE_CHUNK + ) + self.staging_buffer = torch.empty( + ( + staging_page_capacity, + self.layer_num, + 1, + self.indexer_page_stride_size, + ), + dtype=self.indexer_dtype, + device=self.device_pool.device, + ) + def get_hybrid_pool_buffer(self): return [self.index_k_with_scale_buffer] @@ -3278,15 +3460,26 @@ class DSAIndexerPoolHost(HostKVCache): num_layers=self.layer_num, ) elif self.layout == "page_first": - transfer_kv_all_layer_mla_lf_pf( - src_layers=self.index_k_device_ptrs, - dst=self.index_k_with_scale_buffer, - src_indices=device_page_indices, - dst_indices=host_page_indices, - item_size=self.indexer_page_stride_size, - dst_layout_dim=self.indexer_layout_dim, - num_layers=self.layer_num, - ) + if self.can_use_write_back_jit: + jit_transfer_hicache_all_layer_mla_staged_lf_pf( + ptr_src=self.index_k_device_ptrs, + src_indices=device_page_indices, + dst_indices=host_page_indices, + staging=self.staging_buffer, + dst=self.index_k_with_scale_buffer, + page_size=1, + element_size=self.indexer_page_stride_size, + ) + else: + transfer_kv_all_layer_mla_lf_pf( + src_layers=self.index_k_device_ptrs, + dst=self.index_k_with_scale_buffer, + src_indices=device_page_indices, + dst_indices=host_page_indices, + item_size=self.indexer_page_stride_size, + dst_layout_dim=self.indexer_layout_dim, + num_layers=self.layer_num, + ) else: raise ValueError(f"Unsupported layout: {self.layout}") elif io_backend == "direct": diff --git a/test/registered/jit/test_hicache.py b/test/registered/jit/test_hicache.py index 0d5ae5ba4..7fa17c4cd 100644 --- a/test/registered/jit/test_hicache.py +++ b/test/registered/jit/test_hicache.py @@ -3,6 +3,7 @@ import sys import pytest 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 ( ALLOC_MEMORY_FUNCS, @@ -255,11 +256,15 @@ def _run_page_first_staged_write_back_mha( layout: str, element_dim: int, page_count: int ) -> None: pool_size = PAGE_SIZE * (page_count + 8) + head_num = ( + element_dim // 128 if element_dim >= 128 and element_dim % 128 == 0 else 1 + ) + head_dim = element_dim // head_num device_pool = MHATokenToKVPool( size=pool_size, page_size=PAGE_SIZE, - head_num=element_dim // 128, - head_dim=128, + head_num=head_num, + head_dim=head_dim, dtype=torch.bfloat16, layer_num=NUM_LAYERS, device=DEVICE, @@ -270,7 +275,12 @@ def _run_page_first_staged_write_back_mha( device_pool=device_pool, layout=layout, ) - assert host_pool.can_use_jit + assert can_use_write_back_jit_kernel( + element_size=element_dim * host_pool.dtype.itemsize, + ) + assert host_pool.can_use_write_back_jit + if element_dim * host_pool.dtype.itemsize % 128 != 0: + assert not host_pool.can_use_jit assert host_pool.staging_page_capacity > 0 if page_count > 64: assert host_pool.staging_page_capacity < page_count @@ -297,6 +307,14 @@ def _run_page_first_staged_write_back_mha( device_indices = _token_indices_for_pages(device_pages, dtype=src_index_dtype) host_indices = _token_indices_for_pages(host_pages, device="cpu") assert not host_indices.is_cuda + expected_k = [ + device_pool.k_buffer[layer_id][device_indices.to(dtype=torch.int64)].cpu() + for layer_id in range(NUM_LAYERS) + ] + expected_v = [ + device_pool.v_buffer[layer_id][device_indices.to(dtype=torch.int64)].cpu() + for layer_id in range(NUM_LAYERS) + ] host_pool.backup_from_device_all_layer( device_pool, host_indices, device_indices, "kernel" @@ -329,6 +347,25 @@ def _run_page_first_staged_write_back_mha( _assert_page_filled(host_pool.k_data_refs[layer_id], untouched_page, -7) _assert_page_filled(host_pool.v_data_refs[layer_id], untouched_page, -11) + for layer_id in range(NUM_LAYERS): + device_pool.k_buffer[layer_id].zero_() + device_pool.v_buffer[layer_id].zero_() + load_indices = device_indices.to(dtype=torch.int64) + host_indices_load = _token_indices_for_pages(host_pages) + for layer_id in range(NUM_LAYERS): + host_pool.load_to_device_per_layer( + device_pool, host_indices_load, load_indices, layer_id, "kernel" + ) + torch.cuda.synchronize() + + for layer_id in range(NUM_LAYERS): + assert torch.equal( + device_pool.k_buffer[layer_id][load_indices].cpu(), expected_k[layer_id] + ) + assert torch.equal( + device_pool.v_buffer[layer_id][load_indices].cpu(), expected_v[layer_id] + ) + def _run_page_first_staged_write_back_mla( layout: str, element_dim: int, page_count: int @@ -349,7 +386,12 @@ def _run_page_first_staged_write_back_mla( device_pool=device_pool, layout=layout, ) - assert host_pool.can_use_jit + assert can_use_write_back_jit_kernel( + element_size=element_dim * host_pool.dtype.itemsize, + ) + assert host_pool.can_use_write_back_jit + if element_dim * host_pool.dtype.itemsize % 128 != 0: + assert not host_pool.can_use_jit assert host_pool.staging_page_capacity > 0 if page_count > 64: assert host_pool.staging_page_capacity < page_count @@ -374,6 +416,10 @@ def _run_page_first_staged_write_back_mla( device_indices = _token_indices_for_pages(device_pages, dtype=src_index_dtype) host_indices = _token_indices_for_pages(host_pages, device="cpu") assert not host_indices.is_cuda + expected = [ + device_pool.kv_buffer[layer_id][device_indices.to(dtype=torch.int64)].cpu() + for layer_id in range(NUM_LAYERS) + ] host_pool.backup_from_device_all_layer( device_pool, host_indices, device_indices, "kernel" @@ -397,6 +443,21 @@ def _run_page_first_staged_write_back_mla( for untouched_page in [0, page_count + 1]: _assert_page_filled(host_pool.data_refs[layer_id], untouched_page, -13) + for layer_id in range(NUM_LAYERS): + device_pool.kv_buffer[layer_id].zero_() + load_indices = device_indices.to(dtype=torch.int64) + host_indices_load = _token_indices_for_pages(host_pages) + for layer_id in range(NUM_LAYERS): + host_pool.load_to_device_per_layer( + device_pool, host_indices_load, load_indices, layer_id, "kernel" + ) + torch.cuda.synchronize() + + for layer_id in range(NUM_LAYERS): + assert torch.equal( + device_pool.kv_buffer[layer_id][load_indices].cpu(), expected[layer_id] + ) + @pytest.mark.parametrize("layout", LAYOUTS) @pytest.mark.parametrize("element_dim", MHA_ELEMENT_DIMS) @@ -428,5 +489,13 @@ def test_hicache_page_first_staged_write_back_mla( _run_page_first_staged_write_back_mla(layout, element_dim, page_count) +def test_hicache_page_first_staged_write_back_mha_staged_only_alignment() -> None: + _run_page_first_staged_write_back_mha("page_first", 72, 65) + + +def test_hicache_page_first_staged_write_back_mla_staged_only_alignment() -> None: + _run_page_first_staged_write_back_mla("page_first", 72, 65) + + if __name__ == "__main__": sys.exit(pytest.main([__file__, "-v", "-s"])) diff --git a/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_hicache_pp_kl.py b/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_hicache_pp_kl.py index 9f3da52d3..789c3b704 100644 --- a/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_hicache_pp_kl.py +++ b/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_hicache_pp_kl.py @@ -30,8 +30,8 @@ def _assert_pp_decode_cached_tokens(result, history_len, output_len, label): class TestUnifiedQwen3HiCachePP(UnifiedRadixTreeTestMixin, CustomTestCase): """Qwen3-32B + HiCache + PP + UnifiedRadixCache.""" - hicache_io_backend = "direct" - hicache_mem_layout = "page_first_direct" + hicache_io_backend = "kernel" + hicache_mem_layout = "page_first" max_running_requests = 2 kl_threshold = 0.005 gsm8k_threshold = 0.7 diff --git a/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_cp.py b/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_cp.py index 0a0ab3184..e103e8d73 100644 --- a/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_cp.py +++ b/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_cp.py @@ -19,8 +19,8 @@ QWEN3_30B_MODEL = "Qwen/Qwen3-30B-A3B-FP8" class TestUnifiedQwen3HiCacheCP(UnifiedRadixTreeTestMixin, CustomTestCase): """Qwen3-30B-A3B-FP8 + HiCache + CP + UnifiedRadixCache.""" - hicache_io_backend = "direct" - hicache_mem_layout = "page_first_direct" + hicache_io_backend = "kernel" + hicache_mem_layout = "page_first" max_running_requests = 32 kl_threshold = 0.005 gsm8k_threshold = 0.7 diff --git a/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_dsv4.py b/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_dsv4.py index d83ffb330..14e60dd6c 100644 --- a/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_dsv4.py +++ b/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_dsv4.py @@ -146,9 +146,9 @@ class TestUnifiedDeepSeekV4FlashHiCacheL3(AccuracyTwoPassMixin, CustomTestCase): "--hicache-storage-prefetch-policy", "wait_complete", "--hicache-io-backend", - "direct", + "kernel", "--hicache-mem-layout", - "page_first_direct", + "page_first", "--hicache-storage-backend", "file", "--swa-full-tokens-ratio", @@ -210,9 +210,9 @@ class TestUnifiedDeepSeekV4FlashEagleHiCacheL3(AccuracyTwoPassMixin, CustomTestC "--hicache-storage-prefetch-policy", "wait_complete", "--hicache-io-backend", - "direct", + "kernel", "--hicache-mem-layout", - "page_first_direct", + "page_first", "--hicache-storage-backend", "file", "--enable-cache-report", diff --git a/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_mamba.py b/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_mamba.py index a46a7f2ca..591707712 100644 --- a/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_mamba.py +++ b/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_mamba.py @@ -108,9 +108,9 @@ class TestUnifiedMambaHiCache(UnifiedRadixTreeTestMixin, CustomTestCase): "--hicache-write-policy", "write_through", "--hicache-io-backend", - "direct", + "kernel", "--hicache-mem-layout", - "page_first_direct", + "page_first", "--max-total-tokens", "12000", "--max-mamba-cache-size", @@ -169,9 +169,9 @@ class TestUnifiedMambaHiCacheL3(AccuracyTwoPassMixin, CustomTestCase): "--hicache-storage-prefetch-policy", "wait_complete", "--hicache-io-backend", - "direct", + "kernel", "--hicache-mem-layout", - "page_first_direct", + "page_first", "--hicache-storage-backend", "file", "--max-mamba-cache-size", diff --git a/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_nightly.py b/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_nightly.py index 0b8302439..74e337a9b 100644 --- a/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_nightly.py +++ b/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_nightly.py @@ -170,9 +170,9 @@ class TestGLM5HiRadixCacheL3Accuracy(AccuracyTwoPassMixin, CustomTestCase): "--hicache-storage-prefetch-policy", "wait_complete", "--hicache-io-backend", - "direct", + "kernel", "--hicache-mem-layout", - "page_first_direct", + "page_first", "--hicache-storage-backend", "file", ], @@ -218,9 +218,9 @@ class TestGLM5UnifiedRadixCacheL3Accuracy(AccuracyTwoPassMixin, CustomTestCase): "--hicache-storage-prefetch-policy", "wait_complete", "--hicache-io-backend", - "direct", + "kernel", "--hicache-mem-layout", - "page_first_direct", + "page_first", "--hicache-storage-backend", "file", ], 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 new file mode 100644 index 000000000..3c613338b --- /dev/null +++ b/test/registered/unit/mem_cache/test_hicache_staged_write_back_dispatch.py @@ -0,0 +1,926 @@ +"""Unit tests for HiCache staged write-back host-pool dispatch.""" + +import unittest +from contextlib import contextmanager +from types import SimpleNamespace +from unittest import mock + +import torch + +from sglang.srt.managers import cache_controller as manager_cache_controller +from sglang.srt.managers.cache_controller import CacheOperation as ManagerCacheOperation +from sglang.srt.managers.cache_controller import ( + HiCacheController, +) +from sglang.srt.mem_cache.hicache_storage import PoolName, PoolTransfer +from sglang.srt.mem_cache.hybrid_cache import hybrid_cache_controller +from sglang.srt.mem_cache.hybrid_cache.hybrid_cache_controller import ( + CacheOperation, + HybridCacheController, +) +from sglang.srt.mem_cache.memory_pool_host import ( + DeepSeekV4PagedHostPool, + DeepSeekV4StateHostPool, + DSAIndexerPoolHost, + HostPoolGroup, + LogicalHostPool, + MambaPoolHost, + MHATokenToKVPoolHost, + MLATokenToKVPoolHost, + PoolEntry, +) +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" + + +def _indices(start: int, end: int) -> torch.Tensor: + return torch.arange(start, end, dtype=torch.int64) + + +def _ptr_key_from_layers(src_layers) -> tuple[int, ...]: + return tuple(int(src_layers[i].data_ptr()) for i in range(len(src_layers))) + + +def _ptr_key_from_tensor(ptrs: torch.Tensor) -> tuple[int, ...]: + return tuple(int(ptr) for ptr in ptrs.cpu().tolist()) + + +def _cpu_staged_lf_pf_copy( + src_registry, + *, + ptr_src, + src_indices, + dst_indices, + dst, + **_, +): + src_layers = src_registry[_ptr_key_from_tensor(ptr_src)] + src_indices = src_indices.to(dtype=torch.int64, device="cpu") + dst_indices = dst_indices.to(dtype=torch.int64, device="cpu") + for layer_id, src in enumerate(src_layers): + dst[dst_indices, layer_id] = src[src_indices] + + +def _cpu_staged_mha_lf_pf_copy( + src_registry, + *, + k_ptr_src, + v_ptr_src, + src_indices, + dst_indices, + dst_k, + dst_v, + **_, +): + k_src_layers = src_registry[_ptr_key_from_tensor(k_ptr_src)] + v_src_layers = src_registry[_ptr_key_from_tensor(v_ptr_src)] + src_indices = src_indices.to(dtype=torch.int64, device="cpu") + dst_indices = dst_indices.to(dtype=torch.int64, device="cpu") + for layer_id, (k_src, v_src) in enumerate(zip(k_src_layers, v_src_layers)): + dst_k[dst_indices, layer_id] = k_src[src_indices] + dst_v[dst_indices, layer_id] = v_src[src_indices] + + +def _cpu_jit_one_layer_mha_copy( + *, + k_cache_dst, + v_cache_dst, + k_cache_src, + v_cache_src, + indices_dst, + indices_src, + **_, +): + indices_dst = indices_dst.to(dtype=torch.int64, device="cpu") + indices_src = indices_src.to(dtype=torch.int64, device="cpu") + k_cache_dst[indices_dst] = k_cache_src[indices_src] + v_cache_dst[indices_dst] = v_cache_src[indices_src] + + +def _cpu_jit_one_layer_mla_copy( + *, + cache_dst, + cache_src, + indices_dst, + indices_src, + **_, +): + indices_dst = indices_dst.to(dtype=torch.int64, device="cpu") + indices_src = indices_src.to(dtype=torch.int64, device="cpu") + cache_dst[indices_dst] = cache_src[indices_src] + + +def _cpu_per_layer_pf_lf_copy( + *, + src, + dst, + src_indices, + dst_indices, + layer_id, + **_, +): + src_indices = src_indices.to(dtype=torch.int64, device="cpu") + dst_indices = dst_indices.to(dtype=torch.int64, device="cpu") + dst[dst_indices] = src[src_indices, layer_id] + + +class _FakeEvent: + def record(self): + pass + + def wait(self, stream): + pass + + +class _FakeDeviceModule: + Event = _FakeEvent + + @staticmethod + @contextmanager + def stream(stream): + yield + + +class TestHiCacheStagedWriteBackDispatch(unittest.TestCase): + def _patched_transfers(self, src_registry=None): + staged_side_effect = None + if src_registry is not None: + staged_side_effect = lambda **kwargs: _cpu_staged_lf_pf_copy( + src_registry, **kwargs + ) + return ( + mock.patch( + f"{MEMORY_POOL_HOST_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", + create=True, + ), + mock.patch( + f"{MEMORY_POOL_HOST_MODULE}.transfer_kv_per_layer_mla_pf_lf", + side_effect=_cpu_per_layer_pf_lf_copy, + create=True, + ), + ) + + def test_mha_backup_then_load_roundtrip_uses_staged(self): + layer_num = 2 + head_num = 1 + head_dim = 4 + host_indices = _indices(0, 4) + device_indices = _indices(4, 8) + k_layers = [ + (torch.arange(8 * head_num * head_dim, dtype=torch.uint8) + layer_id * 40) + .reshape(8, head_num, head_dim) + .clone() + for layer_id in range(layer_num) + ] + v_layers = [ + ( + torch.arange(8 * head_num * head_dim, dtype=torch.uint8) + + 100 + + layer_id * 40 + ) + .reshape(8, head_num, head_dim) + .clone() + for layer_id in range(layer_num) + ] + expected_k = [layer[device_indices].clone() for layer in k_layers] + expected_v = [layer[device_indices].clone() for layer in v_layers] + device_pool = SimpleNamespace( + k_buffer=k_layers, + v_buffer=v_layers, + k_data_ptrs=torch.tensor( + [layer.data_ptr() for layer in k_layers], dtype=torch.uint64 + ), + v_data_ptrs=torch.tensor( + [layer.data_ptr() for layer in v_layers], dtype=torch.uint64 + ), + ) + + host = MHATokenToKVPoolHost.__new__(MHATokenToKVPoolHost) + host.layout = "page_first" + host.page_size = 1 + host.layer_num = layer_num + host.head_num = head_num + host.head_dim = head_dim + host.element_dim = head_num * head_dim + host.token_stride_size = host.element_dim + host.layout_dim = host.token_stride_size * layer_num + host.dtype = torch.uint8 + host.can_use_jit = True + host.can_use_write_back_jit = True + host.kv_buffer = torch.zeros( + 2, 8, layer_num, head_num, head_dim, dtype=torch.uint8 + ) + host.k_data_refs = [host.k_buffer.transpose(0, 1)[i] for i in range(layer_num)] + host.v_data_refs = [host.v_buffer.transpose(0, 1)[i] for i in range(layer_num)] + host.staging_k_buffer = torch.empty( + 4, layer_num, head_num, head_dim, dtype=torch.uint8 + ) + host.staging_v_buffer = torch.empty_like(host.staging_k_buffer) + src_registry = { + _ptr_key_from_layers(k_layers): k_layers, + _ptr_key_from_layers(v_layers): v_layers, + } + + with ( + mock.patch( + f"{MEMORY_POOL_HOST_MODULE}.jit_transfer_hicache_all_layer_staged_lf_pf", + side_effect=lambda **kwargs: _cpu_staged_mha_lf_pf_copy( + src_registry, **kwargs + ), + ) as staged, + mock.patch( + f"{MEMORY_POOL_HOST_MODULE}.transfer_kv_all_layer_lf_pf", + create=True, + ) as fallback, + mock.patch( + f"{MEMORY_POOL_HOST_MODULE}.jit_transfer_hicache_one_layer", + side_effect=_cpu_jit_one_layer_mha_copy, + ) as load, + mock.patch( + f"{MEMORY_POOL_HOST_MODULE}.can_use_write_back_jit_kernel", + return_value=True, + ) as can_use_write_back_jit_kernel, + ): + host.backup_from_device_all_layer( + device_pool, host_indices, device_indices, io_backend="kernel" + ) + for layer in k_layers + v_layers: + layer.zero_() + for layer_id in range(layer_num): + host.load_to_device_per_layer( + device_pool, + host_indices, + device_indices, + layer_id, + io_backend="kernel", + ) + + self.assertEqual(staged.call_count, 1) + self.assertEqual(fallback.call_count, 0) + self.assertEqual(load.call_count, layer_num) + can_use_write_back_jit_kernel.assert_not_called() + for layer_id in range(layer_num): + self.assertTrue( + torch.equal(k_layers[layer_id][device_indices], expected_k[layer_id]) + ) + self.assertTrue( + torch.equal(v_layers[layer_id][device_indices], expected_v[layer_id]) + ) + self.assertTrue( + torch.equal(host.k_buffer[host_indices, layer_id], expected_k[layer_id]) + ) + self.assertTrue( + torch.equal(host.v_buffer[host_indices, layer_id], expected_v[layer_id]) + ) + + def test_mla_backup_then_load_roundtrip_uses_staged(self): + layer_num = 2 + kv_cache_dim = 5 + host_indices = _indices(0, 4) + device_indices = _indices(4, 8) + device_layers = [ + (torch.arange(8 * kv_cache_dim, dtype=torch.uint8) + layer_id * 50) + .reshape(8, 1, kv_cache_dim) + .clone() + for layer_id in range(layer_num) + ] + expected = [layer[device_indices].clone() for layer in device_layers] + device_pool = SimpleNamespace( + kv_buffer=device_layers, + data_ptrs=torch.tensor( + [layer.data_ptr() for layer in device_layers], dtype=torch.uint64 + ), + ) + + host = MLATokenToKVPoolHost.__new__(MLATokenToKVPoolHost) + host.layout = "page_first" + host.page_size = 1 + host.layer_num = layer_num + host.kv_cache_dim = kv_cache_dim + host.token_stride_size = kv_cache_dim + host.layout_dim = host.token_stride_size * layer_num + host.dtype = torch.uint8 + host.can_use_jit = True + host.can_use_write_back_jit = True + host.kv_buffer = torch.zeros(8, layer_num, 1, kv_cache_dim, dtype=torch.uint8) + host.data_refs = [host.kv_buffer.transpose(0, 1)[i] for i in range(layer_num)] + host.staging_buffer = torch.empty( + 4, layer_num, 1, kv_cache_dim, dtype=torch.uint8 + ) + src_registry = {_ptr_key_from_layers(device_layers): device_layers} + + staged_patch, fallback_patch, _ = self._patched_transfers(src_registry) + with ( + staged_patch as staged, + fallback_patch as fallback, + mock.patch( + f"{MEMORY_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", + return_value=True, + ) as can_use_write_back_jit_kernel, + ): + host.backup_from_device_all_layer( + device_pool, host_indices, device_indices, io_backend="kernel" + ) + for layer in device_layers: + layer.zero_() + for layer_id in range(layer_num): + host.load_to_device_per_layer( + device_pool, + host_indices, + device_indices, + layer_id, + io_backend="kernel", + ) + + self.assertEqual(staged.call_count, 1) + self.assertEqual(fallback.call_count, 0) + self.assertEqual(load.call_count, layer_num) + can_use_write_back_jit_kernel.assert_not_called() + for layer_id, layer in enumerate(device_layers): + self.assertTrue(torch.equal(layer[device_indices], expected[layer_id])) + self.assertTrue( + torch.equal(host.kv_buffer[host_indices, layer_id], expected[layer_id]) + ) + + def test_mamba_backup_then_load_roundtrip_uses_staged(self): + num_layers = 2 + host_indices = _indices(0, 4) + device_indices = _indices(4, 8) + temporal = torch.arange(num_layers * 8 * 3, dtype=torch.uint8).reshape( + num_layers, 8, 1, 3 + ) + conv = (torch.arange(num_layers * 8 * 2, dtype=torch.uint8) + 97).reshape( + num_layers, 8, 1, 2 + ) + device_pool = SimpleNamespace( + mamba_cache=SimpleNamespace(temporal=temporal.clone(), conv=[conv.clone()]) + ) + expected_temporal = device_pool.mamba_cache.temporal[:, device_indices].clone() + expected_conv = device_pool.mamba_cache.conv[0][:, device_indices].clone() + + host = MambaPoolHost.__new__(MambaPoolHost) + host.layout = "page_first" + host.num_mamba_layers = num_layers + host.device_pool = SimpleNamespace(device="cpu") + host.temporal_buffer = torch.zeros(8, num_layers, 1, 3, dtype=torch.uint8) + host.conv_buffer = [ + torch.zeros(8, num_layers, 1, 2, dtype=torch.uint8), + ] + host.conv_state_shapes = [(2,)] + host.temporal_staging_buffer = torch.empty( + 4, num_layers, 1, 3, dtype=torch.uint8 + ) + host.conv_staging_buffers = [ + torch.empty(4, num_layers, 1, 2, dtype=torch.uint8), + ] + host._temporal_can_use_jit = True + host._conv_can_use_jit = [True] + host.can_use_write_back_jit = True + host.temporal_device_ptrs = torch.tensor( + [layer.data_ptr() for layer in device_pool.mamba_cache.temporal], + dtype=torch.uint64, + ) + host.conv_device_ptrs = [ + torch.tensor( + [layer.data_ptr() for layer in device_pool.mamba_cache.conv[0]], + dtype=torch.uint64, + ) + ] + + src_registry = { + _ptr_key_from_layers(device_pool.mamba_cache.temporal): list( + device_pool.mamba_cache.temporal + ), + _ptr_key_from_layers(device_pool.mamba_cache.conv[0]): list( + device_pool.mamba_cache.conv[0] + ), + } + + staged_patch, fallback_patch, load_patch = self._patched_transfers(src_registry) + with staged_patch as staged, fallback_patch as fallback, load_patch as load: + host.backup_from_device_all_layer( + device_pool, host_indices, device_indices, io_backend="kernel" + ) + device_pool.mamba_cache.temporal.zero_() + device_pool.mamba_cache.conv[0].zero_() + for layer_id in range(num_layers): + host.load_to_device_per_layer( + device_pool, + host_indices, + device_indices, + layer_id, + io_backend="kernel", + ) + + self.assertEqual(staged.call_count, 2) + self.assertEqual(fallback.call_count, 0) + self.assertEqual(load.call_count, 4) + self.assertTrue( + torch.equal( + device_pool.mamba_cache.temporal[:, device_indices], expected_temporal + ) + ) + self.assertTrue( + torch.equal( + device_pool.mamba_cache.conv[0][:, device_indices], expected_conv + ) + ) + + def test_deepseek_v4_paged_pool_backup_then_load_roundtrip_uses_staged(self): + layer_num = 2 + slot_page_size = 2 + host_indices = torch.tensor([0, 1, 4, 5], dtype=torch.int64) + device_indices = torch.tensor([2, 3, 6, 7], dtype=torch.int64) + host_rows = torch.tensor([0, 2], dtype=torch.int64) + device_rows = torch.tensor([1, 3], dtype=torch.int64) + device_buffers = [ + (torch.arange(5 * 4, dtype=torch.uint8) + layer_id * 50).reshape(5, 4) + for layer_id in range(layer_num) + ] + expected = [buffer[device_rows].clone() for buffer in device_buffers] + + host = DeepSeekV4PagedHostPool.__new__(DeepSeekV4PagedHostPool) + host.pool_name = "c4" + host.layout = "page_first" + host.slot_page_size = slot_page_size + host.layer_num = layer_num + host.item_bytes = 4 + host.dtype = torch.uint8 + host.device_buffers = device_buffers + host.device_ptrs = torch.tensor( + [buffer.data_ptr() for buffer in device_buffers], dtype=torch.uint64 + ) + host.kv_buffer = torch.zeros( + 4, host.layer_num, host.item_bytes, dtype=torch.uint8 + ) + host.staging_buffer = torch.empty( + 4, host.layer_num, host.item_bytes, dtype=torch.uint8 + ) + host.can_use_jit = False + host.can_use_write_back_jit = True + src_registry = {_ptr_key_from_layers(device_buffers): device_buffers} + + staged_patch, fallback_patch, load_patch = self._patched_transfers(src_registry) + with staged_patch as staged, fallback_patch as fallback, load_patch as load: + host.backup_from_device_all_layer( + device_pool=None, + host_indices=host_indices, + device_indices=device_indices, + io_backend="kernel", + ) + for buffer in device_buffers: + buffer.zero_() + for layer_id in range(layer_num): + host.load_to_device_per_layer( + device_pool=None, + host_indices=host_indices, + device_indices=device_indices, + layer_id=layer_id, + io_backend="kernel", + ) + + self.assertEqual(staged.call_count, 1) + self.assertEqual(fallback.call_count, 0) + self.assertEqual(load.call_count, layer_num) + for layer_id, buffer in enumerate(device_buffers): + self.assertTrue(torch.equal(buffer[device_rows], expected[layer_id])) + self.assertTrue( + torch.equal(host.kv_buffer[host_rows, layer_id], expected[layer_id]) + ) + + def test_deepseek_v4_state_pool_backup_then_load_roundtrip_uses_staged(self): + layer_num = 2 + swa_page_size = 2 + host_indices = torch.tensor([0, 1, 4, 5], dtype=torch.int64) + device_indices = torch.tensor([2, 3, 6, 7], dtype=torch.int64) + host_rows = torch.tensor([0, 2], dtype=torch.int64) + device_rows = torch.tensor([1, 3], dtype=torch.int64) + device_page_views = [ + (torch.arange(5 * 5, dtype=torch.uint8) + layer_id * 60).reshape(5, 5) + for layer_id in range(layer_num) + ] + expected = [buffer[device_rows].clone() for buffer in device_page_views] + + host = DeepSeekV4StateHostPool.__new__(DeepSeekV4StateHostPool) + host.pool_name = "c4_state" + host.layout = "page_first" + host.swa_page_size = swa_page_size + host.layer_num = layer_num + host.state_page_bytes = 5 + host.dtype = torch.uint8 + host.device_page_views = device_page_views + host.device_ptrs = torch.tensor( + [buffer.data_ptr() for buffer in device_page_views], dtype=torch.uint64 + ) + host.kv_buffer = torch.zeros( + 4, host.layer_num, host.state_page_bytes, dtype=torch.uint8 + ) + host.staging_buffer = torch.empty( + 4, host.layer_num, host.state_page_bytes, dtype=torch.uint8 + ) + host.can_use_jit = False + host.can_use_write_back_jit = True + src_registry = {_ptr_key_from_layers(device_page_views): device_page_views} + + staged_patch, fallback_patch, load_patch = self._patched_transfers(src_registry) + with staged_patch as staged, fallback_patch as fallback, load_patch as load: + host.backup_from_device_all_layer( + device_pool=None, + host_indices=host_indices, + device_indices=device_indices, + io_backend="kernel", + ) + for buffer in device_page_views: + buffer.zero_() + for layer_id in range(layer_num): + host.load_to_device_per_layer( + device_pool=None, + host_indices=host_indices, + device_indices=device_indices, + layer_id=layer_id, + io_backend="kernel", + ) + + self.assertEqual(staged.call_count, 1) + self.assertEqual(fallback.call_count, 0) + self.assertEqual(load.call_count, layer_num) + for layer_id, buffer in enumerate(device_page_views): + self.assertTrue(torch.equal(buffer[device_rows], expected[layer_id])) + self.assertTrue( + torch.equal(host.kv_buffer[host_rows, layer_id], expected[layer_id]) + ) + + def test_dsa_indexer_backup_then_load_roundtrip_uses_staged(self): + layer_num = 2 + page_size = 2 + host_indices = torch.tensor([0, 1, 4, 5], dtype=torch.int64) + device_indices = torch.tensor([2, 3, 6, 7], dtype=torch.int64) + host_page_indices = torch.tensor([0, 2], dtype=torch.int64) + device_page_indices = torch.tensor([1, 3], dtype=torch.int64) + indexer_page_stride_size = 8 + device_layers = [ + ( + torch.arange(5 * indexer_page_stride_size, dtype=torch.uint8) + + layer_id * 70 + ).reshape(5, 1, indexer_page_stride_size) + for layer_id in range(layer_num) + ] + expected = [buffer[device_page_indices].clone() for buffer in device_layers] + device_pool = SimpleNamespace(index_k_with_scale_buffer=device_layers) + + host = DSAIndexerPoolHost.__new__(DSAIndexerPoolHost) + host.layout = "page_first" + host.page_size = page_size + host.layer_num = layer_num + host.indexer_page_stride_size = indexer_page_stride_size + host.indexer_layout_dim = host.layer_num * host.indexer_page_stride_size + host.index_k_device_ptrs = torch.tensor( + [buffer.data_ptr() for buffer in device_layers], dtype=torch.uint64 + ) + host.index_k_with_scale_buffer = torch.zeros( + 4, host.layer_num, 1, host.indexer_page_stride_size, dtype=torch.uint8 + ) + host.staging_buffer = torch.empty( + 4, host.layer_num, 1, host.indexer_page_stride_size, dtype=torch.uint8 + ) + host.can_use_jit = False + host.can_use_write_back_jit = True + src_registry = {_ptr_key_from_layers(device_layers): device_layers} + + staged_patch, fallback_patch, load_patch = self._patched_transfers(src_registry) + with staged_patch as staged, fallback_patch as fallback, load_patch as load: + host.backup_from_device_all_layer( + device_pool=device_pool, + host_indices=host_indices, + device_indices=device_indices, + io_backend="kernel", + ) + for buffer in device_layers: + buffer.zero_() + for layer_id in range(layer_num): + host.load_to_device_per_layer( + device_pool=device_pool, + host_indices=host_indices, + device_indices=device_indices, + layer_id=layer_id, + io_backend="kernel", + ) + + self.assertEqual(staged.call_count, 1) + self.assertEqual(fallback.call_count, 0) + self.assertEqual(load.call_count, layer_num) + for layer_id, buffer in enumerate(device_layers): + self.assertTrue( + torch.equal(buffer[device_page_indices], expected[layer_id]) + ) + self.assertTrue( + torch.equal( + host.index_k_with_scale_buffer[host_page_indices, layer_id], + expected[layer_id], + ) + ) + + def test_logical_host_pool_preserves_page_first_group_layout(self): + logical_host_pool = LogicalHostPool(8, 2, layout="page_first") + group = HostPoolGroup( + [ + PoolEntry( + name=PoolName.KV, + host_pool=logical_host_pool, + device_pool=None, + layer_mapper=lambda _: 0, + is_primary_index_anchor=True, + ) + ] + ) + + self.assertEqual(group.layout, "page_first") + self.assertTrue(group.can_use_write_back_jit) + + def test_write_back_jit_hybrid_write_keeps_extra_host_indices_on_cpu(self): + captured = {} + + class FakeHostGroup: + layout = "page_first" + can_use_write_back_jit = True + + def backup_from_device_all_layer( + self, + device_pool, + host_indices, + device_indices, + io_backend, + pool_transfers=None, + ): + captured["host_indices"] = host_indices + captured["pool_transfers"] = pool_transfers + + controller = HybridCacheController.__new__(HybridCacheController) + controller.write_queue = [ + CacheOperation( + host_indices=_indices(0, 4), + device_indices=_indices(4, 8), + node_id=1, + pool_transfers=[ + PoolTransfer( + name=PoolName.DEEPSEEK_V4_C4, + host_indices=_indices(0, 4), + device_indices=_indices(4, 8), + ) + ], + ) + ] + controller.io_backend = "kernel" + controller.mem_pool_host = FakeHostGroup() + controller.mem_pool_device = None + controller.has_draft = False + controller.write_stream = object() + controller.ack_write_queue = [] + controller._record_transfer_indices_on_stream = lambda *args: None + controller.move_hybrid_indices = mock.Mock( + side_effect=AssertionError( + "write-back JIT kernel write should not move indices" + ) + ) + + with mock.patch.object( + hybrid_cache_controller, "device_module", _FakeDeviceModule + ): + controller.start_writing() + + controller.move_hybrid_indices.assert_not_called() + self.assertEqual(captured["host_indices"].device.type, "cpu") + self.assertEqual(captured["pool_transfers"][0].host_indices.device.type, "cpu") + + def test_hybrid_write_moves_indices_without_write_back_jit(self): + captured = {} + + class FakeHostGroup: + layout = "page_first" + can_use_write_back_jit = False + + def backup_from_device_all_layer( + self, + device_pool, + host_indices, + device_indices, + io_backend, + pool_transfers=None, + ): + captured["host_indices"] = host_indices + captured["pool_transfers"] = pool_transfers + + op = CacheOperation( + host_indices=_indices(0, 4), + device_indices=_indices(4, 8), + node_id=1, + pool_transfers=[ + PoolTransfer( + name=PoolName.DEEPSEEK_V4_C4, + host_indices=_indices(0, 4), + device_indices=_indices(4, 8), + ) + ], + ) + controller = HybridCacheController.__new__(HybridCacheController) + controller.write_queue = [op] + controller.io_backend = "kernel" + controller.mem_pool_host = FakeHostGroup() + controller.mem_pool_device = None + controller.has_draft = False + controller.write_stream = object() + controller.ack_write_queue = [] + controller._record_transfer_indices_on_stream = lambda *args: None + controller.move_hybrid_indices = mock.Mock( + return_value=(op.host_indices, op.device_indices, op.pool_transfers) + ) + + with mock.patch.object( + hybrid_cache_controller, "device_module", _FakeDeviceModule + ): + controller.start_writing() + + controller.move_hybrid_indices.assert_called_once() + self.assertEqual(captured["host_indices"].device.type, "cpu") + self.assertEqual(captured["pool_transfers"][0].host_indices.device.type, "cpu") + + def test_hybrid_write_moves_indices_without_page_first_layout(self): + captured = {} + + class FakeHostGroup: + layout = "layer_first" + can_use_write_back_jit = True + + def backup_from_device_all_layer( + self, + device_pool, + host_indices, + device_indices, + io_backend, + pool_transfers=None, + ): + captured["host_indices"] = host_indices + captured["pool_transfers"] = pool_transfers + + op = CacheOperation( + host_indices=_indices(0, 4), + device_indices=_indices(4, 8), + node_id=1, + pool_transfers=[ + PoolTransfer( + name=PoolName.DEEPSEEK_V4_C4, + host_indices=_indices(0, 4), + device_indices=_indices(4, 8), + ) + ], + ) + controller = HybridCacheController.__new__(HybridCacheController) + controller.write_queue = [op] + controller.io_backend = "kernel" + controller.mem_pool_host = FakeHostGroup() + controller.mem_pool_device = None + controller.has_draft = False + controller.write_stream = object() + controller.ack_write_queue = [] + controller._record_transfer_indices_on_stream = lambda *args: None + controller.move_hybrid_indices = mock.Mock( + return_value=(op.host_indices, op.device_indices, op.pool_transfers) + ) + + with mock.patch.object( + hybrid_cache_controller, "device_module", _FakeDeviceModule + ): + controller.start_writing() + + controller.move_hybrid_indices.assert_called_once() + self.assertEqual(captured["host_indices"].device.type, "cpu") + self.assertEqual(captured["pool_transfers"][0].host_indices.device.type, "cpu") + + def test_write_back_jit_cache_controller_keeps_host_indices_on_cpu(self): + captured = {} + + class FakeHostPool: + layout = "page_first" + can_use_write_back_jit = True + + def backup_from_device_all_layer( + self, device_pool, host_indices, device_indices, io_backend + ): + captured["host_indices"] = host_indices + + controller = HiCacheController.__new__(HiCacheController) + controller.write_queue = [ + ManagerCacheOperation( + host_indices=_indices(0, 4), + device_indices=_indices(4, 8), + node_id=1, + ) + ] + controller.io_backend = "kernel" + controller.mem_pool_host = FakeHostPool() + controller.mem_pool_device = None + controller.has_draft = False + controller.write_stream = object() + controller.ack_write_queue = [] + controller.move_indices = mock.Mock( + side_effect=AssertionError( + "write-back JIT kernel write should not move indices" + ) + ) + + with mock.patch.object( + manager_cache_controller, "device_module", _FakeDeviceModule + ): + controller.start_writing() + + controller.move_indices.assert_not_called() + self.assertEqual(captured["host_indices"].device.type, "cpu") + + def test_cache_controller_moves_indices_without_write_back_jit(self): + captured = {} + + class FakeHostPool: + layout = "page_first" + can_use_write_back_jit = False + + def backup_from_device_all_layer( + self, device_pool, host_indices, device_indices, io_backend + ): + captured["host_indices"] = host_indices + + op = ManagerCacheOperation( + host_indices=_indices(0, 4), + device_indices=_indices(4, 8), + node_id=1, + ) + controller = HiCacheController.__new__(HiCacheController) + controller.write_queue = [op] + controller.io_backend = "kernel" + controller.mem_pool_host = FakeHostPool() + controller.mem_pool_device = None + controller.has_draft = False + controller.write_stream = object() + controller.ack_write_queue = [] + controller.move_indices = mock.Mock( + return_value=(op.host_indices, op.device_indices) + ) + + with mock.patch.object( + manager_cache_controller, "device_module", _FakeDeviceModule + ): + controller.start_writing() + + controller.move_indices.assert_called_once() + self.assertEqual(captured["host_indices"].device.type, "cpu") + + def test_cache_controller_moves_indices_without_page_first_layout(self): + captured = {} + + class FakeHostPool: + layout = "layer_first" + can_use_write_back_jit = True + + def backup_from_device_all_layer( + self, device_pool, host_indices, device_indices, io_backend + ): + captured["host_indices"] = host_indices + + op = ManagerCacheOperation( + host_indices=_indices(0, 4), + device_indices=_indices(4, 8), + node_id=1, + ) + controller = HiCacheController.__new__(HiCacheController) + controller.write_queue = [op] + controller.io_backend = "kernel" + controller.mem_pool_host = FakeHostPool() + controller.mem_pool_device = None + controller.has_draft = False + controller.write_stream = object() + controller.ack_write_queue = [] + controller.move_indices = mock.Mock( + return_value=(op.host_indices, op.device_indices) + ) + + with mock.patch.object( + manager_cache_controller, "device_module", _FakeDeviceModule + ): + controller.start_writing() + + controller.move_indices.assert_called_once() + self.assertEqual(captured["host_indices"].device.type, "cpu") + + +if __name__ == "__main__": + unittest.main()