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

[HiCache]Support hybrid pool staged H2D kernel (#28434)

Browse Source 
Co-authored-by: hzh0425 <hzh0425@apache.org>
This commit is contained in:
huangtingwei
2026-06-19 09:48:03 +08:00
committed by GitHub
co-authored by hzh0425
parent 05ee93c44f
commit 6b7ecca663
12 changed files with 1319 additions and 82 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/jit_kernel/hicache.py
+47 -4
View File
@@ -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,
python/sglang/srt/managers/cache_controller.py
+6 -4
View File
@@ -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(
python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py
+6 -4
View File
@@ -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
python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py
+3 -1
View File
@@ -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,
python/sglang/srt/mem_cache/memory_pool_host.py
+242 -49
View File
@@ -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":
test/registered/jit/test_hicache.py
+73 -4
View File
@@ -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"]))
test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_hicache_pp_kl.py
+2 -2
View File
@@ -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
test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_cp.py
+2 -2
View File
@@ -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
test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_dsv4.py
+4 -4
View File
@@ -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",
test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_mamba.py
+4 -4
View File
@@ -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",
test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_nightly.py
+4 -4
View File
@@ -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",
],
test/registered/unit/mem_cache/test_hicache_staged_write_back_dispatch.py
+926
View File
@@ -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()