HiCache Storage: generate hash when inserting new nodes (#9053)

python/sglang/srt/mem_cache/hiradix_cache.py

@@ -151,7 +151,7 @@ class HiRadixCache(RadixCache):
 
     def write_backup_storage(self, node: TreeNode):
         operation_id = self.cache_controller.write_storage(
-            node.host_value, node.key, node.parent.get_last_hash_value()
+            node.host_value, node.key, node.hash_value
         )
         self.ongoing_backup[operation_id] = node
         node.protect_host()
@@ -414,18 +414,18 @@ class HiRadixCache(RadixCache):
                 group=self.tp_group,
             )
         for _ in range(queue_size.item()):
-            ack_id, hash_value, completed_tokens = (
-                self.cache_controller.ack_backup_queue.get()
-            )
+            ack_id, completed_tokens = self.cache_controller.ack_backup_queue.get()
             host_node = self.ongoing_backup[ack_id]
-            if completed_tokens == 0:
-                host_node.hash_value = None
-            elif completed_tokens < len(host_node.key):
-                # backup is only partially successful, split the node
-                new_node = self._split_node(host_node.key, host_node, completed_tokens)
-                new_node.hash_value = hash_value
-            else:
-                host_node.hash_value = hash_value
+
+            if completed_tokens > 0:
+                if completed_tokens < len(host_node.key):
+                    # backup is only partially successful, split the node
+                    new_node = self._split_node(
+                        host_node.key, host_node, completed_tokens
+                    )
+                    new_node.backuped_storage = True
+                else:
+                    host_node.backuped_storage = True
             host_node.release_host()
             del self.ongoing_backup[ack_id]
 
@@ -717,6 +717,21 @@ class HiRadixCache(RadixCache):
             node.children[child_key] = new_node
             self.evictable_size_ += len(value)
 
+            if self.enable_storage:
+                last_hash = node.get_last_hash_value()
+                assert (node == self.root_node) or (
+                    last_hash is not None
+                ), "Parent node must have a hash value with storage enabled"
+                new_node.hash_value = []
+                for idx in range(0, len(key), self.page_size):
+                    new_node.hash_value.append(
+                        self.cache_controller.get_hash_str(
+                            key[idx : idx + self.page_size],
+                            prior_hash=last_hash,
+                        )
+                    )
+                    last_hash = new_node.hash_value[-1]
+
             if self.cache_controller.write_policy != "write_back":
                 self.inc_hit_count(new_node)
         return total_prefix_length