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

feat: add page-aligned index validation for CP HiCache direct transfers

Browse Source 
Add validate_page_aligned_token_indices utility and apply it across
HiCache write/load paths, NSA indexer transfers, and CUDA direct copy
kernels to reject malformed (partial, misaligned, non-contiguous) page
groups before they reach native transfer code. Also validate supported
CP HiCache backend/layout combinations at server startup.
This commit is contained in:
ThomasX
2026-05-10 02:02:54 +08:00
parent f38937723e
commit 3a5928ef53
10 changed files with 472 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/unit/managers/test_hicache_controller_cp.py
+63
View File
@@ -195,6 +195,32 @@ class TestHiCacheControllerCPWrite(CustomTestCase):
self.assertEqual(host_pool.alloc_calls, [4])
self.assertEqual(host_pool.backups[0][1].tolist(), [4, 5, 6, 7])
def test_cp_write_rejects_incomplete_owned_physical_page(self):
host_pool = FakeHostPool(torch.tensor([100, 101, 102, 103], dtype=torch.int64))
controller = self.make_controller(host_pool, cp_rank=1)
logical_locs = torch.tensor([8, 9, 10], dtype=torch.int64)
with self.assertRaisesRegex(ValueError, "physical_device_indices.*whole pages"):
controller.write(logical_locs, node_id=21)
def test_cp_write_rejects_non_contiguous_owned_physical_page(self):
host_pool = FakeHostPool(torch.tensor([100, 101, 102, 103], dtype=torch.int64))
controller = self.make_controller(host_pool, cp_rank=1)
logical_locs = torch.tensor([8, 9, 11, 10], dtype=torch.int64)
with self.assertRaisesRegex(
ValueError, "physical_device_indices.*contiguous page spans"
):
controller.write(logical_locs, node_id=22)
def test_cp_write_rejects_non_contiguous_host_page(self):
host_pool = FakeHostPool(torch.tensor([100, 101, 103, 102], dtype=torch.int64))
controller = self.make_controller(host_pool, cp_rank=1)
logical_locs = torch.arange(8, 12, dtype=torch.int64)
with self.assertRaisesRegex(ValueError, "host_indices.*contiguous page spans"):
controller.write(logical_locs, node_id=23)
def test_cp_write_zero_owned_returns_metadata_and_noop_ack(self):
host_pool = FakeHostPool(torch.tensor([], dtype=torch.int64))
controller = self.make_controller(host_pool, cp_rank=3)
@@ -315,6 +341,43 @@ 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_cp_load_rejects_non_contiguous_physical_device_page(self):
host_pool = FakeHostPool(torch.tensor([100, 101, 102, 103], dtype=torch.int64))
allocator = FakeAllocator(
alloc_result=torch.tensor(
[64, 65, 66, 67, 68, 69, 71, 70, 72, 73, 74, 75, 76, 77, 78, 79],
dtype=torch.int64,
)
)
controller = self.make_controller(host_pool, allocator=allocator, cp_rank=1)
node = TreeNode()
node.host_len = 16
node.cp_hicache = CpHiCacheNodeMetadata(
logical_len=16,
owned_positions=torch.tensor([4, 5, 6, 7], dtype=torch.int64),
host_indices=torch.tensor([100, 101, 102, 103], dtype=torch.int64),
)
with self.assertRaisesRegex(
ValueError, "physical_device_indices.*contiguous page spans"
):
controller.load_cp([node], node_id=31)
def test_cp_load_rejects_non_contiguous_host_page(self):
host_pool = FakeHostPool(torch.tensor([100, 101, 102, 103], dtype=torch.int64))
allocator = FakeAllocator(alloc_result=torch.arange(64, 80, dtype=torch.int64))
controller = self.make_controller(host_pool, allocator=allocator, cp_rank=1)
node = TreeNode()
node.host_len = 16
node.cp_hicache = CpHiCacheNodeMetadata(
logical_len=16,
owned_positions=torch.tensor([4, 5, 6, 7], dtype=torch.int64),
host_indices=torch.tensor([100, 101, 103, 102], dtype=torch.int64),
)
with self.assertRaisesRegex(ValueError, "host_indices.*contiguous page spans"):
controller.load_cp([node], node_id=32)
if __name__ == "__main__":
main()