From 3c533acec6cdb0b0e44ab0f26411e499225fef02 Mon Sep 17 00:00:00 2001 From: Jincong Chen Date: Mon, 10 Aug 2026 17:16:42 +0800 Subject: [PATCH] [Hicache][2/2]Support Mamba branching in Unified Radix Cache with HiCache (#33639) --- .../sglang/srt/mem_cache/memory_pool_host.py | 20 +++---- .../components/mamba_component.py | 8 ++- .../unified_cache/components/swa_component.py | 3 ++ .../components/tree_component.py | 4 ++ .../unified_cache/unified_tree_core.py | 52 ++++++++++++++----- .../srt/mem_cache/unified_radix_cache.py | 7 +-- .../test_unified_radix_cache_unittest.py | 19 +++++++ 7 files changed, 88 insertions(+), 25 deletions(-) diff --git a/python/sglang/srt/mem_cache/memory_pool_host.py b/python/sglang/srt/mem_cache/memory_pool_host.py index 51aec93d4..4555ced64 100644 --- a/python/sglang/srt/mem_cache/memory_pool_host.py +++ b/python/sglang/srt/mem_cache/memory_pool_host.py @@ -1718,15 +1718,17 @@ class HostPoolGroup: pool_transfers: Optional[list] = None, ) -> None: # 1. Anchor (KV) backup - anchor_host_indices, anchor_device_indices = self._normalize_backup_indices( - self.anchor_entry, host_indices, device_indices, io_backend - ) - self.anchor_entry.host_pool.backup_from_device_all_layer( - self.anchor_entry.device_pool, - anchor_host_indices, - anchor_device_indices, - io_backend, - ) + # A zero-length anchor denotes a component-only backup. + if host_indices.numel() > 0: + anchor_host_indices, anchor_device_indices = self._normalize_backup_indices( + self.anchor_entry, host_indices, device_indices, io_backend + ) + self.anchor_entry.host_pool.backup_from_device_all_layer( + self.anchor_entry.device_pool, + anchor_host_indices, + anchor_device_indices, + io_backend, + ) # 2. Extra pool backup for transfer in pool_transfers or []: entry = self.entry_map.get(transfer.name) diff --git a/python/sglang/srt/mem_cache/unified_cache/components/mamba_component.py b/python/sglang/srt/mem_cache/unified_cache/components/mamba_component.py index 1f96b6b1f..11773c07f 100644 --- a/python/sglang/srt/mem_cache/unified_cache/components/mamba_component.py +++ b/python/sglang/srt/mem_cache/unified_cache/components/mamba_component.py @@ -73,6 +73,10 @@ class MambaComponent(TreeComponent): # HiCache state self._mamba_pool_host = None # set to host mamba pool when HiCache enabled + def needs_incremental_backup(self, node: UnifiedTreeNode) -> bool: + data = node.component_data[self.component_type] + return data.value is not None and data.host_value is None + def _inc_session_coverage(self, session_id: str, leaf: UnifiedTreeNode) -> None: cd = leaf.component_data[self.component_type] cd.session_ref += 1 @@ -156,7 +160,9 @@ class MambaComponent(TreeComponent): # Full KV may extend beyond the latest reusable Mamba state. The branching # point is the last Mamba-cache-chunk-aligned position within the Full-KV hit - # that lies beyond the current Mamba boundary. + # that lies beyond the current Mamba boundary. With HiCache, incremental + # persistence of a new branching state is currently write-through only; + # write-back eviction may discard the device-only state. aligned_seqlen = ( result.full_kv_hit_length // self.mamba_cache_chunk_size ) * self.mamba_cache_chunk_size diff --git a/python/sglang/srt/mem_cache/unified_cache/components/swa_component.py b/python/sglang/srt/mem_cache/unified_cache/components/swa_component.py index d35ef0731..1074c7cbf 100644 --- a/python/sglang/srt/mem_cache/unified_cache/components/swa_component.py +++ b/python/sglang/srt/mem_cache/unified_cache/components/swa_component.py @@ -76,6 +76,9 @@ class SWAComponent(TreeComponent): component_type = ComponentType.SWA + def needs_incremental_backup(self, node: UnifiedTreeNode) -> bool: + return False + def reset_session_state(self) -> None: super().reset_session_state() self._session_leaf_covered_len = {} diff --git a/python/sglang/srt/mem_cache/unified_cache/components/tree_component.py b/python/sglang/srt/mem_cache/unified_cache/components/tree_component.py index bc7b83a6a..44de7926d 100644 --- a/python/sglang/srt/mem_cache/unified_cache/components/tree_component.py +++ b/python/sglang/srt/mem_cache/unified_cache/components/tree_component.py @@ -315,6 +315,10 @@ class TreeComponent(ABC): cd = node.component_data[self.component_type] return cd.value is None and cd.host_value is not None + def needs_incremental_backup(self, node: UnifiedTreeNode) -> bool: + """Whether this component has new device data missing from Host.""" + return False + def refresh_lru( self, phase: LRURefreshPhase, diff --git a/python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py b/python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py index d1acd88b1..94805e2eb 100644 --- a/python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py +++ b/python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py @@ -1008,8 +1008,31 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface): ) state.phase = _InsertPhase.TAIL + def _needs_incremental_component_backup(self, node: UnifiedTreeNode) -> bool: + return any( + component.needs_incremental_backup(node) + for component in self.components + if component.component_type != BASE_COMPONENT_TYPE + ) + + def _should_backup_after_insert(self, state: _InsertWalkState) -> bool: + """Check whether the insert target needs a Host backup.""" + if state.is_new_leaf: + return self._inc_hit_count_and_check( + state.target_node, state.params.chunked + ) + + node = state.target_node + return ( + self.enable_hicache + and not self.is_write_back + and node.backuped + and node.write_through_pending_id is None + and self._needs_incremental_component_backup(node) + ) + def _insert_tail_step(self, state: _InsertWalkState) -> None: - """Refresh the LRUs and append the terminal new-leaf backup.""" + """Refresh the LRUs and append terminal backup actions.""" if state.target_node is not self.root_node: for component in self.components: if component.component_type == BASE_COMPONENT_TYPE: @@ -1018,9 +1041,7 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface): LRURefreshPhase.INSERT_END, state.target_node, self.root_node ) - if state.is_new_leaf and self._inc_hit_count_and_check( - state.target_node, state.params.chunked - ): + if self._should_backup_after_insert(state): state.pending_actions.append( self._build_backup_kv_action(state.target_node) ) @@ -1760,12 +1781,18 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface): return self._build_backup_spec(self.node_by_id(node_id)) def _build_backup_spec(self, node: UnifiedTreeNode): - """Gather device value backup spec.""" + """Gather missing Full and component transfers for Host backup.""" device_value = node.component_data[BASE_COMPONENT_TYPE].value + assert device_value is not None + if node.backuped: + device_value = device_value[:0] + comp_xfers: dict[ComponentType, list] = {} for comp in self.components: if comp.component_type == BASE_COMPONENT_TYPE: continue + if node.component_data[comp.component_type].host_value is not None: + continue t = comp.build_hicache_transfers(node, CacheTransferPhase.BACKUP_HOST) if t: comp_xfers[comp.component_type] = t @@ -1894,13 +1921,14 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface): """Commit a successful backup to the node.""" node = self.node_by_id(node_id) cache_actions: list[CacheAction | ComponentAction] = [] - kv_xfer = PoolTransfer(name=PoolName.KV, host_indices=host_indices) - self.components_by_type[BASE_COMPONENT_TYPE].commit_hicache_transfer( - node, - CacheTransferPhase.BACKUP_HOST, - transfers=[kv_xfer], - cache_actions=cache_actions, - ) + if host_indices.numel() > 0: + kv_xfer = PoolTransfer(name=PoolName.KV, host_indices=host_indices) + self.components_by_type[BASE_COMPONENT_TYPE].commit_hicache_transfer( + node, + CacheTransferPhase.BACKUP_HOST, + transfers=[kv_xfer], + cache_actions=cache_actions, + ) for ct, xfers in comp_xfers.items(): self.components_by_type[ct].commit_hicache_transfer( node, diff --git a/python/sglang/srt/mem_cache/unified_radix_cache.py b/python/sglang/srt/mem_cache/unified_radix_cache.py index cb7b9a499..f657897bc 100644 --- a/python/sglang/srt/mem_cache/unified_radix_cache.py +++ b/python/sglang/srt/mem_cache/unified_radix_cache.py @@ -919,10 +919,11 @@ class UnifiedRadixCache(BasePrefixCache): """Run a backup action top-down, stopping at the first failed backup.""" written = 0 for node_id in action.node_ids: - # Overlapping chain actions: skip already-backed nodes. - if self.tree_core.is_backuped(node_id): - continue device_value, comp_xfers = self.tree_core.build_backup_spec(node_id) + # Overlapping chain actions may revisit nodes with Full KV already + # backed up. Skip only when no transfer remains. + if device_value.numel() == 0 and not comp_xfers: + continue sidecar_xfers = self._build_backup_sidecar(device_value, comp_xfers) host_indices = self._execute_kv_backup( node_id, device_value, comp_xfers, sidecar_xfers 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 a8887b715..2e17223ed 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 @@ -3828,6 +3828,10 @@ class UnifiedRadixCacheSuite: leaf, device_frees, host_frees, target=EvictLayer.HOST ) cache._free_values(device_frees, host_frees) + full_host_pool = cache.cache_controller.mem_pool_host + mamba_host_pool = cache.components[ComponentType.MAMBA]._mamba_pool_host + full_available_before = full_host_pool.available_size() + mamba_available_before = mamba_host_pool.available_size() result = cache.match_prefix(MatchPrefixParams(key=RadixKey(array("q", tokens)))) @@ -3845,12 +3849,27 @@ class UnifiedRadixCacheSuite: req_to_token_pool, tokens[:branching_seqlen], ) + cache.writing_check(write_back=True) + # Full was already backed up, so only one Mamba slot is allocated. + self.assertEqual( + full_host_pool.available_size(), + full_available_before, + ) + self.assertEqual( + mamba_host_pool.available_size(), + mamba_available_before - 1, + ) + second_match = cache.match_prefix( MatchPrefixParams(key=RadixKey(array("q", tokens))) ) self.assertEqual(len(second_match.device_indices), branching_seqlen) self.assertIsNone(second_match.mamba_branching_seqlen) + branching_node = cache.resolve_node_handle(second_match.last_device_node) + self.assertIsNotNone( + branching_node.component_data[ComponentType.MAMBA].host_value + ) def test_scheduler_hicache_full_mamba_init_load_back_appends_new_indices(self): if not self.cfg.has_mamba or self.cfg.has_swa or self.cfg.page_size != 1: