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.

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,
             )
         )

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
         )

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}"
+        )

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"}