feat: split and evict CP HiCache metadata

python/sglang/srt/mem_cache/hiradix_cache.py

@@ -1084,6 +1084,9 @@ class HiRadixCache(RadixCache):
         return parent
 
     def _evict_host_for_physical_slots(self, required_host_slots: int) -> int:
+        if required_host_slots <= 0:
+            return 0
+
         leaves = list(self.evictable_host_leaves)
         eviction_heap = [
             (self.eviction_strategy.get_priority(node), node) for node in leaves
@@ -1098,6 +1101,9 @@ class HiRadixCache(RadixCache):
             if not x.evicted:
                 continue
 
+            if not self._node_backuped(x):
+                continue
+
             if x.pin_expiry > 0 and time.monotonic() > x.pin_expiry:
                 self._clear_pin(x)
 
@@ -1124,6 +1130,9 @@ class HiRadixCache(RadixCache):
         return num_evicted
 
     def evict_host(self, num_tokens: int):
+        if self._uses_cp_hicache:
+            return self._evict_host_for_physical_slots(num_tokens)
+
         leaves = list(self.evictable_host_leaves)
         eviction_heap = [
             (self.eviction_strategy.get_priority(node), node) for node in leaves
@@ -1596,7 +1605,14 @@ class HiRadixCache(RadixCache):
         else:
             new_node.value = child.value[:split_len].clone()
             child.value = child.value[split_len:].clone()
-        if child.backuped:
+        if self._uses_cp_hicache:
+            if self._node_backuped(child):
+                new_node.cp_hicache, child.cp_hicache = child.cp_hicache.split(
+                    split_len
+                )
+                new_node.host_len = split_len
+                child.host_len = child.host_len - split_len
+        elif child.backuped:
             new_node.host_value = child.host_value[:split_len].clone()
             child.host_value = child.host_value[split_len:].clone()
 

test/registered/unit/mem_cache/test_cp_hicache_metadata.py

@@ -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()