feat: split and evict CP HiCache metadata

2026-05-08 01:50:07 +08:00
parent bebbbfd0e8
commit f1a7e605e8
2 changed files with 91 additions and 1 deletions
@@ -306,5 +306,79 @@ class TestHiRadixCacheCPBackup(CustomTestCase):
        self.assertEqual(node.cp_hicache.host_indices.tolist(), [55])


+class TestHiRadixCacheCPSplitEvict(CustomTestCase):
+    def test_split_node_splits_cp_metadata_by_owned_positions(self):
+        cache = HiRadixCache.__new__(HiRadixCache)
+        cache._uses_cp_hicache = True
+        cache.get_child_key_fn = lambda key: key.token_ids[0]
+        cache.page_size = 1
+
+        root = TreeNode()
+        root.key = RadixKey([])
+        child = TreeNode()
+        child.parent = root
+        child.key = RadixKey(list(range(10)))
+        child.value = None
+        child.hash_value = None
+        child.host_len = 10
+        child.cp_hicache = 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),
+        )
+        root.children[0] = child
+
+        new_node = cache._split_node(child.key, child, 5)
+
+        self.assertEqual(new_node.host_len, 5)
+        self.assertEqual(child.host_len, 5)
+        self.assertEqual(new_node.cp_hicache.owned_positions.tolist(), [0, 2])
+        self.assertEqual(new_node.cp_hicache.host_indices.tolist(), [20, 21])
+        self.assertEqual(child.cp_hicache.owned_positions.tolist(), [0, 4])
+        self.assertEqual(child.cp_hicache.host_indices.tolist(), [22, 23])
+
+    def test_cp_host_eviction_uses_physical_freed_slots_for_progress(self):
+        cache = HiRadixCache.__new__(HiRadixCache)
+        cache._uses_cp_hicache = True
+        cache.root_node = TreeNode()
+        cache.root_node.key = RadixKey([])
+        cache.evictable_host_leaves = set()
+        cache.get_child_key_fn = lambda key: key.token_ids[0]
+        cache.eviction_strategy = type(
+            "Strategy", (), {"get_priority": lambda self, node: 0}
+        )()
+        cache._clear_pin = lambda node: None
+        cache._record_remove_event = lambda node: None
+        cache._update_host_leaf_status = lambda node: None
+        freed = []
+        cache.cache_controller = type(
+            "Controller",
+            (),
+            {
+                "evict_host": lambda self, indices: freed.append(indices.clone())
+                or len(indices)
+            },
+        )()
+        node = TreeNode()
+        node.parent = cache.root_node
+        node.key = RadixKey([1, 2, 3, 4])
+        node.value = None
+        node.host_len = 4
+        node.cp_hicache = CpHiCacheNodeMetadata(
+            logical_len=4,
+            owned_positions=torch.tensor([1], dtype=torch.int64),
+            host_indices=torch.tensor([70], dtype=torch.int64),
+        )
+        cache.root_node.children[1] = node
+        cache.evictable_host_leaves.add(node)
+
+        physical_freed = cache._evict_host_for_physical_slots(1)
+
+        self.assertEqual(physical_freed, 1)
+        self.assertEqual(freed[0].tolist(), [70])
+        self.assertEqual(node.host_len, 0)
+        self.assertIsNone(node.cp_hicache)
+
+
 if __name__ == "__main__":
    unittest.main()