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:

test/registered/unit/mem_cache/test_hiradix_cache_unit.py

@@ -0,0 +1,123 @@
+"""Unit tests for srt/mem_cache/hiradix_cache.py KV cache events."""
+
+import os
+import unittest
+from array import array
+
+import torch
+
+from sglang.srt.disaggregation.kv_events import BlockStored, StorageMedium
+from sglang.srt.mem_cache.allocator import TokenToKVPoolAllocator
+from sglang.srt.mem_cache.base_prefix_cache import InsertParams, MatchPrefixParams
+from sglang.srt.mem_cache.cache_init_params import CacheInitParams
+from sglang.srt.mem_cache.hiradix_cache import HiRadixCache
+from sglang.srt.mem_cache.memory_pool import MHATokenToKVPool, ReqToTokenPool
+from sglang.srt.mem_cache.radix_cache import RadixKey
+from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cuda_ci(est_time=15, stage="base-b", runner_config="1-gpu-small")
+
+PAGE_SIZE = 2
+
+
+class TestHiRadixCacheKVEvents(CustomTestCase):
+    @classmethod
+    def setUpClass(cls):
+        if not torch.cuda.is_available():
+            raise unittest.SkipTest("CUDA is required for HiRadixCache tests.")
+        os.environ.setdefault("MASTER_ADDR", "127.0.0.1")
+        os.environ.setdefault("MASTER_PORT", "29601")
+        if not torch.distributed.is_initialized():
+            torch.distributed.init_process_group(backend="gloo", rank=0, world_size=1)
+
+    def _build_cache(self):
+        server_args = ServerArgs(
+            model_path="dummy",
+            page_size=PAGE_SIZE,
+            hicache_io_backend="direct",
+            hicache_write_policy="write_through",
+        )
+        set_global_server_args_for_scheduler(server_args)
+        req_to_token_pool = ReqToTokenPool(
+            size=10,
+            max_context_len=512,
+            device="cuda",
+            enable_memory_saver=False,
+        )
+        kv_pool = MHATokenToKVPool(
+            size=256,
+            page_size=PAGE_SIZE,
+            dtype=torch.bfloat16,
+            head_num=2,
+            head_dim=64,
+            layer_num=4,
+            device="cuda",
+            enable_memory_saver=False,
+        )
+        allocator = TokenToKVPoolAllocator(
+            size=256,
+            dtype=torch.bfloat16,
+            device="cuda",
+            kvcache=kv_pool,
+            need_sort=False,
+        )
+        params = CacheInitParams(
+            req_to_token_pool=req_to_token_pool,
+            token_to_kv_pool_allocator=allocator,
+            page_size=PAGE_SIZE,
+            disable=False,
+            enable_kv_cache_events=True,
+            tp_cache_group=torch.distributed.group.WORLD,
+        )
+        cache = HiRadixCache(params, server_args)
+        # Disable hit-count-driven write-through; tests back up explicitly.
+        cache.write_through_threshold = 1 << 30
+        return cache, allocator
+
+    def _insert(self, cache, allocator, tokens):
+        key = RadixKey(array("q", tokens))
+        value = allocator.alloc(len(tokens))
+        self.assertIsNotNone(value)
+        return cache.insert(InsertParams(key=key, value=value[: len(tokens)]))
+
+    def _leaf_for(self, cache, tokens):
+        match = cache.match_prefix(MatchPrefixParams(key=RadixKey(array("q", tokens))))
+        self.assertIsNot(match.last_device_node, cache.root_node)
+        return match.last_device_node
+
+    def _stored_cpu_events(self, cache):
+        return [
+            e
+            for e in cache.take_events()
+            if isinstance(e, BlockStored) and e.medium == StorageMedium.CPU
+        ]
+
+    def test_split_pending_write_through_publishes_fragments(self):
+        cache, allocator = self._build_cache()
+        cache.take_events()
+
+        self._insert(cache, allocator, [1, 2, 3, 4])
+        node = self._leaf_for(cache, [1, 2, 3, 4])
+        backed_up = cache.write_backup(node, write_back=True)
+        self.assertGreater(backed_up, 0)
+
+        # Split the node while its write-through DMA is still pending.
+        self._insert(cache, allocator, [1, 2, 5, 6])
+        self.assertEqual(self._stored_cpu_events(cache), [])
+
+        cache.writing_check(write_back=True)
+
+        # Both split fragments must be published, with intact parentage.
+        stored_cpu = self._stored_cpu_events(cache)
+        self.assertEqual(
+            [list(e.token_ids) for e in stored_cpu],
+            [[1, 2], [3, 4]],
+        )
+        self.assertIsNone(stored_cpu[0].parent_block_hash)
+        self.assertEqual(stored_cpu[1].parent_block_hash, stored_cpu[0].block_hashes[0])
+
+
+if __name__ == "__main__":
+    unittest.main(verbosity=2)

test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py

@@ -462,6 +462,42 @@ class TestUnifiedRadixCacheKVEvents(CustomTestCase):
         removed_cpu = self._removed_events(tree, StorageMedium.CPU)
         self.assertCountEqual([e.block_hashes[0] for e in removed_cpu], stored_hashes)
 
+    def test_hicache_split_pending_write_through_publishes_fragments(self):
+        tree, allocator, _ = build_fixture(self.cfg, enable_kv_cache_events=True)
+        self._init_hicache(tree)
+        tree.take_events()
+
+        self._insert(tree, allocator, [1, 2, 3, 4])
+        node = self._leaf_for(tree, [1, 2, 3, 4])
+        backed_up = tree.write_backup(node, write_back=True)
+        self.assertGreater(backed_up, 0)
+
+        # Split the node while its write-through DMA is still pending.
+        self._insert(tree, allocator, [1, 2, 5, 6])
+        self.assertEqual(self._stored_events(tree, StorageMedium.CPU), [])
+
+        # Each fragment must also be persisted to L3 on ack: lock_node only
+        # holds the suffix after the split.
+        tree.enable_storage = True
+        with mock.patch.object(tree, "write_backup_storage") as backup_storage:
+            tree.writing_check(write_back=True)
+        self.assertEqual(
+            [
+                list(call.args[0].key.token_ids)
+                for call in backup_storage.call_args_list
+            ],
+            [[1, 2], [3, 4]],
+        )
+
+        # Both split fragments must be published, with intact parentage.
+        stored_cpu = self._stored_events(tree, StorageMedium.CPU)
+        self.assertEqual(
+            [list(e.token_ids) for e in stored_cpu],
+            [[1, 2], [3, 4]],
+        )
+        self.assertIsNone(stored_cpu[0].parent_block_hash)
+        self.assertEqual(stored_cpu[1].parent_block_hash, stored_cpu[0].block_hashes[0])
+
     def test_hicache_reinsert_evicted_node_emits_gpu_store(self):
         tree, allocator, _ = build_fixture(self.cfg, enable_kv_cache_events=True)
         self._init_hicache(tree)