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.
This commit is contained in:
2026-05-10 02:02:54 +08:00
parent f38937723e
commit 3a5928ef53
10 changed files with 472 additions and 18 deletions
+34
View File
@@ -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<at::Tensor>& src_layers,
std::vector<at::Tensor> 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>();
int64_t* dst_indices_ptr = dst_indices_cpu.data_ptr<int64_t>();
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) {