From 978fb6ed1a1f113a46201439453dfc1313368109 Mon Sep 17 00:00:00 2001 From: ishandhanani <82981111+ishandhanani@users.noreply.github.com> Date: Wed, 3 Jun 2026 23:52:25 +0200 Subject: [PATCH] hicache kv events: publish split write-through fragments (#27072) --- python/sglang/srt/mem_cache/hiradix_cache.py | 63 +++++++-- python/sglang/srt/mem_cache/radix_cache.py | 1 + .../srt/mem_cache/unified_radix_cache.py | 85 +++++++++--- .../unit/mem_cache/test_hiradix_cache_unit.py | 123 ++++++++++++++++++ .../test_unified_radix_cache_unittest.py | 36 +++++ 5 files changed, 279 insertions(+), 29 deletions(-) create mode 100644 test/registered/unit/mem_cache/test_hiradix_cache_unit.py diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index aef02c2ca..5143d95fc 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/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: diff --git a/python/sglang/srt/mem_cache/radix_cache.py b/python/sglang/srt/mem_cache/radix_cache.py index 41c392fec..65c5b70e4 100644 --- a/python/sglang/srt/mem_cache/radix_cache.py +++ b/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 diff --git a/python/sglang/srt/mem_cache/unified_radix_cache.py b/python/sglang/srt/mem_cache/unified_radix_cache.py index c9ba6ec86..36406875f 100644 --- a/python/sglang/srt/mem_cache/unified_radix_cache.py +++ b/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: diff --git a/test/registered/unit/mem_cache/test_hiradix_cache_unit.py b/test/registered/unit/mem_cache/test_hiradix_cache_unit.py new file mode 100644 index 000000000..e7cd78329 --- /dev/null +++ b/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) diff --git a/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py b/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py index 4cb531965..f590ada12 100644 --- a/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py +++ b/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)