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

Route page-first direct HiCache loads through TAI

Browse Source 
Use tai-kernel for direct page_first_direct per-layer H2D load across MHA, MLA, and NSA indexer pools. This keeps SGLang off the sgl-kernel cudaMemcpyBatchAsync path that crashes on CUDA 13 while preserving fail-fast behavior when the required TAI op is unavailable.

Constraint: remote CUDA 13 stack crashes in sgl-kernel PF->LF direct load via cuMemcpyBatchAsync_v2

Rejected: Silent fallback to sgl-kernel or Python loop | fallbacks would hide either a crash-prone ABI path or a large performance regression

Confidence: high

Scope-risk: moderate

Directive: page_first_direct direct load must remain fail-fast if tai_kernel.nsa_prefill.transfer_kv_per_layer_direct_pf_lf is missing

Tested: remote g0034 PYTHONPATH=python pytest -q test/registered/unit/managers/test_hicache_controller_cp.py: 55 passed, 3 warnings

Tested: remote g0034 CUDA smoke for MLATokenToKVPoolHost.load_to_device_per_layer with direct/page_first_direct passed

Not-tested: full SGLang ETE server after the final commit
This commit is contained in:
laoyao0822
2026-05-28 02:37:56 +08:00
parent 40a8de5fd1
commit 67d52346de
2 changed files with 329 additions and 30 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -40,7 +40,6 @@ if not (_is_npu or _is_xpu or _is_mps):
transfer_kv_all_layer_mla_lf_pf,
transfer_kv_direct,
transfer_kv_per_layer,
transfer_kv_per_layer_direct_pf_lf,
transfer_kv_per_layer_mla,
transfer_kv_per_layer_mla_pf_lf,
transfer_kv_per_layer_pf_lf,
@@ -66,6 +65,38 @@ def _load_tai_transfer_kv_per_layer_mla_lf_pf():
) from exc
@lru_cache(maxsize=1)
def _load_tai_transfer_kv_per_layer_direct_lf_pf():
try:
from tai_kernel.nsa_prefill import transfer_kv_per_layer_direct_lf_pf
return transfer_kv_per_layer_direct_lf_pf
except Exception as exc:
raise RuntimeError(
"[CP_HICACHE_FAILFAST][missing_tai_page_first_direct_lf_pf] "
"direct+page_first_direct per-layer D2H backup requires "
"tai_kernel.nsa_prefill.transfer_kv_per_layer_direct_lf_pf. "
"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:
from tai_kernel.nsa_prefill import transfer_kv_per_layer_direct_pf_lf
return transfer_kv_per_layer_direct_pf_lf
except Exception as exc:
raise RuntimeError(
"[CP_HICACHE_FAILFAST][missing_tai_page_first_direct_pf_lf] "
"direct+page_first_direct per-layer H2D load requires "
"tai_kernel.nsa_prefill.transfer_kv_per_layer_direct_pf_lf. "
"Build/sync tai-kernel; this path intentionally does not fall back "
"to the sgl-kernel direct copy path because it is unsafe on CUDA 13."
) from exc
def synchronized(func):
@wraps(func)
def wrapper(self, *args, **kwargs):
@@ -476,7 +507,7 @@ class MHATokenToKVPoolHost(HostKVCache):
page_size=self.page_size,
)
elif self.layout == "page_first_direct":
transfer_kv_per_layer_direct_pf_lf(
_load_tai_transfer_kv_per_layer_direct_pf_lf()(
src_ptrs=[self.k_buffer, self.v_buffer],
dst_ptrs=[
device_pool.k_buffer[layer_id],
@@ -640,17 +671,17 @@ class MHATokenToKVPoolHost(HostKVCache):
page_size=self.page_size,
)
elif self.layout == "page_first_direct":
for host_index, device_index in zip(
host_indices.cpu().tolist(), device_indices.cpu().tolist()
):
host_page = host_index // self.page_size
host_offset = host_index % self.page_size
self.k_buffer[host_page, layer_id, host_offset] = (
device_pool.k_buffer[layer_id][device_index]
)
self.v_buffer[host_page, layer_id, host_offset] = (
device_pool.v_buffer[layer_id][device_index]
)
_load_tai_transfer_kv_per_layer_direct_lf_pf()(
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":
@@ -993,7 +1024,7 @@ class MLATokenToKVPoolHost(HostKVCache):
page_size=self.page_size,
)
elif self.layout == "page_first_direct":
transfer_kv_per_layer_direct_pf_lf(
_load_tai_transfer_kv_per_layer_direct_pf_lf()(
src_ptrs=[self.kv_buffer],
dst_ptrs=[device_pool.kv_buffer[layer_id]],
src_indices=host_indices,
@@ -1126,14 +1157,14 @@ class MLATokenToKVPoolHost(HostKVCache):
page_size=self.page_size,
)
elif self.layout == "page_first_direct":
for host_index, device_index in zip(
host_indices.cpu().tolist(), device_indices.cpu().tolist()
):
host_page = host_index // self.page_size
host_offset = host_index % self.page_size
self.kv_buffer[host_page, layer_id, host_offset] = (
device_pool.kv_buffer[layer_id][device_index]
)
_load_tai_transfer_kv_per_layer_direct_lf_pf()(
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":
@@ -1398,7 +1429,7 @@ class NSATokenToKVPoolHost(MLATokenToKVPoolHost):
page_size=1,
)
elif self.layout == "page_first_direct":
transfer_kv_per_layer_direct_pf_lf(
_load_tai_transfer_kv_per_layer_direct_pf_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,
@@ -1500,13 +1531,14 @@ class NSATokenToKVPoolHost(MLATokenToKVPoolHost):
page_size=1,
)
elif self.layout == "page_first_direct":
for host_page, device_page in zip(
host_page_indices.cpu().tolist(),
device_page_indices.cpu().tolist(),
):
self.index_k_with_scale_buffer[host_page, layer_id, 0] = (
device_pool.index_k_with_scale_buffer[layer_id][device_page]
)
_load_tai_transfer_kv_per_layer_direct_lf_pf()(
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: