ThomasX
3a5928ef53
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.
47 lines
1.5 KiB
Python
47 lines
1.5 KiB
Python
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}"
|
|
)
|