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:

267
test/registered/unit/managers/test_hicache_controller_cp.py
View File

@@ -109,6 +109,7 @@ from sglang.srt.managers.cache_controller import HiCacheController
from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout
from sglang.srt.mem_cache.hiradix_cache import CpHiCacheNodeMetadata
from sglang.srt.mem_cache.memory_pool_host import (
MHATokenToKVPoolHost,
MLATokenToKVPoolHost,
NSATokenToKVPoolHost,
)
@@ -270,6 +271,272 @@ class TestPageFirstPerLayerBackupTaiKernel(CustomTestCase):
self.assertEqual(kwargs["item_size"], 32)
self.assertEqual(kwargs["dst_layout_dim"], 96)
def test_mla_page_first_direct_per_layer_backup_uses_direct_lf_pf(self):
calls = []
def fake_direct(src_ptrs, dst_ptrs, src_indices, dst_indices, **kwargs):
calls.append(
(src_ptrs, dst_ptrs, src_indices.clone(), dst_indices.clone(), kwargs)
)
host_pool = MLATokenToKVPoolHost.__new__(MLATokenToKVPoolHost)
host_pool.layout = "page_first_direct"
host_pool.page_size = 4
host_pool.kv_buffer = torch.empty((8, 3, 4, 1, 16), dtype=torch.uint8)
device_pool = type("DevicePool", (), {})()
device_pool.kv_buffer = torch.empty((3, 32, 1, 16), dtype=torch.uint8)
host_indices = torch.tensor([4, 5, 6, 7], dtype=torch.int64)
device_indices = torch.tensor([12, 13, 14, 15], dtype=torch.int64)
with patch(
"sglang.srt.mem_cache.memory_pool_host._load_tai_transfer_kv_per_layer_direct_lf_pf",
return_value=fake_direct,
):
host_pool.backup_from_device_per_layer(
device_pool,
host_indices,
device_indices,
layer_id=2,
io_backend="direct",
)
self.assertEqual(len(calls), 1)
src_ptrs, dst_ptrs, src_indices, dst_indices, kwargs = calls[0]
self.assertEqual(len(src_ptrs), 1)
self.assertEqual(src_ptrs[0].data_ptr(), device_pool.kv_buffer[2].data_ptr())
self.assertEqual(len(dst_ptrs), 1)
self.assertEqual(dst_ptrs[0].data_ptr(), host_pool.kv_buffer.data_ptr())
self.assertEqual(src_indices.tolist(), [12, 13, 14, 15])
self.assertEqual(dst_indices.tolist(), [4, 5, 6, 7])
self.assertEqual(kwargs["layer_id"], 2)
self.assertEqual(kwargs["page_size"], 4)
def test_mha_page_first_direct_per_layer_backup_uses_direct_lf_pf(self):
calls = []
def fake_direct(src_ptrs, dst_ptrs, src_indices, dst_indices, **kwargs):
calls.append(
(src_ptrs, dst_ptrs, src_indices.clone(), dst_indices.clone(), kwargs)
)
host_pool = MHATokenToKVPoolHost.__new__(MHATokenToKVPoolHost)
host_pool.layout = "page_first_direct"
host_pool.page_size = 4
host_pool.kv_buffer = torch.empty((2, 8, 3, 4, 2, 8), dtype=torch.uint8)
device_pool = type("DevicePool", (), {})()
device_pool.k_buffer = torch.empty((3, 32, 2, 8), dtype=torch.uint8)
device_pool.v_buffer = torch.empty((3, 32, 2, 8), dtype=torch.uint8)
host_indices = torch.tensor([4, 5, 6, 7], dtype=torch.int64)
device_indices = torch.tensor([12, 13, 14, 15], dtype=torch.int64)
with patch(
"sglang.srt.mem_cache.memory_pool_host._load_tai_transfer_kv_per_layer_direct_lf_pf",
return_value=fake_direct,
):
host_pool.backup_from_device_per_layer(
device_pool,
host_indices,
device_indices,
layer_id=2,
io_backend="direct",
)
self.assertEqual(len(calls), 1)
src_ptrs, dst_ptrs, src_indices, dst_indices, kwargs = calls[0]
self.assertEqual(len(src_ptrs), 2)
self.assertEqual(src_ptrs[0].data_ptr(), device_pool.k_buffer[2].data_ptr())
self.assertEqual(src_ptrs[1].data_ptr(), device_pool.v_buffer[2].data_ptr())
self.assertEqual(len(dst_ptrs), 2)
self.assertEqual(dst_ptrs[0].data_ptr(), host_pool.k_buffer.data_ptr())
self.assertEqual(dst_ptrs[1].data_ptr(), host_pool.v_buffer.data_ptr())
self.assertEqual(src_indices.tolist(), [12, 13, 14, 15])
self.assertEqual(dst_indices.tolist(), [4, 5, 6, 7])
self.assertEqual(kwargs["layer_id"], 2)
self.assertEqual(kwargs["page_size"], 4)
def test_nsa_indexer_page_first_direct_per_layer_backup_uses_direct_lf_pf(self):
calls = []
def fake_direct(src_ptrs, dst_ptrs, src_indices, dst_indices, **kwargs):
calls.append(
(src_ptrs, dst_ptrs, src_indices.clone(), dst_indices.clone(), kwargs)
)
host_pool = NSATokenToKVPoolHost.__new__(NSATokenToKVPoolHost)
host_pool.layout = "page_first_direct"
host_pool.page_size = 4
host_pool.index_k_with_scale_buffer = torch.empty(
(8, 3, 1, 32), dtype=torch.uint8
)
device_pool = type("DevicePool", (), {})()
device_pool.index_k_with_scale_buffer = torch.empty(
(3, 8, 32), dtype=torch.uint8
)
host_indices = torch.tensor([4, 5, 6, 7], dtype=torch.int64)
device_indices = torch.tensor([12, 13, 14, 15], dtype=torch.int64)
with patch(
"sglang.srt.mem_cache.memory_pool_host._load_tai_transfer_kv_per_layer_direct_lf_pf",
return_value=fake_direct,
):
host_pool._backup_indexer_from_device_per_layer(
device_pool,
host_indices,
device_indices,
layer_id=1,
io_backend="direct",
)
self.assertEqual(len(calls), 1)
src_ptrs, dst_ptrs, src_indices, dst_indices, kwargs = calls[0]
self.assertEqual(len(src_ptrs), 1)
self.assertEqual(
src_ptrs[0].data_ptr(),
device_pool.index_k_with_scale_buffer[1].data_ptr(),
)
self.assertEqual(len(dst_ptrs), 1)
self.assertEqual(
dst_ptrs[0].data_ptr(),
host_pool.index_k_with_scale_buffer.data_ptr(),
)
self.assertEqual(src_indices.tolist(), [3])
self.assertEqual(dst_indices.tolist(), [1])
self.assertEqual(kwargs["layer_id"], 1)
self.assertEqual(kwargs["page_size"], 1)
def test_mla_page_first_direct_per_layer_load_uses_tai_direct_pf_lf(self):
calls = []
def fake_direct(src_ptrs, dst_ptrs, src_indices, dst_indices, **kwargs):
calls.append(
(src_ptrs, dst_ptrs, src_indices.clone(), dst_indices.clone(), kwargs)
)
host_pool = MLATokenToKVPoolHost.__new__(MLATokenToKVPoolHost)
host_pool.layout = "page_first_direct"
host_pool.page_size = 4
host_pool.kv_buffer = torch.empty((8, 3, 4, 1, 16), dtype=torch.uint8)
device_pool = type("DevicePool", (), {})()
device_pool.kv_buffer = torch.empty((3, 32, 1, 16), dtype=torch.uint8)
host_indices = torch.tensor([4, 5, 6, 7], dtype=torch.int64)
device_indices = torch.tensor([12, 13, 14, 15], dtype=torch.int64)
with patch(
"sglang.srt.mem_cache.memory_pool_host._load_tai_transfer_kv_per_layer_direct_pf_lf",
return_value=fake_direct,
):
host_pool.load_to_device_per_layer(
device_pool,
host_indices,
device_indices,
layer_id=2,
io_backend="direct",
)
self.assertEqual(len(calls), 1)
src_ptrs, dst_ptrs, src_indices, dst_indices, kwargs = calls[0]
self.assertEqual(len(src_ptrs), 1)
self.assertEqual(src_ptrs[0].data_ptr(), host_pool.kv_buffer.data_ptr())
self.assertEqual(len(dst_ptrs), 1)
self.assertEqual(dst_ptrs[0].data_ptr(), device_pool.kv_buffer[2].data_ptr())
self.assertEqual(src_indices.tolist(), [4, 5, 6, 7])
self.assertEqual(dst_indices.tolist(), [12, 13, 14, 15])
self.assertEqual(kwargs["layer_id"], 2)
self.assertEqual(kwargs["page_size"], 4)
def test_mha_page_first_direct_per_layer_load_uses_tai_direct_pf_lf(self):
calls = []
def fake_direct(src_ptrs, dst_ptrs, src_indices, dst_indices, **kwargs):
calls.append(
(src_ptrs, dst_ptrs, src_indices.clone(), dst_indices.clone(), kwargs)
)
host_pool = MHATokenToKVPoolHost.__new__(MHATokenToKVPoolHost)
host_pool.layout = "page_first_direct"
host_pool.page_size = 4
host_pool.kv_buffer = torch.empty((2, 8, 3, 4, 2, 8), dtype=torch.uint8)
device_pool = type("DevicePool", (), {})()
device_pool.k_buffer = torch.empty((3, 32, 2, 8), dtype=torch.uint8)
device_pool.v_buffer = torch.empty((3, 32, 2, 8), dtype=torch.uint8)
host_indices = torch.tensor([4, 5, 6, 7], dtype=torch.int64)
device_indices = torch.tensor([12, 13, 14, 15], dtype=torch.int64)
with patch(
"sglang.srt.mem_cache.memory_pool_host._load_tai_transfer_kv_per_layer_direct_pf_lf",
return_value=fake_direct,
):
host_pool.load_to_device_per_layer(
device_pool,
host_indices,
device_indices,
layer_id=2,
io_backend="direct",
)
self.assertEqual(len(calls), 1)
src_ptrs, dst_ptrs, src_indices, dst_indices, kwargs = calls[0]
self.assertEqual(len(src_ptrs), 2)
self.assertEqual(src_ptrs[0].data_ptr(), host_pool.k_buffer.data_ptr())
self.assertEqual(src_ptrs[1].data_ptr(), host_pool.v_buffer.data_ptr())
self.assertEqual(len(dst_ptrs), 2)
self.assertEqual(dst_ptrs[0].data_ptr(), device_pool.k_buffer[2].data_ptr())
self.assertEqual(dst_ptrs[1].data_ptr(), device_pool.v_buffer[2].data_ptr())
self.assertEqual(src_indices.tolist(), [4, 5, 6, 7])
self.assertEqual(dst_indices.tolist(), [12, 13, 14, 15])
self.assertEqual(kwargs["layer_id"], 2)
self.assertEqual(kwargs["page_size"], 4)
def test_nsa_indexer_page_first_direct_per_layer_load_uses_tai_direct_pf_lf(self):
calls = []
def fake_direct(src_ptrs, dst_ptrs, src_indices, dst_indices, **kwargs):
calls.append(
(src_ptrs, dst_ptrs, src_indices.clone(), dst_indices.clone(), kwargs)
)
host_pool = NSATokenToKVPoolHost.__new__(NSATokenToKVPoolHost)
host_pool.layout = "page_first_direct"
host_pool.page_size = 4
host_pool.index_k_with_scale_buffer = torch.empty(
(8, 3, 1, 32), dtype=torch.uint8
)
device_pool = type("DevicePool", (), {})()
device_pool.index_k_with_scale_buffer = torch.empty(
(3, 8, 32), dtype=torch.uint8
)
host_indices = torch.tensor([4, 5, 6, 7], dtype=torch.int64)
device_indices = torch.tensor([12, 13, 14, 15], dtype=torch.int64)
with patch(
"sglang.srt.mem_cache.memory_pool_host._load_tai_transfer_kv_per_layer_direct_pf_lf",
return_value=fake_direct,
):
host_pool._load_indexer_to_device_per_layer(
device_pool,
host_indices,
device_indices,
layer_id=1,
io_backend="direct",
)
self.assertEqual(len(calls), 1)
src_ptrs, dst_ptrs, src_indices, dst_indices, kwargs = calls[0]
self.assertEqual(len(src_ptrs), 1)
self.assertEqual(
src_ptrs[0].data_ptr(),
host_pool.index_k_with_scale_buffer.data_ptr(),
)
self.assertEqual(len(dst_ptrs), 1)
self.assertEqual(
dst_ptrs[0].data_ptr(),
device_pool.index_k_with_scale_buffer[1].data_ptr(),
)
self.assertEqual(src_indices.tolist(), [1])
self.assertEqual(dst_indices.tolist(), [3])
self.assertEqual(kwargs["layer_id"], 1)
self.assertEqual(kwargs["page_size"], 1)
class FakeAllocator:
def __init__(self, alloc_result=None):