chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 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()
test/registered/unit/mem_cache/test_nsa_pool_host_unit.py
+75 -12
View File
@@ -10,11 +10,12 @@ from sglang.srt.mem_cache.memory_pool_host import (
)
from sglang.srt.utils import is_cuda, is_hip, is_npu, is_xpu
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.test_utils import CustomTestCase
register_cuda_ci(est_time=3, suite="stage-b-test-1-gpu-small")
class TestNSAHiCacheTransfer(unittest.TestCase):
class TestNSAHiCacheTransfer(CustomTestCase):
def setUp(self):
if not torch.cuda.is_available():
self.skipTest("CUDA is required for NSA host transfer tests.")
@@ -36,7 +37,7 @@ class TestNSAHiCacheTransfer(unittest.TestCase):
]
return torch.cat(parts, dim=0)
def _run_device_to_host_indexer_copy(self, io_backend: str):
def _run_device_to_host_indexer_copy(self, io_backend: str, layout: str):
page_size = 1 if is_hip() else 64
layer_num = 2
size = page_size * 4
@@ -63,7 +64,7 @@ class TestNSAHiCacheTransfer(unittest.TestCase):
host_to_device_ratio=2.0,
host_size=0,
page_size=page_size,
layout="layer_first",
layout=layout,
pin_memory=pin_memory,
device="cpu",
)
@@ -105,26 +106,88 @@ class TestNSAHiCacheTransfer(unittest.TestCase):
for host_page, device_page in zip(
host_pages.tolist(), device_pages.tolist()
):
got = host_pool.index_k_with_scale_buffer[layer_id][host_page].cpu()
if layout == "page_first_direct":
got = host_pool.index_k_with_scale_buffer[host_page, layer_id].cpu()
else:
got = host_pool.index_k_with_scale_buffer[layer_id][host_page].cpu()
expected = device_pool.index_k_with_scale_buffer[layer_id][
device_page
].cpu()
self.assertTrue(torch.equal(got, expected))
host_start = host_page * page_size
device_start = device_page * page_size
got_kv = host_pool.kv_buffer[layer_id][
host_start : host_start + page_size
].cpu()
expected_kv = device_pool.kv_buffer[layer_id][
device_start : device_start + page_size
].cpu()
if layout == "page_first_direct":
got_kv = host_pool.kv_buffer[host_page, layer_id].cpu()
else:
host_start = host_page * page_size
got_kv = host_pool.kv_buffer[layer_id][
host_start : host_start + page_size
].cpu()
self.assertTrue(torch.equal(got_kv, expected_kv))
def test_device_to_host_indexer_kernel(self):
self._run_device_to_host_indexer_copy(io_backend="kernel")
def test_device_to_host_indexer_kernel_layer_first(self):
self._run_device_to_host_indexer_copy(
io_backend="kernel", layout="layer_first"
)
def test_device_to_host_indexer_direct(self):
self._run_device_to_host_indexer_copy(io_backend="direct")
def test_device_to_host_indexer_direct_layer_first(self):
self._run_device_to_host_indexer_copy(
io_backend="direct", layout="layer_first"
)
def test_device_to_host_indexer_direct_page_first_direct(self):
self._run_device_to_host_indexer_copy(
io_backend="direct", layout="page_first_direct"
)
class TestNSAIndexerPageIndices(CustomTestCase):
def make_host_pool_stub(self, page_size: int):
host_pool = NSATokenToKVPoolHost.__new__(NSATokenToKVPoolHost)
host_pool.page_size = page_size
return host_pool
def test_indexer_page_indices_accepts_valid_page_spans(self):
host_pool = self.make_host_pool_stub(page_size=4)
host_pages, device_pages = host_pool._get_indexer_page_indices(
torch.tensor([8, 9, 10, 11], dtype=torch.int64),
torch.tensor([16, 17, 18, 19], dtype=torch.int64),
)
self.assertEqual(host_pages.tolist(), [2])
self.assertEqual(device_pages.tolist(), [4])
def test_indexer_page_indices_rejects_partial_host_page(self):
host_pool = self.make_host_pool_stub(page_size=4)
with self.assertRaisesRegex(ValueError, "host_indices.*whole pages"):
host_pool._get_indexer_page_indices(
torch.tensor([8, 9, 10], dtype=torch.int64),
torch.tensor([16, 17, 18], dtype=torch.int64),
)
def test_indexer_page_indices_rejects_misaligned_device_page(self):
host_pool = self.make_host_pool_stub(page_size=4)
with self.assertRaisesRegex(
ValueError, "device_indices.*start at page boundaries"
):
host_pool._get_indexer_page_indices(
torch.tensor([8, 9, 10, 11], dtype=torch.int64),
torch.tensor([17, 18, 19, 20], dtype=torch.int64),
)
def test_indexer_page_indices_rejects_non_contiguous_host_page(self):
host_pool = self.make_host_pool_stub(page_size=4)
with self.assertRaisesRegex(ValueError, "host_indices.*contiguous page spans"):
host_pool._get_indexer_page_indices(
torch.tensor([8, 9, 11, 10], dtype=torch.int64),
torch.tensor([16, 17, 18, 19], dtype=torch.int64),
)
if __name__ == "__main__":
test/registered/unit/mem_cache/test_page_index_utils.py
+64
View File
@@ -0,0 +1,64 @@
import unittest
import torch
from sglang.srt.mem_cache.page_index_utils import (
validate_page_aligned_token_indices,
)
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase
register_cpu_ci(est_time=1, suite="stage-a-test-cpu")
class TestPageIndexUtils(CustomTestCase):
def test_valid_empty_indices(self):
validate_page_aligned_token_indices(
torch.empty((0,), dtype=torch.int64), 4, "indices"
)
def test_valid_single_page(self):
validate_page_aligned_token_indices(
torch.tensor([8, 9, 10, 11], dtype=torch.int64), 4, "indices"
)
def test_valid_multiple_pages(self):
validate_page_aligned_token_indices(
torch.tensor([4, 5, 6, 7, 16, 17, 18, 19], dtype=torch.int64),
4,
"indices",
)
def test_rejects_non_1d_tensor(self):
with self.assertRaisesRegex(ValueError, "must be a 1-D tensor"):
validate_page_aligned_token_indices(
torch.tensor([[0, 1], [2, 3]], dtype=torch.int64), 4, "indices"
)
def test_rejects_non_positive_page_size(self):
with self.assertRaisesRegex(ValueError, "page_size must be positive"):
validate_page_aligned_token_indices(
torch.tensor([0, 1, 2, 3], dtype=torch.int64), 0, "indices"
)
def test_rejects_partial_page(self):
with self.assertRaisesRegex(ValueError, "whole pages"):
validate_page_aligned_token_indices(
torch.tensor([8, 9, 10], dtype=torch.int64), 4, "indices"
)
def test_rejects_non_page_start(self):
with self.assertRaisesRegex(ValueError, "start at page boundaries"):
validate_page_aligned_token_indices(
torch.tensor([9, 10, 11, 12], dtype=torch.int64), 4, "indices"
)
def test_rejects_non_contiguous_group(self):
with self.assertRaisesRegex(ValueError, "contiguous page spans"):
validate_page_aligned_token_indices(
torch.tensor([8, 9, 11, 10], dtype=torch.int64), 4, "indices"
)
if __name__ == "__main__":
unittest.main()
test/registered/unit/server_args/test_server_args.py
+75
View File
@@ -371,6 +371,26 @@ class TestHiCacheArgs(CustomTestCase):
if expected_decode_backend is not None:
self.assertEqual(args.decode_attention_backend, expected_decode_backend)
def _make_cp_hicache_args(self, **overrides) -> ServerArgs:
args = self._make_args(
enable_hierarchical_cache=True,
enable_nsa_prefill_context_parallel=True,
enable_nsa_prefill_cp_shared_kv=True,
nsa_prefill_cp_mode="in-seq-split",
disaggregation_mode="prefill",
page_size=64,
enable_hisparse=False,
)
for key, value in overrides.items():
setattr(args, key, value)
return args
def _normalize_and_validate_cp_hicache_args(self, **overrides) -> ServerArgs:
args = self._make_cp_hicache_args(**overrides)
args._handle_hicache()
args._handle_cp_hicache_layout_validation()
return args
def test_cp_shared_kv_rejects_hicache_storage_backend(self):
with (
self.assertRaisesRegex(
@@ -468,6 +488,61 @@ class TestHiCacheArgs(CustomTestCase):
self.assertEqual(args.decode_attention_backend, "triton")
def test_cp_hicache_accepts_supported_backend_layout_pairs(self):
cases = [
("kernel", "layer_first"),
("kernel", "page_first"),
("direct", "layer_first"),
("direct", "page_first_direct"),
]
for io_backend, mem_layout in cases:
with self.subTest(io_backend=io_backend, mem_layout=mem_layout):
args = self._normalize_and_validate_cp_hicache_args(
hicache_io_backend=io_backend,
hicache_mem_layout=mem_layout,
)
self.assertEqual(args.hicache_io_backend, io_backend)
self.assertEqual(args.hicache_mem_layout, mem_layout)
def test_cp_hicache_normalizes_supported_alias_pairs(self):
args = self._normalize_and_validate_cp_hicache_args(
hicache_io_backend="kernel",
hicache_mem_layout="page_first_direct",
)
self.assertEqual(args.hicache_io_backend, "direct")
self.assertEqual(args.hicache_mem_layout, "page_first_direct")
args = self._normalize_and_validate_cp_hicache_args(
hicache_io_backend="direct",
hicache_mem_layout="page_first",
)
self.assertEqual(args.hicache_io_backend, "direct")
self.assertEqual(args.hicache_mem_layout, "page_first_direct")
def test_cp_hicache_rejects_page_head_layout(self):
with self.assertRaisesRegex(ValueError, "CP shared KV HiCache.*page_head"):
self._normalize_and_validate_cp_hicache_args(
hicache_io_backend="kernel",
hicache_mem_layout="page_head",
)
def test_cp_hicache_rejects_cuda_kv_split_layout(self):
with self.assertRaisesRegex(
ValueError, "CP shared KV HiCache.*page_first_kv_split"
):
self._normalize_and_validate_cp_hicache_args(
hicache_io_backend="kernel",
hicache_mem_layout="page_first_kv_split",
)
def test_cp_hicache_rejects_kernel_ascend_backend(self):
with self.assertRaisesRegex(ValueError, "CP shared KV HiCache.*kernel_ascend"):
self._normalize_and_validate_cp_hicache_args(
hicache_io_backend="kernel_ascend",
hicache_mem_layout="page_first_kv_split",
)
if __name__ == "__main__":
unittest.main()