hicache kv events: publish split write-through fragments (#27072)

python/sglang/srt/mem_cache/hiradix_cache.py

@@ -714,8 +714,7 @@ class HiRadixCache(RadixCache):
         if host_indices is not None:
             node.host_value = host_indices.clone()
             assert len(node.host_value) > 0
-            # Record backup_len for ack-time walk-and-concat after split.
-            self.ongoing_write_through[node.id] = (node, len(node.key))
+            self._track_write_through_node(node, len(node.key))
             if not write_back:
                 self.inc_lock_ref(node)
         else:
@@ -723,6 +722,50 @@ class HiRadixCache(RadixCache):
 
         return len(host_indices)
 
+    def _track_write_through_node(self, node: TreeNode, backup_len: int) -> None:
+        node.write_through_pending_id = node.id
+        self.ongoing_write_through[node.id] = (node, backup_len, [node])
+
+    def _replace_pending_write_through_node(
+        self, old_node: TreeNode, new_nodes: List[TreeNode]
+    ) -> None:
+        ack_id = old_node.write_through_pending_id
+        if ack_id is None:
+            return
+
+        pending = self.ongoing_write_through.get(ack_id)
+        if pending is None:
+            return
+
+        lock_node, backup_len, publish_nodes = pending
+        updated_nodes = []
+        replaced = False
+        for node in publish_nodes:
+            if node is old_node:
+                updated_nodes.extend(new_nodes)
+                replaced = True
+            else:
+                updated_nodes.append(node)
+
+        if not replaced:
+            return
+
+        for node in new_nodes:
+            node.write_through_pending_id = ack_id
+        self.ongoing_write_through[ack_id] = (lock_node, backup_len, updated_nodes)
+
+    def _finish_write_through_ack(self, ack_id: int, *, release_lock: bool) -> None:
+        lock_node, backup_len, publish_nodes = self.ongoing_write_through.pop(ack_id)
+        for node in publish_nodes:
+            if node.write_through_pending_id == ack_id:
+                node.write_through_pending_id = None
+            # DMA confirmed -- block is now on host.
+            self._record_store_event(node, medium=StorageMedium.CPU)
+        if self.enable_storage:
+            self.write_backup_storage(lock_node, backup_len)
+        if release_lock:
+            self.dec_lock_ref(lock_node)
+
     def write_backup_storage(self, node: TreeNode, backup_len: Optional[int] = None):
         # Recover pre-split data via walk-and-concat if node was split.
         # prefix_keys anchored at chain top to avoid double-counting.
@@ -802,11 +845,7 @@ class HiRadixCache(RadixCache):
                 for _, finish_event, ack_list in self.cache_controller.ack_write_queue:
                     finish_event.synchronize()
                     for ack_id in ack_list:
-                        node, backup_len = self.ongoing_write_through.pop(ack_id)
-                        # DMA confirmed -- block is now on host.
-                        self._record_store_event(node, medium=StorageMedium.CPU)
-                        if self.enable_storage:
-                            self.write_backup_storage(node, backup_len)
+                        self._finish_write_through_ack(ack_id, release_lock=False)
                 self.cache_controller.ack_write_queue.clear()
                 assert len(self.ongoing_write_through) == 0
             return
@@ -829,12 +868,7 @@ class HiRadixCache(RadixCache):
             _, finish_event, ack_list = self.cache_controller.ack_write_queue.pop(0)
             finish_event.synchronize()
             for ack_id in ack_list:
-                node, backup_len = self.ongoing_write_through.pop(ack_id)
-                # DMA confirmed -- block is now on host.
-                self._record_store_event(node, medium=StorageMedium.CPU)
-                self.dec_lock_ref(node)
-                if self.enable_storage:
-                    self.write_backup_storage(node, backup_len)
+                self._finish_write_through_ack(ack_id, release_lock=True)
             finish_count -= 1
 
     def loading_check(self):
@@ -1506,6 +1540,9 @@ class HiRadixCache(RadixCache):
         child.key = child.key[split_len:]
         new_node.parent.children[key.child_key(self.page_size)] = new_node
 
+        if child.backuped:
+            self._replace_pending_write_through_node(child, [new_node, child])
+
         return new_node
 
     def insert(self, params: InsertParams) -> InsertResult:

python/sglang/srt/mem_cache/radix_cache.py

@@ -214,6 +214,7 @@ class TreeNode:
         self.host_ref_counter = 0
         # store the host indices of KV cache
         self.host_value: Optional[torch.Tensor] = None
+        self.write_through_pending_id: Optional[int] = None
         # store hash values of each pages
         self.hash_value: Optional[List[str]] = None
         # priority for priority-aware eviction

python/sglang/srt/mem_cache/unified_radix_cache.py

@@ -95,6 +95,7 @@ class UnifiedTreeNode:
         )
         self.id = UnifiedTreeNode.counter
         UnifiedTreeNode.counter += 1
+        self.write_through_pending_id: Optional[int] = None
 
     def component(self, component_type: ComponentType) -> ComponentData:
         return self.component_data[component_type]
@@ -369,7 +370,12 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache):
             for ct in self.tree_components
         }
         self.ongoing_write_through: dict[
-            int, tuple[UnifiedTreeNode, Optional[DecLockRefParams]]
+            int,
+            tuple[
+                UnifiedTreeNode,
+                Optional[DecLockRefParams],
+                list[UnifiedTreeNode],
+            ],
         ] = {}
         self.ongoing_load_back: dict[int, tuple[UnifiedTreeNode, DecLockRefParams]] = {}
         self.enable_storage = False
@@ -915,6 +921,9 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache):
             component.redistribute_on_node_split(new_parent=new_node, child=child)
         new_node.parent.children[key.child_key(self.page_size)] = new_node
 
+        if child.backuped:
+            self._replace_pending_write_through_node(child, [new_node, child])
+
         self._for_each_component_lru(
             new_node, UnifiedLRUList.insert_mru, skip_existing=True
         )
@@ -1466,9 +1475,63 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache):
         lock_params = None
         if not write_back:
             lock_params = self.inc_lock_ref(node).to_dec_params()
-        self.ongoing_write_through[node.id] = (node, lock_params)
+        self._track_write_through_node(node, lock_params)
         return len(host_indices)
 
+    def _track_write_through_node(
+        self,
+        node: UnifiedTreeNode,
+        lock_params: Optional[DecLockRefParams],
+    ) -> None:
+        node.write_through_pending_id = node.id
+        self.ongoing_write_through[node.id] = (node, lock_params, [node])
+
+    def _replace_pending_write_through_node(
+        self, old_node: UnifiedTreeNode, new_nodes: list[UnifiedTreeNode]
+    ) -> None:
+        ack_id = old_node.write_through_pending_id
+        if ack_id is None:
+            return
+
+        pending = self.ongoing_write_through.get(ack_id)
+        if pending is None:
+            return
+
+        lock_node, lock_params, publish_nodes = pending
+        updated_nodes = []
+        replaced = False
+        for node in publish_nodes:
+            if node is old_node:
+                updated_nodes.extend(new_nodes)
+                replaced = True
+            else:
+                updated_nodes.append(node)
+
+        if not replaced:
+            return
+
+        for node in new_nodes:
+            node.write_through_pending_id = ack_id
+        self.ongoing_write_through[ack_id] = (
+            lock_node,
+            lock_params,
+            updated_nodes,
+        )
+
+    def _finish_write_through_ack(self, ack_id: int) -> None:
+        lock_node, lock_params, publish_nodes = self.ongoing_write_through.pop(ack_id)
+        for node in publish_nodes:
+            if node.write_through_pending_id == ack_id:
+                node.write_through_pending_id = None
+            self._record_store_event(node, medium=StorageMedium.CPU)
+        if lock_params is not None:
+            self.dec_lock_ref(lock_node, lock_params)
+        if self.enable_storage:
+            # Back up each fragment: after a split, lock_node only holds the
+            # suffix; the prefix fragment must be persisted as well.
+            for node in publish_nodes:
+                self.write_backup_storage(node)
+
     def load_back(
         self,
         best_match_node: UnifiedTreeNode,
@@ -2141,14 +2204,8 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache):
                 for _, finish_event, ack_list in cc.ack_write_queue:
                     finish_event.synchronize()
                     for ack_id in ack_list:
-                        entry = self.ongoing_write_through.pop(ack_id, None)
-                        if entry is not None:
-                            node, params = entry
-                            self._record_store_event(node, medium=StorageMedium.CPU)
-                            if params is not None:
-                                self.dec_lock_ref(node, params)
-                            if self.enable_storage:
-                                self.write_backup_storage(node)
+                        if ack_id in self.ongoing_write_through:
+                            self._finish_write_through_ack(ack_id)
                 cc.ack_write_queue.clear()
                 assert len(self.ongoing_write_through) == 0
             return
@@ -2172,11 +2229,7 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache):
             _, finish_event, ack_list = cc.ack_write_queue.pop(0)
             finish_event.synchronize()
             for ack_id in ack_list:
-                node, params = self.ongoing_write_through.pop(ack_id)
-                self._record_store_event(node, medium=StorageMedium.CPU)
-                self.dec_lock_ref(node, params)
-                if self.enable_storage:
-                    self.write_backup_storage(node)
+                self._finish_write_through_ack(ack_id)
             finish_count -= 1
 
     def loading_check(self) -> None:
@@ -2626,7 +2679,7 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache):
                 )
 
         # ── PART 5: Ongoing Operations ──
-        for nid, (n, _) in self.ongoing_write_through.items():
+        for nid, (n, _, _) in self.ongoing_write_through.items():
             if n not in all_node_set:
                 E(f"[Ongoing] write_through node {nid} not in tree")
             elif n.component_data[FCT].lock_ref <= 0: