test: harden CP HiCache metadata invariants
This commit is contained in:
@@ -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
|
||||
|
||||
@@ -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(
|
||||
|
||||
Reference in New Issue
Block a user