Enable CP HiCache direct transfers to use layer-page host layout
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
@@ -105,7 +105,7 @@ for _schema in (
|
||||
if "already" not in str(exc).lower() and "duplicate" not in str(exc).lower():
|
||||
raise
|
||||
|
||||
from sglang.srt.managers.cache_controller import HiCacheController
|
||||
from sglang.srt.managers.cache_controller import CacheOperation, 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 (
|
||||
@@ -449,6 +449,139 @@ class TestPageFirstPerLayerBackupTaiKernel(CustomTestCase):
|
||||
self.assertEqual(kwargs["layer_id"], 1)
|
||||
self.assertEqual(kwargs["page_size"], 1)
|
||||
|
||||
def test_mla_layer_page_first_per_layer_backup_uses_direct_lf_lpf(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 = "layer_page_first"
|
||||
host_pool.page_size = 4
|
||||
host_pool.kv_buffer = torch.empty((3, 8, 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_lpf",
|
||||
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_layer_page_first_per_layer_backup_uses_direct_lf_lpf(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 = "layer_page_first"
|
||||
host_pool.page_size = 4
|
||||
host_pool.kv_buffer = torch.empty((2, 3, 8, 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_lpf",
|
||||
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_layer_page_first_per_layer_backup_uses_direct_lf_lpf(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 = "layer_page_first"
|
||||
host_pool.page_size = 4
|
||||
host_pool.index_k_with_scale_buffer = torch.empty(
|
||||
(3, 8, 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_lpf",
|
||||
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 = []
|
||||
|
||||
@@ -582,6 +715,139 @@ class TestPageFirstPerLayerBackupTaiKernel(CustomTestCase):
|
||||
self.assertEqual(kwargs["layer_id"], 1)
|
||||
self.assertEqual(kwargs["page_size"], 1)
|
||||
|
||||
def test_mla_layer_page_first_per_layer_load_uses_tai_direct_lpf_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 = "layer_page_first"
|
||||
host_pool.page_size = 4
|
||||
host_pool.kv_buffer = torch.empty((3, 8, 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_lpf_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_layer_page_first_per_layer_load_uses_tai_direct_lpf_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 = "layer_page_first"
|
||||
host_pool.page_size = 4
|
||||
host_pool.kv_buffer = torch.empty((2, 3, 8, 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_lpf_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_layer_page_first_per_layer_load_uses_tai_direct_lpf_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 = "layer_page_first"
|
||||
host_pool.page_size = 4
|
||||
host_pool.index_k_with_scale_buffer = torch.empty(
|
||||
(3, 8, 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_lpf_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)
|
||||
|
||||
def test_nsa_indexer_load_reuses_precomputed_page_indices_across_layers(self):
|
||||
calls = []
|
||||
|
||||
@@ -1845,6 +2111,23 @@ class TestHiCacheControllerCPLoad(TestHiCacheControllerCPWrite):
|
||||
self.assertEqual(queued_op.host_indices.tolist(), [100, 101, 102, 103])
|
||||
self.assertEqual(queued_op.device_indices.tolist(), [20, 21, 22, 23])
|
||||
|
||||
def test_direct_layer_page_first_move_indices_keeps_host_order_and_cpu_device_indices(self):
|
||||
host_pool = FakeHostPool(torch.empty((0,), dtype=torch.int64))
|
||||
host_pool.layout = "layer_page_first"
|
||||
controller = self.make_controller(host_pool, cp_rank=1)
|
||||
op = CacheOperation(
|
||||
host_indices=torch.tensor([12, 8, 9, 10], dtype=torch.int64),
|
||||
device_indices=torch.tensor([32, 28, 29, 30], dtype=torch.int64),
|
||||
node_id=7,
|
||||
)
|
||||
|
||||
host_indices, device_indices = controller.move_indices(op, host_pool)
|
||||
|
||||
self.assertEqual(host_indices.tolist(), [12, 8, 9, 10])
|
||||
self.assertEqual(host_indices.device.type, "cpu")
|
||||
self.assertEqual(device_indices.tolist(), [32, 28, 29, 30])
|
||||
self.assertEqual(device_indices.device.type, "cpu")
|
||||
|
||||
def test_cp_load_rejects_non_contiguous_physical_device_page(self):
|
||||
host_pool = FakeHostPool(torch.tensor([100, 101, 102, 103], dtype=torch.int64))
|
||||
allocator = FakeAllocator(
|
||||
|
||||
Reference in New Issue
Block a user