[Hicache][2/2]Support Mamba branching in Unified Radix Cache with HiCache (#33639)

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)

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

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 = {}

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,

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,

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

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: