From 3a5928ef538833b41e23798f7488c772cb7ee6d5 Mon Sep 17 00:00:00 2001 From: ThomasX Date: Sun, 10 May 2026 02:02:54 +0800 Subject: [PATCH] feat: add page-aligned index validation for CP HiCache direct transfers 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. --- .../sglang/srt/managers/cache_controller.py | 27 +++++- .../sglang/srt/mem_cache/memory_pool_host.py | 9 +- .../sglang/srt/mem_cache/page_index_utils.py | 46 ++++++++++ python/sglang/srt/server_args.py | 28 ++++++ sgl-kernel/csrc/kvcacheio/transfer.cu | 34 ++++++++ sgl-kernel/tests/test_kvcacheio.py | 57 ++++++++++++ .../managers/test_hicache_controller_cp.py | 63 ++++++++++++++ .../unit/mem_cache/test_nsa_pool_host_unit.py | 87 ++++++++++++++++--- .../unit/mem_cache/test_page_index_utils.py | 64 ++++++++++++++ .../unit/server_args/test_server_args.py | 75 ++++++++++++++++ 10 files changed, 472 insertions(+), 18 deletions(-) create mode 100644 python/sglang/srt/mem_cache/page_index_utils.py create mode 100644 test/registered/unit/mem_cache/test_page_index_utils.py diff --git a/python/sglang/srt/managers/cache_controller.py b/python/sglang/srt/managers/cache_controller.py index e38c46c61..c2d4252b7 100644 --- a/python/sglang/srt/managers/cache_controller.py +++ b/python/sglang/srt/managers/cache_controller.py @@ -40,6 +40,7 @@ from sglang.srt.layers.dp_attention import ( ) from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout from sglang.srt.mem_cache.memory_pool import MLATokenToKVPool, NSATokenToKVPool +from sglang.srt.mem_cache.page_index_utils import validate_page_aligned_token_indices from sglang.srt.utils import get_device_module logger = logging.getLogger(__name__) @@ -735,6 +736,16 @@ class HiCacheController: event.record() self.ack_load_queue.append(HiCacheAck(event, event, [node_id])) + def _validate_cp_hicache_page_indices( + self, + host_indices: torch.Tensor, + device_indices: torch.Tensor, + ) -> None: + validate_page_aligned_token_indices(host_indices, self.page_size, "host_indices") + validate_page_aligned_token_indices( + device_indices, self.page_size, "physical_device_indices" + ) + def _write_cp( self, device_indices: torch.Tensor, @@ -762,6 +773,11 @@ class HiCacheController: host_indices = self.mem_pool_host.alloc(len(physical_device_indices)) if host_indices is None: return HiCacheWriteFailure(required_host_slots=len(physical_device_indices)) + try: + self._validate_cp_hicache_page_indices(host_indices, physical_device_indices) + except Exception: + self.mem_pool_host.free(host_indices) + raise self.write_queue.append( CacheOperation(host_indices, physical_device_indices, node_id, priority) @@ -817,10 +833,17 @@ class HiCacheController: self._append_completed_load_ack(node_id) return device_indices + host_indices = torch.cat(host_chunks) + physical_device_indices = torch.cat(physical_chunks) + try: + self._validate_cp_hicache_page_indices(host_indices, physical_device_indices) + except Exception: + self.mem_pool_device_allocator.free(device_indices) + raise self.load_queue.append( CacheOperation( - torch.cat(host_chunks), - torch.cat(physical_chunks), + host_indices, + physical_device_indices, node_id, ) ) diff --git a/python/sglang/srt/mem_cache/memory_pool_host.py b/python/sglang/srt/mem_cache/memory_pool_host.py index ea220bef0..aa4a72fdb 100644 --- a/python/sglang/srt/mem_cache/memory_pool_host.py +++ b/python/sglang/srt/mem_cache/memory_pool_host.py @@ -23,6 +23,7 @@ from sglang.srt.mem_cache.memory_pool import ( MLATokenToKVPool, NSATokenToKVPool, ) +from sglang.srt.mem_cache.page_index_utils import validate_page_aligned_token_indices from sglang.srt.utils import is_cuda, is_mps, is_npu, is_xpu _is_cuda = is_cuda() @@ -1169,10 +1170,10 @@ class NSATokenToKVPoolHost(MLATokenToKVPoolHost): def _get_indexer_page_indices(self, host_indices, device_indices): if host_indices.numel() == 0: return host_indices, device_indices - if host_indices.numel() % self.page_size != 0: - raise ValueError( - "Index buffer transfer expects page-aligned indices for NSA." - ) + validate_page_aligned_token_indices(host_indices, self.page_size, "host_indices") + validate_page_aligned_token_indices( + device_indices, self.page_size, "device_indices" + ) host_page_indices = ( host_indices.reshape(-1, self.page_size)[:, 0] // self.page_size ) diff --git a/python/sglang/srt/mem_cache/page_index_utils.py b/python/sglang/srt/mem_cache/page_index_utils.py new file mode 100644 index 000000000..48c51cf75 --- /dev/null +++ b/python/sglang/srt/mem_cache/page_index_utils.py @@ -0,0 +1,46 @@ +from __future__ import annotations + +import torch + + +def validate_page_aligned_token_indices( + indices: torch.Tensor, + page_size: int, + name: str, +) -> None: + """Validate that token indices are grouped as complete contiguous pages. + + HiCache page-first direct transfers treat every `page_size` entries as one + page and derive the page id from the first token. A malformed group can + make native direct-copy code address the wrong page, so validate the exact + invariant before entering transfer code. + """ + + if page_size <= 0: + raise ValueError(f"page_size must be positive, got {page_size}") + if indices.dim() != 1: + raise ValueError( + f"{name} must be a 1-D tensor, got shape={tuple(indices.shape)}" + ) + if indices.numel() == 0: + return + if indices.numel() % page_size != 0: + raise ValueError( + f"{name} must contain whole pages: numel={indices.numel()} page_size={page_size}" + ) + + page_view = indices.reshape(-1, page_size) + starts = page_view[:, 0] + if torch.any(torch.remainder(starts, page_size) != 0): + raise ValueError( + f"{name} page groups must start at page boundaries: page_size={page_size}" + ) + + expected_offsets = torch.arange( + page_size, device=indices.device, dtype=indices.dtype + ) + expected = starts[:, None] + expected_offsets + if not torch.equal(page_view, expected): + raise ValueError( + f"{name} page groups must be contiguous page spans: page_size={page_size}" + ) diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 82e0f14e0..9acb65d1d 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -798,6 +798,7 @@ class ServerArgs: # Handle Hicache settings. self._handle_hicache() + self._handle_cp_hicache_layout_validation() # Handle data parallelism. self._handle_data_parallelism() @@ -917,6 +918,33 @@ class ServerArgs: "Disable hicache_storage_backend or disable CP shared KV." ) + def _handle_cp_hicache_layout_validation(self): + if not ( + self.enable_nsa_prefill_cp_shared_kv and self.enable_hierarchical_cache + ): + return + + supported_pairs = { + ("kernel", "layer_first"), + ("kernel", "page_first"), + ("direct", "layer_first"), + ("direct", "page_first_direct"), + } + current_pair = (self.hicache_io_backend, self.hicache_mem_layout) + if current_pair in supported_pairs: + return + + raise ValueError( + "CP shared KV HiCache supports only " + "kernel/layer_first, kernel/page_first, direct/layer_first, " + "and direct/page_first_direct after HiCache normalization. " + f"Got hicache_io_backend={self.hicache_io_backend!r} " + f"hicache_mem_layout={self.hicache_mem_layout!r}. " + "page_head is MHA/Mooncake-specific, page_first_kv_split is " + "kernel_ascend-specific, and kernel_ascend CP shared KV HiCache " + "is not implemented in the CUDA NSA host path." + ) + def _handle_deprecated_args(self): # Handle deprecated tool call parsers deprecated_tool_call_parsers = {"qwen25": "qwen", "glm45": "glm"} diff --git a/sgl-kernel/csrc/kvcacheio/transfer.cu b/sgl-kernel/csrc/kvcacheio/transfer.cu index 7b1e264e0..a5e990e4f 100644 --- a/sgl-kernel/csrc/kvcacheio/transfer.cu +++ b/sgl-kernel/csrc/kvcacheio/transfer.cu @@ -686,6 +686,38 @@ inline void transfer_page_direct( /* non_blocking= */ true); } +inline void validate_page_first_direct_indices( + const int64_t* indices_ptr, int64_t num_indices, int64_t page_size, const char* name) { + if (num_indices == 0) { + return; + } + for (int64_t i = 0; i < num_indices; i += page_size) { + const int64_t start = indices_ptr[i]; + TORCH_CHECK( + start % page_size == 0, + name, + " page groups must start at page boundaries: group=", + i / page_size, + " start=", + start, + " page_size=", + page_size); + for (int64_t offset = 1; offset < page_size; ++offset) { + TORCH_CHECK( + indices_ptr[i + offset] == start + offset, + name, + " page groups must be contiguous page spans: group=", + i / page_size, + " offset=", + offset, + " expected=", + start + offset, + " got=", + indices_ptr[i + offset]); + } + } +} + void transfer_kv_direct( const std::vector& src_layers, std::vector dst_layers, @@ -749,6 +781,8 @@ inline void transfer_kv_page_first_direct_impl( const int64_t num_pages = src_indices_cpu.size(0) / page_size; int64_t* src_indices_ptr = src_indices_cpu.data_ptr(); int64_t* dst_indices_ptr = dst_indices_cpu.data_ptr(); + validate_page_first_direct_indices(src_indices_ptr, src_indices_cpu.numel(), page_size, "src_indices"); + validate_page_first_direct_indices(dst_indices_ptr, dst_indices_cpu.numel(), page_size, "dst_indices"); auto fallback_to_page_copy = [&]() { if constexpr (IsLf2Pf) { diff --git a/sgl-kernel/tests/test_kvcacheio.py b/sgl-kernel/tests/test_kvcacheio.py index 16fc7e826..b9a3e2f60 100644 --- a/sgl-kernel/tests/test_kvcacheio.py +++ b/sgl-kernel/tests/test_kvcacheio.py @@ -33,6 +33,15 @@ def ref_copy_with_indices_pf_direct( ][layer_id].to(dst_pool.device) +def make_page_indices(page_ids, page_size): + return torch.cat( + [ + torch.arange(p * page_size, (p + 1) * page_size, dtype=torch.int64) + for p in page_ids + ] + ) + + def ref_copy_with_indices_page_head( src_pool, dst_pool, @@ -522,6 +531,54 @@ def test_transfer_kv_pf_direct( torch.set_default_dtype(original_dtype) +@pytest.mark.skipif(is_hip(), reason="HIP uses the fallback path for this direct validation") +def test_transfer_kv_all_layer_direct_lf_pf_rejects_misaligned_page_start(): + page_size = 4 + item_size = 8 + num_layers = 1 + src_pool = torch.randn(num_layers, 32, item_size, device="cuda") + dst_pool = torch.zeros(8, num_layers, page_size, item_size).pin_memory() + src_indices = torch.tensor([1, 2, 3, 4], dtype=torch.int64) + dst_indices = make_page_indices([0], page_size) + + with pytest.raises(RuntimeError, match="page boundaries"): + transfer_kv_all_layer_direct_lf_pf( + [src_pool[0]], [dst_pool], src_indices, dst_indices, page_size + ) + + +@pytest.mark.skipif(is_hip(), reason="HIP uses the fallback path for this direct validation") +def test_transfer_kv_all_layer_direct_lf_pf_rejects_non_contiguous_page_group(): + page_size = 4 + item_size = 8 + num_layers = 1 + src_pool = torch.randn(num_layers, 32, item_size, device="cuda") + dst_pool = torch.zeros(8, num_layers, page_size, item_size).pin_memory() + src_indices = torch.tensor([4, 5, 7, 6], dtype=torch.int64) + dst_indices = make_page_indices([0], page_size) + + with pytest.raises(RuntimeError, match="contiguous page spans"): + transfer_kv_all_layer_direct_lf_pf( + [src_pool[0]], [dst_pool], src_indices, dst_indices, page_size + ) + + +@pytest.mark.skipif(is_hip(), reason="HIP uses the fallback path for this direct validation") +def test_transfer_kv_per_layer_direct_pf_lf_rejects_misaligned_host_page_start(): + page_size = 4 + item_size = 8 + layer_id = 0 + src_pool = torch.randn(8, 1, page_size, item_size).pin_memory() + dst_pool = torch.zeros(1, 32, item_size, device="cuda") + src_indices = torch.tensor([1, 2, 3, 4], dtype=torch.int64) + dst_indices = make_page_indices([0], page_size).to("cuda") + + with pytest.raises(RuntimeError, match="page boundaries"): + transfer_kv_per_layer_direct_pf_lf( + [src_pool], [dst_pool[0]], src_indices, dst_indices, layer_id, page_size + ) + + @pytest.mark.skipif(is_hip(), reason="HIP is not supported for this test") @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) @pytest.mark.parametrize("num_items_to_transfer", [256, 1024]) diff --git a/test/registered/unit/managers/test_hicache_controller_cp.py b/test/registered/unit/managers/test_hicache_controller_cp.py index 02159f779..6616c8fdb 100644 --- a/test/registered/unit/managers/test_hicache_controller_cp.py +++ b/test/registered/unit/managers/test_hicache_controller_cp.py @@ -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() diff --git a/test/registered/unit/mem_cache/test_nsa_pool_host_unit.py b/test/registered/unit/mem_cache/test_nsa_pool_host_unit.py index f75819557..b35798485 100644 --- a/test/registered/unit/mem_cache/test_nsa_pool_host_unit.py +++ b/test/registered/unit/mem_cache/test_nsa_pool_host_unit.py @@ -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__": diff --git a/test/registered/unit/mem_cache/test_page_index_utils.py b/test/registered/unit/mem_cache/test_page_index_utils.py new file mode 100644 index 000000000..9398508bb --- /dev/null +++ b/test/registered/unit/mem_cache/test_page_index_utils.py @@ -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() diff --git a/test/registered/unit/server_args/test_server_args.py b/test/registered/unit/server_args/test_server_args.py index 0fafd503d..7937c0589 100644 --- a/test/registered/unit/server_args/test_server_args.py +++ b/test/registered/unit/server_args/test_server_args.py @@ -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()