test: harden CP HiCache metadata invariants

This commit is contained in:
2026-05-08 01:06:58 +08:00
parent 5fbe43381f
commit 1ede8bb999
2 changed files with 35 additions and 7 deletions

View File

@@ -64,8 +64,12 @@ class CpHiCacheNodeMetadata:
def __post_init__(self):
if self.logical_len < 0:
raise ValueError(f"logical_len must be non-negative, got {self.logical_len}")
self.owned_positions = self.owned_positions.to(device="cpu", dtype=torch.int64)
self.host_indices = self.host_indices.to(device="cpu", dtype=torch.int64)
self.owned_positions = self.owned_positions.to(
device="cpu", dtype=torch.int64
).detach().clone()
self.host_indices = self.host_indices.to(
device="cpu", dtype=torch.int64
).detach().clone()
if self.owned_positions.numel() != self.host_indices.numel():
raise ValueError(
"owned_positions and host_indices must have same length, got "
@@ -76,8 +80,10 @@ class CpHiCacheNodeMetadata:
self.owned_positions >= self.logical_len
):
raise ValueError("owned_positions must be in [0, logical_len)")
if torch.any(self.owned_positions[1:] < self.owned_positions[:-1]):
raise ValueError("owned_positions must be sorted")
if torch.any(self.owned_positions[1:] <= self.owned_positions[:-1]):
raise ValueError(
"owned_positions must be sorted and strictly increasing"
)
def split(
self, split_len: int

View File

@@ -1,6 +1,6 @@
import sys
import unittest
from unittest.mock import MagicMock, patch
from unittest.mock import MagicMock
import torch
@@ -10,8 +10,7 @@ for _mod in ("sgl_kernel", "sgl_kernel.kvcacheio"):
if _mod not in sys.modules:
sys.modules[_mod] = MagicMock()
from sglang.srt.mem_cache.hiradix_cache import CpHiCacheNodeMetadata, HiRadixCache
from sglang.srt.mem_cache.radix_cache import RadixKey, TreeNode
from sglang.srt.mem_cache.hiradix_cache import CpHiCacheNodeMetadata
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase
@@ -62,6 +61,21 @@ class TestCpHiCacheNodeMetadata(CustomTestCase):
self.assertEqual(metadata.owned_positions.device.type, "cpu")
self.assertEqual(metadata.host_indices.device.type, "cpu")
def test_metadata_does_not_alias_input_tensors(self):
owned_positions = torch.tensor([1, 3], dtype=torch.int64)
host_indices = torch.tensor([10, 11], dtype=torch.int64)
metadata = CpHiCacheNodeMetadata(
logical_len=4,
owned_positions=owned_positions,
host_indices=host_indices,
)
owned_positions[0] = 2
host_indices[0] = 12
self.assertEqual(metadata.owned_positions.tolist(), [1, 3])
self.assertEqual(metadata.host_indices.tolist(), [10, 11])
def test_invalid_split_raises(self):
metadata = CpHiCacheNodeMetadata(
logical_len=4,
@@ -80,6 +94,14 @@ class TestCpHiCacheNodeMetadata(CustomTestCase):
host_indices=torch.tensor([9, 10], dtype=torch.int64),
)
def test_duplicate_positions_raise(self):
with self.assertRaisesRegex(ValueError, "strictly increasing"):
CpHiCacheNodeMetadata(
logical_len=4,
owned_positions=torch.tensor([1, 1], dtype=torch.int64),
host_indices=torch.tensor([9, 10], dtype=torch.int64),
)
def test_length_mismatch_raises(self):
with self.assertRaisesRegex(ValueError, "same length"):
CpHiCacheNodeMetadata(