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.
65 lines
2.2 KiB
Python
65 lines
2.2 KiB
Python
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()
|