fix: treat CP HiCache metadata as backup state
This commit is contained in:
@@ -10,6 +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.base_prefix_cache import EvictParams
|
||||
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
|
||||
@@ -183,6 +184,13 @@ class FakeEvictionStrategy:
|
||||
return 0
|
||||
|
||||
|
||||
class FakeEvictDeviceController:
|
||||
write_policy = "write_through"
|
||||
|
||||
def evict_device(self, device_indices):
|
||||
return len(device_indices)
|
||||
|
||||
|
||||
class TestHiRadixCacheCPBackup(CustomTestCase):
|
||||
def test_node_backuped_uses_cp_metadata(self):
|
||||
cache = HiRadixCache.__new__(HiRadixCache)
|
||||
@@ -260,6 +268,43 @@ class TestHiRadixCacheCPBackup(CustomTestCase):
|
||||
self.assertNotIn(1, root.children)
|
||||
self.assertEqual(node.host_len, 16)
|
||||
|
||||
def test_evict_demotes_cp_backed_node_without_deleting_radix_child(self):
|
||||
cache = HiRadixCache.__new__(HiRadixCache)
|
||||
cache._uses_cp_hicache = True
|
||||
cache.cache_controller = FakeEvictDeviceController()
|
||||
cache.evictable_leaves = set()
|
||||
cache.evictable_host_leaves = set()
|
||||
cache.eviction_strategy = FakeEvictionStrategy()
|
||||
cache.evictable_size_ = 4
|
||||
cache.protected_size_ = 0
|
||||
cache._record_remove_event = lambda node: (_ for _ in ()).throw(
|
||||
AssertionError("must not delete backed radix child")
|
||||
)
|
||||
|
||||
root = TreeNode()
|
||||
root.key = RadixKey(token_ids=[], extra_key=None)
|
||||
root.value = []
|
||||
cache.root_node = root
|
||||
node = TreeNode()
|
||||
node.parent = root
|
||||
node.key = RadixKey(token_ids=[1, 2, 3, 4], extra_key=None)
|
||||
node.value = torch.arange(4, dtype=torch.int64)
|
||||
node.host_len = 4
|
||||
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] = node
|
||||
cache.evictable_leaves.add(node)
|
||||
|
||||
cache.evict(EvictParams(num_tokens=4))
|
||||
|
||||
self.assertIn(1, root.children)
|
||||
self.assertIsNone(node.value)
|
||||
self.assertIsNotNone(node.cp_hicache)
|
||||
self.assertEqual(node.cp_hicache.host_indices.tolist(), [55])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
Reference in New Issue
Block a user