From 7d47f40a96b43e2f818173e914dd8a09d8edaa99 Mon Sep 17 00:00:00 2001
From: Zhangheng <hzh0425@apache.org>
Date: Fri, 17 Apr 2026 12:09:41 +0800
Subject: [PATCH] [UnifiedRadixTree]: Add HiCache hook interface for
 TreeComponent (#22924)

---
 .../tree_component.py                         | 37 +++++++++++++++++++
 1 file changed, 37 insertions(+)

diff --git a/python/sglang/srt/mem_cache/unified_cache_components/tree_component.py b/python/sglang/srt/mem_cache/unified_cache_components/tree_component.py
index 739db523b..31a2bad5d 100644
--- a/python/sglang/srt/mem_cache/unified_cache_components/tree_component.py
+++ b/python/sglang/srt/mem_cache/unified_cache_components/tree_component.py
@@ -17,6 +17,7 @@ from sglang.srt.mem_cache.base_prefix_cache import (
     MatchPrefixParams,
     MatchResult,
 )
+from sglang.srt.mem_cache.hicache_storage import PoolTransfer
 
 if TYPE_CHECKING:
     from sglang.srt.managers.schedule_batch import Req
@@ -62,6 +63,16 @@ class ComponentData:
     value: Optional[torch.Tensor] = None
     lock_ref: int = 0
     metadata: dict[str, Any] = dataclasses.field(default_factory=dict)
+    host_value: Optional[torch.Tensor] = None
+    host_lock_ref: int = 0
+
+
+class CacheTransferPhase(str, Enum):
+
+    BACKUP_HOST = "backup_host"  # D→H
+    LOAD_BACK = "load_back"  # H→D
+    BACKUP_STORAGE = "backup_storage"  # H→Storage
+    PREFETCH = "prefetch"  # Storage→H
 
 
 def get_and_increase_time_counter() -> float64:
@@ -290,3 +301,29 @@ class TreeComponent(ABC):
         ``insert_params`` is None only on the disabled path; on early-return
         paths it is still provided so components can free their resources."""
         pass
+
+    # ---- HiCache Hooks ----
+
+    def build_hicache_transfers(
+        self, node: UnifiedTreeNode, phase: CacheTransferPhase, **kw
+    ) -> Optional[list[PoolTransfer]]:
+        """Build transfer descriptors for this component in the given phase.
+        Returns None if the component has nothing to transfer."""
+        return None
+
+    def commit_hicache_transfer(
+        self,
+        node: UnifiedTreeNode,
+        phase: CacheTransferPhase,
+        transfers: list[PoolTransfer] = (),
+    ) -> None:
+        """Post-transfer bookkeeping: store host indices, update LRU, etc."""
+        pass
+
+    def drive_host_eviction(
+        self, num_tokens: int, tracker: dict[ComponentType, int]
+    ) -> None:
+        """Evict from this component's host-side resources.
+        Called by HostPoolGroup when the host pool is full.
+        Default no-op for components without host storage."""
+        pass