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

Enable CP HiCache direct transfers to use layer-page host layout

Browse Source 
CP shared-KV HiCache transfers are per-layer, so the host layout should match the access pattern instead of forcing page-major strides through every layer. This adds a direct-only layer_page_first layout, routes per-layer KV and NSA index backup/load through the TAI LF<->LPF direct kernels, and keeps storage/page-buffer metadata paths fail-fast until their page-level contract is redesigned.\n\nThe direct controller keeps host indices in caller order for both page_first_direct and layer_page_first because the TAI direct path requires CPU index descriptors and owns descriptor coalescing. All-layer backup intentionally loops over per-layer direct kernels rather than using the sgl-kernel all-layer direct ABI.\n\nConstraint: layer_page_first is currently host-only CP HiCache; storage backends assume page-major contiguous page metadata.\nConstraint: TAI LPF direct kernels require CPU int64 page indices and complete page spans.\nRejected: silently fallback to SM copy when TAI LPF kernels are missing | that hides production performance regressions.\nRejected: support storage page metadata in this commit | LPF requires a layer-page-level storage contract, not a one-pointer-per-page contract.\nConfidence: medium\nScope-risk: moderate\nDirective: Do not enable storage or kernel backend for layer_page_first without redesigning page-buffer metadata and adding remote ETE coverage.\nTested: local py_compile for touched runtime files.\nTested: remote py_compile in g0034 container for touched runtime files.\nTested: remote targeted pytest: 5 passed for parser/storage/layout/move_indices smoke coverage.\nNot-tested: full CP HiCache ETE with --hicache-mem-layout layer_page_first after this commit step.\nNot-tested: combined CUDA roundtrip tests in one pytest process; previous independent runs passed but combined run exposed a host-memory registration lifecycle issue.
This commit is contained in:
laoyao0822
2026-06-10 02:03:07 +08:00
parent 5e22279670
commit 24da983ff5
8 changed files with 1796 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
python/sglang/srt/managers/cache_controller.py
View File

@@ -1697,7 +1697,7 @@ class HiCacheController:
device_indices = device_indices.cpu()
host_indices, idx = host_indices.sort()
return host_indices, device_indices.index_select(0, idx)
elif mem_pool_host.layout == "page_first_direct":
elif mem_pool_host.layout in ("page_first_direct", "layer_page_first"):
return host_indices, device_indices.cpu()
else:
raise ValueError(

8
python/sglang/srt/mem_cache/hicache_storage.py
View File

@@ -78,6 +78,14 @@ class HiCacheStorage(ABC):
# todo, the page size of storage backend does not have to be the same as the same as host memory pool
def register_mem_pool_host(self, mem_pool_host: HostKVCache):
if getattr(mem_pool_host, "layout", None) == "layer_page_first":
raise RuntimeError(
"[CP_HICACHE_FAILFAST][layer_page_first_storage_unsupported] "
"layer_page_first is currently a host-only direct CP HiCache "
"layout. Storage backends are not supported for this layout "
"because their page-level zero-copy/metadata contracts assume "
"page-major contiguous pages."
)
self.mem_pool_host = mem_pool_host
def batch_get_v1(

202
python/sglang/srt/mem_cache/memory_pool_host.py
View File

@@ -83,6 +83,22 @@ def _load_tai_transfer_kv_per_layer_direct_lf_pf():
) from exc
@lru_cache(maxsize=1)
def _load_tai_transfer_kv_per_layer_direct_lf_lpf():
try:
from tai_kernel.nsa_prefill import transfer_kv_per_layer_direct_lf_lpf
return transfer_kv_per_layer_direct_lf_lpf
except Exception as exc:
raise RuntimeError(
"[CP_HICACHE_FAILFAST][missing_tai_layer_page_first_direct_lf_lpf] "
"direct+layer_page_first per-layer D2H backup requires "
"tai_kernel.nsa_prefill.transfer_kv_per_layer_direct_lf_lpf. "
"Build/sync tai-kernel; this path intentionally does not fall back "
"to an SM-consuming copy kernel."
) from exc
@lru_cache(maxsize=1)
def _load_tai_transfer_kv_per_layer_direct_pf_lf():
try:
@@ -99,6 +115,32 @@ def _load_tai_transfer_kv_per_layer_direct_pf_lf():
) from exc
@lru_cache(maxsize=1)
def _load_tai_transfer_kv_per_layer_direct_lpf_lf():
try:
from tai_kernel.nsa_prefill import transfer_kv_per_layer_direct_lpf_lf
return transfer_kv_per_layer_direct_lpf_lf
except Exception as exc:
raise RuntimeError(
"[CP_HICACHE_FAILFAST][missing_tai_layer_page_first_direct_lpf_lf] "
"direct+layer_page_first per-layer H2D load requires "
"tai_kernel.nsa_prefill.transfer_kv_per_layer_direct_lpf_lf. "
"Build/sync tai-kernel; this path intentionally does not fall back "
"to an SM-consuming copy kernel."
) from exc
def _raise_layer_page_first_page_buffer_meta_unsupported() -> None:
raise RuntimeError(
"[CP_HICACHE_FAILFAST][layer_page_first_page_buffer_meta_unsupported] "
"layer_page_first is currently a host-only direct CP HiCache layout. "
"Storage/zero-copy page buffer metadata is not supported because a "
"single page across all layers is not stored as one page-major "
"contiguous block. Disable storage or use page_first_direct."
)
def synchronized(func):
@wraps(func)
def wrapper(self, *args, **kwargs):
@@ -768,6 +810,15 @@ class MHATokenToKVPoolHost(HostKVCache):
self.head_num,
self.head_dim,
)
elif self.layout == "layer_page_first":
dims = (
2,
self.layer_num,
self.page_num,
self.page_size,
self.head_num,
self.head_dim,
)
elif self.layout == "page_head":
dims = (
2,
@@ -882,6 +933,18 @@ class MHATokenToKVPoolHost(HostKVCache):
layer_id=layer_id,
page_size=self.page_size,
)
elif self.layout == "layer_page_first":
_load_tai_transfer_kv_per_layer_direct_lpf_lf()(
src_ptrs=[self.k_buffer, self.v_buffer],
dst_ptrs=[
device_pool.k_buffer[layer_id],
device_pool.v_buffer[layer_id],
],
src_indices=host_indices,
dst_indices=device_indices,
layer_id=layer_id,
page_size=self.page_size,
)
else:
raise ValueError(f"Unsupported layout: {self.layout}")
elif io_backend == "kernel_ascend":
@@ -986,6 +1049,20 @@ class MHATokenToKVPoolHost(HostKVCache):
layer_id=layer_id,
page_size=self.page_size,
)
elif self.layout == "layer_page_first":
tai_transfer = _load_tai_transfer_kv_per_layer_direct_lf_lpf()
for layer_id in range(self.layer_num):
tai_transfer(
src_ptrs=[
device_pool.k_buffer[layer_id],
device_pool.v_buffer[layer_id],
],
dst_ptrs=[self.k_buffer, self.v_buffer],
src_indices=device_indices,
dst_indices=host_indices,
layer_id=layer_id,
page_size=self.page_size,
)
else:
raise ValueError(f"Unsupported layout: {self.layout}")
elif io_backend == "kernel_ascend":
@@ -1056,6 +1133,18 @@ class MHATokenToKVPoolHost(HostKVCache):
layer_id=layer_id,
page_size=self.page_size,
)
elif self.layout == "layer_page_first":
_load_tai_transfer_kv_per_layer_direct_lf_lpf()(
src_ptrs=[
device_pool.k_buffer[layer_id],
device_pool.v_buffer[layer_id],
],
dst_ptrs=[self.k_buffer, self.v_buffer],
src_indices=device_indices,
dst_indices=host_indices,
layer_id=layer_id,
page_size=self.page_size,
)
else:
raise ValueError(f"Unsupported layout: {self.layout}")
elif io_backend == "kernel_ascend":
@@ -1084,6 +1173,9 @@ class MHATokenToKVPoolHost(HostKVCache):
elif self.layout in ["page_first_direct", "page_head"]:
real_index = index // self.page_size
data_page = self.kv_buffer[:, real_index : real_index + 1, :, :, :, :]
elif self.layout == "layer_page_first":
real_index = index // self.page_size
data_page = self.kv_buffer[:, :, real_index : real_index + 1, :, :, :]
else:
raise ValueError(f"Unsupported layout: {self.layout}")
if flat:
@@ -1122,6 +1214,13 @@ class MHATokenToKVPoolHost(HostKVCache):
2, 1, self.layer_num, self.page_size, self.head_num, self.head_dim
)
)
elif self.layout == "layer_page_first":
real_index = index // self.page_size
self.kv_buffer[:, :, real_index : real_index + 1, :, :, :] = (
data_page.reshape(
2, self.layer_num, 1, self.page_size, self.head_num, self.head_dim
)
)
elif self.layout == "page_head":
real_index = index // self.page_size
self.kv_buffer[:, real_index : real_index + 1, :, :, :, :] = (
@@ -1238,6 +1337,8 @@ class MHATokenToKVPoolHost(HostKVCache):
* self.head_dim
)
element_size_list = [element_size] * len(ptr_list)
elif self.layout == "layer_page_first":
_raise_layer_page_first_page_buffer_meta_unsupported()
else:
raise ValueError(f"Unsupported layout: {self.layout}")
return ptr_list, element_size_list
@@ -1313,6 +1414,14 @@ class MLATokenToKVPoolHost(HostKVCache):
1,
self.kv_cache_dim,
)
elif self.layout == "layer_page_first":
dims = (
self.layer_num,
self.page_num,
self.page_size,
1,
self.kv_cache_dim,
)
# Ascend-specific: Aligns with NPUMLATokenToKVPool layout
# Separately allocate k_buffer and v_buffer for easier data transfer.
elif self.layout == "page_first_kv_split":
@@ -1406,6 +1515,15 @@ class MLATokenToKVPoolHost(HostKVCache):
layer_id=layer_id,
page_size=self.page_size,
)
elif self.layout == "layer_page_first":
_load_tai_transfer_kv_per_layer_direct_lpf_lf()(
src_ptrs=[self.kv_buffer],
dst_ptrs=[device_pool.kv_buffer[layer_id]],
src_indices=host_indices,
dst_indices=device_indices,
layer_id=layer_id,
page_size=self.page_size,
)
else:
raise ValueError(f"Unsupported layout: {self.layout}")
elif io_backend == "kernel_ascend":
@@ -1478,6 +1596,17 @@ class MLATokenToKVPoolHost(HostKVCache):
layer_id=layer_id,
page_size=self.page_size,
)
elif self.layout == "layer_page_first":
tai_transfer = _load_tai_transfer_kv_per_layer_direct_lf_lpf()
for layer_id in range(self.layer_num):
tai_transfer(
src_ptrs=[device_pool.kv_buffer[layer_id]],
dst_ptrs=[self.kv_buffer],
src_indices=device_indices,
dst_indices=host_indices,
layer_id=layer_id,
page_size=self.page_size,
)
else:
raise ValueError(f"Unsupported layout: {self.layout}")
elif io_backend == "kernel_ascend":
@@ -1546,6 +1675,15 @@ class MLATokenToKVPoolHost(HostKVCache):
layer_id=layer_id,
page_size=self.page_size,
)
elif self.layout == "layer_page_first":
_load_tai_transfer_kv_per_layer_direct_lf_lpf()(
src_ptrs=[device_pool.kv_buffer[layer_id]],
dst_ptrs=[self.kv_buffer],
src_indices=device_indices,
dst_indices=host_indices,
layer_id=layer_id,
page_size=self.page_size,
)
else:
raise ValueError(f"Unsupported layout: {self.layout}")
elif io_backend == "kernel_ascend":
@@ -1576,6 +1714,9 @@ class MLATokenToKVPoolHost(HostKVCache):
elif self.layout == "page_first_direct":
real_index = index // self.page_size
data_page = self.kv_buffer[real_index : real_index + 1, :, :, :, :]
elif self.layout == "layer_page_first":
real_index = index // self.page_size
data_page = self.kv_buffer[:, real_index : real_index + 1, :, :, :]
else:
raise ValueError(f"Unsupported layout: {self.layout}")
if flat:
@@ -1619,6 +1760,15 @@ class MLATokenToKVPoolHost(HostKVCache):
1,
self.kv_cache_dim,
)
elif self.layout == "layer_page_first":
real_index = index // self.page_size
self.kv_buffer[:, real_index : real_index + 1, :, :, :] = data_page.reshape(
self.layer_num,
1,
self.page_size,
1,
self.kv_cache_dim,
)
else:
raise ValueError(f"Unsupported layout: {self.layout}")
@@ -1658,6 +1808,8 @@ class MLATokenToKVPoolHost(HostKVCache):
* self.kv_cache_dim
)
element_size_list = [element_size] * len(ptr_list)
elif self.layout == "layer_page_first":
_raise_layer_page_first_page_buffer_meta_unsupported()
else:
raise ValueError(f"Unsupported layout: {self.layout}")
return ptr_list, element_size_list
@@ -1754,6 +1906,27 @@ class NSATokenToKVPoolHost(MLATokenToKVPoolHost):
pin_memory=self.pin_memory,
allocator=self.allocator,
)
elif self.layout == "layer_page_first":
self.index_k_with_scale_buffer = alloc_func(
(
self.layer_num,
self.indexer_page_num,
1,
self.indexer_page_stride_size,
),
dtype=self.indexer_dtype,
device=self.device,
pin_memory=self.pin_memory,
allocator=self.allocator,
)
self.index_k_data_refs = [
self.index_k_with_scale_buffer[i] for i in range(self.layer_num)
]
self.index_k_data_ptrs = torch.tensor(
[x.data_ptr() for x in self.index_k_data_refs],
dtype=torch.uint64,
device=self.device_pool.device,
)
else:
raise ValueError(f"Unsupported layout: {self.layout}")
@@ -1833,6 +2006,15 @@ class NSATokenToKVPoolHost(MLATokenToKVPoolHost):
layer_id=layer_id,
page_size=1,
)
elif self.layout == "layer_page_first":
_load_tai_transfer_kv_per_layer_direct_lpf_lf()(
src_ptrs=[self.index_k_with_scale_buffer],
dst_ptrs=[device_pool.index_k_with_scale_buffer[layer_id]],
src_indices=host_page_indices,
dst_indices=device_page_indices,
layer_id=layer_id,
page_size=1,
)
else:
raise ValueError(f"Unsupported layout: {self.layout}")
else:
@@ -1892,6 +2074,17 @@ class NSATokenToKVPoolHost(MLATokenToKVPoolHost):
layer_id=layer_id,
page_size=1,
)
elif self.layout == "layer_page_first":
tai_transfer = _load_tai_transfer_kv_per_layer_direct_lf_lpf()
for layer_id in range(self.layer_num):
tai_transfer(
src_ptrs=[device_pool.index_k_with_scale_buffer[layer_id]],
dst_ptrs=[self.index_k_with_scale_buffer],
src_indices=device_page_indices,
dst_indices=host_page_indices,
layer_id=layer_id,
page_size=1,
)
else:
raise ValueError(f"Unsupported layout: {self.layout}")
else:
@@ -1943,6 +2136,15 @@ class NSATokenToKVPoolHost(MLATokenToKVPoolHost):
layer_id=layer_id,
page_size=1,
)
elif self.layout == "layer_page_first":
_load_tai_transfer_kv_per_layer_direct_lf_lpf()(
src_ptrs=[device_pool.index_k_with_scale_buffer[layer_id]],
dst_ptrs=[self.index_k_with_scale_buffer],
src_indices=device_page_indices,
dst_indices=host_page_indices,
layer_id=layer_id,
page_size=1,
)
else:
raise ValueError(f"Unsupported layout: {self.layout}")
else:

19
python/sglang/srt/server_args.py
View File

@@ -962,6 +962,7 @@ class ServerArgs:
("kernel", "page_first"),
("direct", "layer_first"),
("direct", "page_first_direct"),
("direct", "layer_page_first"),
}
current_pair = (self.hicache_io_backend, self.hicache_mem_layout)
if current_pair in supported_pairs:
@@ -970,12 +971,14 @@ class ServerArgs:
raise ValueError(
"CP shared KV HiCache supports only "
"kernel/layer_first, kernel/page_first, direct/layer_first, "
"and direct/page_first_direct after HiCache normalization. "
"direct/page_first_direct, and direct/layer_page_first after "
"HiCache normalization. "
f"Got hicache_io_backend={self.hicache_io_backend!r} "
f"hicache_mem_layout={self.hicache_mem_layout!r}. "
"page_head is MHA/Mooncake-specific, page_first_kv_split is "
"kernel_ascend-specific, and kernel_ascend CP shared KV HiCache "
"is not implemented in the CUDA NSA host path."
"is not implemented in the CUDA NSA host path. "
"layer_page_first is supported only with the direct IO backend."
)
def _handle_deprecated_args(self):
@@ -3027,6 +3030,17 @@ class ServerArgs:
)
def _resolve_storage_layout_compatibility(self):
if (
self.hicache_storage_backend is not None
and self.hicache_mem_layout == "layer_page_first"
):
raise ValueError(
"[CP_HICACHE_FAILFAST][layer_page_first_storage_unsupported] "
"layer_page_first is currently a host-only direct CP HiCache "
"layout and does not support storage backends. Disable "
"hicache_storage_backend or use page_first_direct."
)
if (
self.hicache_storage_backend != "mooncake"
or self.hicache_mem_layout != "layer_first"
@@ -5225,6 +5239,7 @@ class ServerArgs:
"layer_first",
"page_first",
"page_first_direct",
"layer_page_first",
"page_first_kv_split",
"page_head",
],