From 67d52346de8cdd7ca97cba943e91f0cd82882f3d Mon Sep 17 00:00:00 2001 From: laoyao0822 Date: Thu, 28 May 2026 02:37:56 +0800 Subject: [PATCH] Route page-first direct HiCache loads through TAI 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 --- .../sglang/srt/mem_cache/memory_pool_host.py | 92 ++++-- .../managers/test_hicache_controller_cp.py | 267 ++++++++++++++++++ 2 files changed, 329 insertions(+), 30 deletions(-) diff --git a/python/sglang/srt/mem_cache/memory_pool_host.py b/python/sglang/srt/mem_cache/memory_pool_host.py index c7ba5b557..fbacc88b4 100644 --- a/python/sglang/srt/mem_cache/memory_pool_host.py +++ b/python/sglang/srt/mem_cache/memory_pool_host.py @@ -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: diff --git a/test/registered/unit/managers/test_hicache_controller_cp.py b/test/registered/unit/managers/test_hicache_controller_cp.py index 9c529d24d..ea3af2f86 100644 --- a/test/registered/unit/managers/test_hicache_controller_cp.py +++ b/test/registered/unit/managers/test_hicache_controller_cp.py @@ -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):