266 lines
9.6 KiB
Python
266 lines
9.6 KiB
Python
import sys
|
|
import unittest
|
|
from unittest.mock import MagicMock
|
|
|
|
import torch
|
|
|
|
# Stub out sgl_kernel before any sglang import so this CPU unit test does not
|
|
# require CUDA extension libraries to be installed.
|
|
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.test.ci.ci_register import register_cpu_ci
|
|
from sglang.test.test_utils import CustomTestCase
|
|
|
|
register_cpu_ci(est_time=2, suite="stage-a-test-cpu")
|
|
|
|
|
|
class TestCpHiCacheNodeMetadata(CustomTestCase):
|
|
def test_split_zero_len_moves_all_positions_to_child(self):
|
|
metadata = CpHiCacheNodeMetadata(
|
|
logical_len=8,
|
|
owned_positions=torch.tensor([1, 3, 7], dtype=torch.int64),
|
|
host_indices=torch.tensor([10, 11, 12], dtype=torch.int64),
|
|
)
|
|
|
|
parent, child = metadata.split(0)
|
|
|
|
self.assertEqual(parent.logical_len, 0)
|
|
self.assertEqual(parent.owned_positions.tolist(), [])
|
|
self.assertEqual(parent.host_indices.tolist(), [])
|
|
self.assertEqual(child.logical_len, 8)
|
|
self.assertEqual(child.owned_positions.tolist(), [1, 3, 7])
|
|
self.assertEqual(child.host_indices.tolist(), [10, 11, 12])
|
|
|
|
def test_split_rebases_child_positions(self):
|
|
metadata = CpHiCacheNodeMetadata(
|
|
logical_len=10,
|
|
owned_positions=torch.tensor([0, 2, 5, 9], dtype=torch.int64),
|
|
host_indices=torch.tensor([20, 21, 22, 23], dtype=torch.int64),
|
|
)
|
|
|
|
parent, child = metadata.split(5)
|
|
|
|
self.assertEqual(parent.logical_len, 5)
|
|
self.assertEqual(parent.owned_positions.tolist(), [0, 2])
|
|
self.assertEqual(parent.host_indices.tolist(), [20, 21])
|
|
self.assertEqual(child.logical_len, 5)
|
|
self.assertEqual(child.owned_positions.tolist(), [0, 4])
|
|
self.assertEqual(child.host_indices.tolist(), [22, 23])
|
|
|
|
def test_zero_owned_metadata_is_valid(self):
|
|
metadata = CpHiCacheNodeMetadata(
|
|
logical_len=64,
|
|
owned_positions=torch.empty((0,), dtype=torch.int32),
|
|
host_indices=torch.empty((0,), dtype=torch.int32),
|
|
)
|
|
|
|
self.assertEqual(metadata.logical_len, 64)
|
|
self.assertEqual(metadata.owned_positions.device.type, "cpu")
|
|
self.assertEqual(metadata.host_indices.device.type, "cpu")
|
|
self.assertEqual(metadata.owned_positions.dtype, torch.int64)
|
|
self.assertEqual(metadata.host_indices.dtype, torch.int64)
|
|
|
|
def test_non_int64_inputs_are_converted(self):
|
|
metadata = CpHiCacheNodeMetadata(
|
|
logical_len=4,
|
|
owned_positions=torch.tensor([1, 3], dtype=torch.int32),
|
|
host_indices=torch.tensor([10, 11], dtype=torch.int32),
|
|
)
|
|
|
|
self.assertEqual(metadata.owned_positions.dtype, torch.int64)
|
|
self.assertEqual(metadata.host_indices.dtype, torch.int64)
|
|
|
|
def test_negative_logical_len_raises(self):
|
|
with self.assertRaisesRegex(ValueError, "logical_len"):
|
|
CpHiCacheNodeMetadata(
|
|
logical_len=-1,
|
|
owned_positions=torch.empty((0,), dtype=torch.int64),
|
|
host_indices=torch.empty((0,), dtype=torch.int64),
|
|
)
|
|
|
|
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,
|
|
owned_positions=torch.tensor([1], dtype=torch.int64),
|
|
host_indices=torch.tensor([9], dtype=torch.int64),
|
|
)
|
|
|
|
with self.assertRaisesRegex(ValueError, "split_len"):
|
|
metadata.split(5)
|
|
|
|
def test_unsorted_positions_raise(self):
|
|
with self.assertRaisesRegex(ValueError, "sorted"):
|
|
CpHiCacheNodeMetadata(
|
|
logical_len=4,
|
|
owned_positions=torch.tensor([2, 1], dtype=torch.int64),
|
|
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(
|
|
logical_len=4,
|
|
owned_positions=torch.tensor([1, 2], dtype=torch.int64),
|
|
host_indices=torch.tensor([9], dtype=torch.int64),
|
|
)
|
|
|
|
def test_out_of_range_positions_raise(self):
|
|
with self.assertRaisesRegex(ValueError, r"\[0, logical_len\)"):
|
|
CpHiCacheNodeMetadata(
|
|
logical_len=4,
|
|
owned_positions=torch.tensor([4], dtype=torch.int64),
|
|
host_indices=torch.tensor([9], dtype=torch.int64),
|
|
)
|
|
|
|
|
|
class FakeWriteFailure:
|
|
metadata = None
|
|
|
|
def __init__(self, required_host_slots):
|
|
self.required_host_slots = required_host_slots
|
|
|
|
|
|
class FakeWriteSuccess:
|
|
required_host_slots = 0
|
|
|
|
def __init__(self, metadata):
|
|
self.metadata = metadata
|
|
|
|
|
|
class FakeWriteController:
|
|
def __init__(self, required_host_slots):
|
|
self.required_host_slots = required_host_slots
|
|
self.calls = 0
|
|
self.write_policy = "write_through"
|
|
self.evicted_host_indices = []
|
|
|
|
def write(self, device_indices, node_id=-1, priority=None):
|
|
self.calls += 1
|
|
if self.calls == 1:
|
|
return FakeWriteFailure(self.required_host_slots)
|
|
return FakeWriteSuccess(
|
|
CpHiCacheNodeMetadata(
|
|
logical_len=len(device_indices),
|
|
owned_positions=torch.tensor([0], dtype=torch.int64),
|
|
host_indices=torch.tensor([99], dtype=torch.int64),
|
|
)
|
|
)
|
|
|
|
def evict_host(self, host_indices):
|
|
self.evicted_host_indices.append(host_indices.clone())
|
|
return len(host_indices)
|
|
|
|
|
|
class FakeEvictionStrategy:
|
|
def get_priority(self, node):
|
|
return 0
|
|
|
|
|
|
class TestHiRadixCacheCPBackup(CustomTestCase):
|
|
def test_node_backuped_uses_cp_metadata(self):
|
|
cache = HiRadixCache.__new__(HiRadixCache)
|
|
cache._uses_cp_hicache = True
|
|
node = TreeNode()
|
|
node.host_len = 8
|
|
node.cp_hicache = CpHiCacheNodeMetadata(
|
|
logical_len=8,
|
|
owned_positions=torch.tensor([1, 2], dtype=torch.int64),
|
|
host_indices=torch.tensor([10, 11], dtype=torch.int64),
|
|
)
|
|
|
|
self.assertTrue(cache._node_backuped(node))
|
|
|
|
def test_inc_hit_count_does_not_rewrite_cp_backed_node(self):
|
|
cache = HiRadixCache.__new__(HiRadixCache)
|
|
cache._uses_cp_hicache = True
|
|
cache.write_through_threshold = 1
|
|
cache.cache_controller = type(
|
|
"Controller", (), {"write_policy": "write_through"}
|
|
)()
|
|
cache.write_backup = lambda node: (_ for _ in ()).throw(
|
|
AssertionError("must not rewrite")
|
|
)
|
|
node = TreeNode()
|
|
node.host_len = 4
|
|
node.cp_hicache = CpHiCacheNodeMetadata(
|
|
logical_len=4,
|
|
owned_positions=torch.tensor([], dtype=torch.int64),
|
|
host_indices=torch.tensor([], dtype=torch.int64),
|
|
)
|
|
|
|
cache._inc_hit_count(node)
|
|
|
|
self.assertEqual(node.hit_count, 1)
|
|
|
|
def test_write_backup_retries_by_required_physical_slots(self):
|
|
cache = HiRadixCache.__new__(HiRadixCache)
|
|
cache._uses_cp_hicache = True
|
|
cache.cache_controller = FakeWriteController(required_host_slots=1)
|
|
cache.evictable_host_leaves = set()
|
|
cache.eviction_strategy = FakeEvictionStrategy()
|
|
cache.get_child_key_fn = lambda key: key.token_ids[0]
|
|
cache._record_remove_event = lambda node: None
|
|
cache.ongoing_write_through = {}
|
|
cache.inc_node_lock_ref = lambda node: None
|
|
|
|
root = TreeNode()
|
|
root.key = RadixKey(token_ids=[], extra_key=None)
|
|
root.value = []
|
|
cache.root_node = root
|
|
evictable_node = TreeNode()
|
|
evictable_node.parent = root
|
|
evictable_node.key = RadixKey(token_ids=[1], extra_key=None)
|
|
evictable_node.value = None
|
|
evictable_node.host_len = 4
|
|
evictable_node.cp_hicache = CpHiCacheNodeMetadata(
|
|
logical_len=4,
|
|
owned_positions=torch.tensor([0], dtype=torch.int64),
|
|
host_indices=torch.tensor([55], dtype=torch.int64),
|
|
)
|
|
root.children[1] = evictable_node
|
|
cache.evictable_host_leaves.add(evictable_node)
|
|
|
|
node = TreeNode()
|
|
node.value = torch.arange(16, dtype=torch.int64)
|
|
|
|
cache.write_backup(node)
|
|
|
|
self.assertEqual(
|
|
cache.cache_controller.evicted_host_indices[0].tolist(), [55]
|
|
)
|
|
self.assertEqual(evictable_node.host_len, 0)
|
|
self.assertIsNone(evictable_node.cp_hicache)
|
|
self.assertNotIn(1, root.children)
|
|
self.assertEqual(node.host_len, 16)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|