fix: implement CP HiCache write retry eviction

python/sglang/srt/mem_cache/hiradix_cache.py

@@ -1073,6 +1073,56 @@ class HiRadixCache(RadixCache):
         self._delete_leaf(node)
         return num_evicted
 
+    def _remove_host_leaf(self, node: TreeNode) -> TreeNode:
+        parent = node.parent
+        key = self.get_child_key_fn(node.key)
+        v = parent.children.pop(key, None)
+        assert v == node, f"parent does not have child key, {key}"
+        if node in self.evictable_host_leaves:
+            self.evictable_host_leaves.remove(node)
+        self._update_host_leaf_status(parent)
+        return parent
+
+    def _evict_host_for_physical_slots(self, required_host_slots: int) -> int:
+        leaves = list(self.evictable_host_leaves)
+        eviction_heap = [
+            (self.eviction_strategy.get_priority(node), node) for node in leaves
+        ]
+        heapq.heapify(eviction_heap)
+
+        num_evicted = 0
+        while num_evicted < required_host_slots and len(eviction_heap):
+            _priority, x = heapq.heappop(eviction_heap)
+            if x == self.root_node:
+                break
+            if not x.evicted:
+                continue
+
+            if x.pin_expiry > 0 and time.monotonic() > x.pin_expiry:
+                self._clear_pin(x)
+
+            if x.host_ref_counter > 0:
+                continue
+
+            host_indices = (
+                x.cp_hicache.host_indices if x.cp_hicache is not None else None
+            )
+            physical_count = len(host_indices) if host_indices is not None else 0
+
+            self._record_remove_event(x)
+            if physical_count > 0:
+                num_evicted += self.cache_controller.evict_host(host_indices)
+            x.host_len = 0
+            x.cp_hicache = None
+            x.host_value = None
+
+            parent = self._remove_host_leaf(x)
+            if len(parent.children) == 0 and parent.evicted:
+                new_priority = self.eviction_strategy.get_priority(parent)
+                heapq.heappush(eviction_heap, (new_priority, parent))
+
+        return num_evicted
+
     def evict_host(self, num_tokens: int):
         leaves = list(self.evictable_host_leaves)
         eviction_heap = [
@@ -1102,16 +1152,11 @@ class HiRadixCache(RadixCache):
             self._record_remove_event(x)
             num_evicted += self.cache_controller.evict_host(x.host_value)
 
-            key = self.get_child_key_fn(x.key)
-            v = x.parent.children.pop(key, None)
-            assert v == x, f"parent does not have child key, {key}"
-            if x in self.evictable_host_leaves:
-                self.evictable_host_leaves.remove(x)
-            self._update_host_leaf_status(x.parent)
+            parent = self._remove_host_leaf(x)
 
-            if len(x.parent.children) == 0 and x.parent.evicted:
-                new_priority = self.eviction_strategy.get_priority(x.parent)
-                heapq.heappush(eviction_heap, (new_priority, x.parent))
+            if len(parent.children) == 0 and parent.evicted:
+                new_priority = self.eviction_strategy.get_priority(parent)
+                heapq.heappush(eviction_heap, (new_priority, parent))
 
     def load_back(
         self, node: TreeNode, mem_quota: Optional[int] = None