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
This commit is contained in:
laoyao0822
@@ -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):
|
||||
|
||||
Reference in New Issue
Block a user