Hicache Storage Layer Prototype (#7704)

python/sglang/srt/mem_cache/radix_cache.py

@@ -55,8 +55,13 @@ class TreeNode:
         self.hit_count = 0
         # indicating the node is loading KV cache from host
         self.loading = False
+        # indicating the node is locked to protect from eviction
+        # incremented when the node is referenced by a storage operation
+        self.host_ref_counter = 0
         # store the host indices of KV cache
         self.host_value: Optional[torch.Tensor] = None
+        # store hash values of each pages
+        self.hash_value: Optional[List[str]] = None
 
         self.id = TreeNode.counter if id is None else id
         TreeNode.counter += 1
@@ -69,6 +74,27 @@ class TreeNode:
     def backuped(self):
         return self.host_value is not None
 
+    @property
+    def backuped_storage(self):
+        return self.hash_value is not None and len(self.hash_value) > 0
+
+    def protect_host(self):
+        """Protect the host value from eviction."""
+        self.host_ref_counter += 1
+
+    def release_host(self):
+        """Release the host value, allowing it to be evicted."""
+        if self.host_ref_counter > 0:
+            self.host_ref_counter -= 1
+        else:
+            raise RuntimeError("Host reference counter is already zero.")
+
+    def get_last_hash_value(self) -> Optional[str]:
+        """Returns the hash value of the last page in this node."""
+        if self.hash_value is None or len(self.hash_value) == 0:
+            return None
+        return self.hash_value[-1]
+
     def __lt__(self, other: "TreeNode"):
         return self.last_access_time < other.last_access_time