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:
@@ -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) {
|
||||
|
||||
Reference in New Issue
Block a user