[UnifiedTree]: Fix UnifiedRadixCache device match semantics with HiCache (#25277)

python/sglang/srt/mem_cache/unified_cache_components/full_component.py

@@ -41,8 +41,15 @@ class FullComponent(TreeComponent):
         # HiCache state: set to host KV pool when HiCache enabled
         self._full_kv_pool_host = None
 
-    def create_match_validator(self) -> Callable[[UnifiedTreeNode], bool]:
+    def create_match_validator(
-        # HiCache: evicted + backuped nodes are valid match boundaries
+        self, match_device_only: bool = False
+    ) -> Callable[[UnifiedTreeNode], bool]:
+        if match_device_only:
+            return (
+                lambda node: node.component_data[self.component_type].value is not None
+            )
+
+        # HiCache: evicted + backuped nodes are valid match boundaries.
         return lambda node: (
             node.component_data[self.component_type].value is not None or node.backuped
         )

python/sglang/srt/mem_cache/unified_cache_components/mamba_component.py

@@ -50,8 +50,13 @@ class MambaComponent(TreeComponent):
         # HiCache state
         self._mamba_pool_host = None  # set to host mamba pool when HiCache enabled
 
-    def create_match_validator(self) -> Callable[[UnifiedTreeNode], bool]:
+    def create_match_validator(
+        self, match_device_only: bool = False
+    ) -> Callable[[UnifiedTreeNode], bool]:
         ct = self.component_type
+        if match_device_only:
+            return lambda node: node.component_data[ct].value is not None
+
         # HiCache: evicted + backuped (host_value present) is also a valid match
         return lambda node: (
             node.component_data[ct].value is not None
@@ -69,7 +74,10 @@ class MambaComponent(TreeComponent):
         req = params.req
         last_node = result.best_match_node
 
-        if len(value_chunks) > best_value_len:
+        # HiCache can still use prefix matches and load back host-backed Mamba
+        # states. We temporarily skip branching-state fill in that mode and can
+        # add a HiCache-aware branching policy later.
+        if self.cache.cache_controller is None and len(value_chunks) > best_value_len:
             chunk_size = get_global_server_args().mamba_cache_chunk_size
             aligned_seqlen = (
                 sum(len(v) for v in value_chunks) // chunk_size

python/sglang/srt/mem_cache/unified_cache_components/swa_component.py

@@ -69,7 +69,9 @@ class SWAComponent(TreeComponent):
         self.cache.lru_lists[ct].insert_mru(node)
         self.cache.component_evictable_size_[ct] += len(value)
 
-    def create_match_validator(self) -> Callable[[UnifiedTreeNode], bool]:
+    def create_match_validator(
+        self, match_device_only: bool = False
+    ) -> Callable[[UnifiedTreeNode], bool]:
         sliding_window_size = self.sliding_window_size
         ct = self.component_type
         state = {"len": float("inf")}
@@ -78,7 +80,7 @@ class SWAComponent(TreeComponent):
             cd = node.component_data[ct]
             # HiCache: a host-only tombstone is a valid match boundary too
             # — load_back will restore SWA from host before use.
-            if cd.value is None and cd.host_value is None:
+            if cd.value is None and (match_device_only or cd.host_value is None):
                 state["len"] = 0
                 return False
             state["len"] += len(node.key)

python/sglang/srt/mem_cache/unified_cache_components/tree_component.py

@@ -116,11 +116,15 @@ class TreeComponent(ABC):
         return len(value) if value is not None else 0
 
     @abstractmethod
-    def create_match_validator(self) -> Callable[[UnifiedTreeNode], bool]:
+    def create_match_validator(
+        self, match_device_only: bool = False
+    ) -> Callable[[UnifiedTreeNode], bool]:
         """Return a per-match stateful predicate that decides whether a node
         is a valid match boundary for this component.
         Called once per match_prefix; the returned closure may carry state.
-        - Full: always True (every node is valid).
+        When match_device_only is true, host-backed nodes must not be accepted
+        as valid match boundaries.
+        - Full: returns True if the node has full component data.
         - SWA: tracks accumulated length since last gap; returns True only
           when the contiguous window reaches swa_sliding_window_size.
         - Mamba: returns True iff the node has mamba component data."""

python/sglang/srt/mem_cache/unified_radix_cache.py

@@ -351,9 +351,18 @@ class UnifiedRadixCache(BasePrefixCache):
         if len(key) == 0:
             return self._empty_match_result
 
-        value, best_match_node, best_value_len = self._match_prefix_helper(key)
+        (
+            value,
+            best_match_node,
+            best_match_device_node,
+            best_match_device_value_len,
+        ) = self._match_prefix_helper(key)
         return self._match_post_processor(
-            params, value, best_match_node, best_value_len
+            params,
+            value,
+            best_match_node,
+            best_match_device_node,
+            best_match_device_value_len,
         )
 
     def insert(self, params: InsertParams) -> InsertResult:
@@ -585,67 +594,53 @@ class UnifiedRadixCache(BasePrefixCache):
 
     # ---- Internal Helpers ----
 
-    def _match_prefix_helper_readonly(
-        self, key: RadixKey
-    ) -> tuple[list[torch.Tensor], UnifiedTreeNode, int]:
-        """Read-only version of _match_prefix_helper that does not split nodes.
-        Only considers fully matched nodes, ignores partial matches.
-
-        Not used yet; reserved for future read-only match operations."""
-        node = self.root_node
-        child_key = key.child_key(self.page_size)
-        value: list[torch.Tensor] = []
-        best_value_len = 0
-        best_match_node = node
-        validators = tuple(
-            comp.create_match_validator() for comp in self._components_tuple
-        )
-
-        def _update_best_if_valid(node):
-            nonlocal best_value_len, best_match_node
-            if all(v(node) for v in validators):
-                best_value_len = len(value)
-                best_match_node = node
-
-        while len(key) > 0 and child_key in node.children:
-            child = node.children[child_key]
-
-            # HiCache: dead node (evicted + not backuped) — stop traversal
-            if child.evicted and not child.backuped:
-                break
-
-            prefix_len = child.key.match(key, page_size=self.page_size)
-            if prefix_len < len(child.key):
-                # Read-only: do not split, ignore partial match and stop
-                break
-
-            if not child.evicted:
-                value.append(child.component_data[BASE_COMPONENT_TYPE].value)
-            node = child
-            _update_best_if_valid(node)
-            key = key[prefix_len:]
-            if len(key):
-                child_key = key.child_key(self.page_size)
-        return value, best_match_node, best_value_len
-
     def _match_prefix_helper(
         self, key: RadixKey
-    ) -> tuple[list[torch.Tensor], UnifiedTreeNode, int]:
+    ) -> tuple[list[torch.Tensor], UnifiedTreeNode, UnifiedTreeNode, int]:
+        # Non-HiCache mode has only device-resident matches, so the scheduler
+        # device anchor follows the best match. In HiCache mode, host-backed
+        # nodes can also match, so we separately track the best device-resident
+        # match for scheduler prefix indices and locking.
         node = self.root_node
         child_key = key.child_key(self.page_size)
         value: list[torch.Tensor] = []
-        best_value_len = 0
         best_match_node = node
+        best_match_device_node = node
+        best_match_device_value_len = 0
+        separate_device_match = self.cache_controller is not None
+        if separate_device_match:
             validators = tuple(
                 comp.create_match_validator() for comp in self._components_tuple
             )
+            device_validators = tuple(
+                comp.create_match_validator(match_device_only=True)
+                for comp in self._components_tuple
+            )
+        else:
+            validators = tuple(
+                comp.create_match_validator(match_device_only=True)
+                for comp in self._components_tuple
+            )
+
+        def _all_valid(validators, node):
+            return all([v(node) for v in validators])
 
         def _update_best_if_valid(node):
-            nonlocal best_value_len, best_match_node
+            nonlocal best_match_node
-            if all(v(node) for v in validators):
+            nonlocal best_match_device_value_len, best_match_device_node
-                best_value_len = len(value)
+            matched = _all_valid(validators, node)
+            if matched:
                 best_match_node = node
 
+            if not separate_device_match:
+                if matched:
+                    best_match_device_value_len = len(value)
+                    best_match_device_node = node
+                return
+            if _all_valid(device_validators, node):
+                best_match_device_value_len = len(value)
+                best_match_device_node = node
+
         while len(key) > 0 and child_key in node.children:
             child = node.children[child_key]
 
@@ -668,14 +663,21 @@ class UnifiedRadixCache(BasePrefixCache):
             key = key[prefix_len:]
             if len(key):
                 child_key = key.child_key(self.page_size)
-        return value, best_match_node, best_value_len
+
+        return (
+            value,
+            best_match_node,
+            best_match_device_node,
+            best_match_device_value_len,
+        )
 
     def _match_post_processor(
         self,
         params: MatchPrefixParams,
         value: list[torch.Tensor],
         best_match_node: UnifiedTreeNode,
-        best_value_len: int,
+        best_match_device_node: UnifiedTreeNode,
+        best_match_device_value_len: int,
     ) -> MatchResult:
         node_update = best_match_node
         for comp in self._components_tuple:
@@ -691,23 +693,21 @@ class UnifiedRadixCache(BasePrefixCache):
             cur_time -= 0.00001
             node_update = node_update.parent
 
-        # Walk up to find last_device_node
+        # Walk up to find last_host_node for full component.
-        last_device_node = best_match_node
+        if self.cache_controller is None:
-        while last_device_node is not self.root_node and last_device_node.evicted:
+            last_host_node = best_match_device_node
-            last_device_node = last_device_node.parent
+        else:
-
-        # Walk up to find last_host_node
             last_host_node = best_match_node
             while last_host_node is not self.root_node and not last_host_node.backuped:
                 last_host_node = last_host_node.parent
 
-        if best_value_len > 0:
+        if best_match_device_value_len > 0:
-            device_indices = torch.cat(value[:best_value_len])
+            device_indices = torch.cat(value[:best_match_device_value_len])
         else:
             device_indices = self._empty_match_result.device_indices
         result = MatchResult(
             device_indices=device_indices,
-            last_device_node=last_device_node,
+            last_device_node=best_match_device_node,
             last_host_node=last_host_node,
             best_match_node=best_match_node,
             host_hit_length=0,
@@ -718,7 +718,7 @@ class UnifiedRadixCache(BasePrefixCache):
                 result=result,
                 params=params,
                 value_chunks=value,
-                best_value_len=best_value_len,
+                best_value_len=best_match_device_value_len,
             )
         return result
 
@@ -1219,10 +1219,10 @@ class UnifiedRadixCache(BasePrefixCache):
         best_match_node: UnifiedTreeNode,
         mem_quota: Optional[int] = None,
         req=None,
-    ) -> Optional[torch.Tensor]:
+    ) -> bool:
         """Load evicted KV data from host back to device (H→D)."""
         if self.cache_controller is None:
-            return None
+            return False
 
         # Build KV transfer
         kv_xfer = self.components[BASE_COMPONENT_TYPE].build_hicache_transfers(
@@ -1255,7 +1255,7 @@ class UnifiedRadixCache(BasePrefixCache):
             mem_quota is not None and kv_tokens > mem_quota + result.delta
         ):
             self.dec_lock_ref(best_match_node, ancestor_lock_params)
-            return None
+            return False
 
         avail = self.token_to_kv_pool_allocator.available_size()
         if avail < kv_tokens:
@@ -1263,7 +1263,7 @@ class UnifiedRadixCache(BasePrefixCache):
             result = self.evict(EvictParams(num_tokens=needed))
             if result.num_tokens_evicted < needed:
                 self.dec_lock_ref(best_match_node, ancestor_lock_params)
-                return None
+                return False
 
         # Load H→D
         aux_xfers = [x for xfers in comp_xfers.values() for x in xfers]
@@ -1276,7 +1276,7 @@ class UnifiedRadixCache(BasePrefixCache):
 
         self.dec_lock_ref(best_match_node, ancestor_lock_params)
         if device_indices is None:
-            return None
+            return False
 
         # Commit: each component gets only its own transfers
         kv_xfer.device_indices = device_indices
@@ -1297,7 +1297,7 @@ class UnifiedRadixCache(BasePrefixCache):
             best_match_node,
             self.inc_lock_ref(best_match_node).to_dec_params(),
         )
-        return device_indices
+        return True
 
     def _build_sidecar_transfers(
         self,
@@ -1432,25 +1432,41 @@ class UnifiedRadixCache(BasePrefixCache):
         best_match_node = params.best_match_node
         mem_quota = params.mem_quota
         req = params.req
+        assert req is not None
+        last_best_match_device_node = req.last_node
+
+        def _collect_new_prefix_indices() -> torch.Tensor:
+            prefix_chunks: list[torch.Tensor] = []
+            node = best_match_node
+            while node is not last_best_match_device_node:
+                value = node.component_data[BASE_COMPONENT_TYPE].value
+                assert value is not None
+                prefix_chunks.append(value)
+                node = node.parent
+            if not prefix_chunks:
+                return self._empty_match_result.device_indices
+            prefix_chunks.reverse()
+            return torch.cat(prefix_chunks)
 
         if best_match_node.evicted or params.host_hit_length > 0:
-            loading_values = self.load_back(best_match_node, mem_quota, req=req)
+            if self.load_back(best_match_node, mem_quota, req=req):
-            if loading_values is not None:
+                new_indices = _collect_new_prefix_indices()
+                if new_indices.numel() == 0:
+                    return (
+                        self._empty_match_result.device_indices,
+                        last_best_match_device_node,
+                    )
+
                 logger.debug(
                     "init_load_back success: loaded %d tokens for node %d",
-                    len(loading_values),
+                    len(new_indices),
                     best_match_node.id,
                 )
-                return loading_values, best_match_node
+                return new_indices, best_match_node
-
-            # Fallback: walk up to non-evicted ancestor
-            # TODO(ispobock): The fallback path is not correct. The last_device_node should consider all the components.
-            while best_match_node is not self.root_node and best_match_node.evicted:
-                best_match_node = best_match_node.parent
 
         return (
             self._empty_match_result.device_indices,
-            best_match_node,
+            last_best_match_device_node,
         )
 
     def check_hicache_events(self) -> None:

test/registered/radix_cache/test_unified_radix_hicache_kl.py

@@ -155,7 +155,7 @@ class TestUnifiedDeepSeekV4FlashHiCache(UnifiedRadixTreeTestMixin, CustomTestCas
                 "--max-total-tokens",
                 "20000",
                 "--max-running-requests",
-                "4",
+                "2",
             ],
             env={
                 "SGLANG_DSV4_FP4_EXPERTS": "0",

test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py

@@ -15,6 +15,7 @@ from sglang.srt.mem_cache.base_prefix_cache import (
     DecLockRefParams,
     EvictParams,
     EvictResult,
+    InitLoadBackParams,
     InsertParams,
     MatchPrefixParams,
     MatchResult,
@@ -908,6 +909,8 @@ class UnifiedRadixCacheSuite:
         """Verify readonly match does not modify tree structure (no split)."""
         if self.cfg.page_size > 1 or self.cfg.has_mamba or self.cfg.has_swa:
             self.skipTest("Full-only page_size=1 only")
+        if not hasattr(UnifiedRadixCache, "_match_prefix_helper_readonly"):
+            self.skipTest("_match_prefix_helper_readonly is not available")
         tree, allocator, req_to_token_pool = build_fixture(self.cfg)
 
         self._insert(tree, allocator, req_to_token_pool, [1, 2, 3, 4, 5])
@@ -922,19 +925,27 @@ class UnifiedRadixCacheSuite:
         self.assertEqual(node_count_before, 2)
 
         tree._match_prefix_helper(RadixKey([1, 2]))
-        value, best_match_node, best_value_len = tree._match_prefix_helper(
+        (
-            RadixKey([1, 2, 3, 4])
+            value,
-        )
+            best_match_node,
+            best_match_device_node,
+            best_value_len,
+        ) = tree._match_prefix_helper(RadixKey([1, 2, 3, 4]))
         self.assertEqual(best_value_len, 2)
         self.assertEqual(best_match_node.key.token_ids, [3, 4])
+        self.assertIs(best_match_device_node, best_match_node)
         node_count_after_regular = count_nodes(tree.root_node)
         self.assertEqual(node_count_after_regular, node_count_before + 2)
 
-        value, best_match_node, best_value_len = tree._match_prefix_helper_readonly(
+        (
-            RadixKey([1, 2, 3])
+            value,
-        )
+            best_match_node,
+            best_match_device_node,
+            best_value_len,
+        ) = tree._match_prefix_helper_readonly(RadixKey([1, 2, 3]))
         self.assertEqual(best_value_len, 1)
         self.assertEqual(best_match_node.key.token_ids, [1, 2])
+        self.assertIs(best_match_device_node, best_match_node)
         node_count_after_readonly = count_nodes(tree.root_node)
         self.assertEqual(node_count_after_readonly, node_count_after_regular)
 
@@ -1258,8 +1269,8 @@ class UnifiedRadixCacheSuite:
     # ================================================================
 
     def _skip_unsupported_hicache_test(self):
-        if self.cfg.has_swa:
+        if self.cfg.has_swa and self.cfg.has_mamba:
-            self.skipTest("HiCache tests do not run on SWA stacks")
+            self.skipTest("HiCache unit fixture does not support SWA + Mamba stacks")
         return False
 
     def _simulate_backup(self, tree, node):
@@ -1343,14 +1354,14 @@ class UnifiedRadixCacheSuite:
                 self._backup_node(tree, node)
 
     def _load_back_node(self, tree, node):
-        device_indices = tree.load_back(node)
+        loaded = tree.load_back(node)
-        self.assertIsNotNone(device_indices)
+        self.assertTrue(loaded)
         producer_id = tree.ready_to_load_host_cache()
         self.assertNotEqual(producer_id, -1)
         for _, finish_event, _ in list(tree.cache_controller.ack_load_queue):
             finish_event.synchronize()
         tree.loading_check()
-        return device_indices
+        return node.component_data[ComponentType.FULL].value
 
     def _get_full_kv_pool(self, allocator):
         kv_pool = allocator.get_kvcache()
@@ -1474,7 +1485,9 @@ class UnifiedRadixCacheSuite:
 
     def test_hicache_partial_match_splits_evicted_backed_up_node(self):
         """Partial matches on host-only nodes must keep the host prefix usable."""
-        tree, allocator, req_to_token_pool = build_fixture(self.cfg)
+        if self._skip_unsupported_hicache_test():
+            return
+        tree, allocator, req_to_token_pool = self._build_hicache_fixture()
         ps = self.cfg.page_size
         seq = self._make_seq(1, 4)
         expected_prefix = seq[: 2 * ps]
@@ -1484,7 +1497,7 @@ class UnifiedRadixCacheSuite:
         self._insert(tree, allocator, req_to_token_pool, seq)
         m = tree.match_prefix(MatchPrefixParams(key=RadixKey(seq)))
         node = m.last_device_node
-        self._simulate_backup(tree, node)
+        self._backup_node(tree, node)
 
         tree.evict(EvictParams(num_tokens=len(seq)))
         self.assertTrue(node.evicted)
@@ -1684,6 +1697,244 @@ class UnifiedRadixCacheSuite:
         chain.reverse()
         return chain
 
+    def _release_ongoing_load_back_locks(self, tree):
+        for node, lock_params in list(tree.ongoing_load_back.values()):
+            tree.dec_lock_ref(node, lock_params)
+        tree.ongoing_load_back.clear()
+
+    def _finish_pending_loads(self, tree):
+        producer_id = tree.ready_to_load_host_cache()
+        self.assertNotEqual(producer_id, -1)
+        for _, finish_event, _ in list(tree.cache_controller.ack_load_queue):
+            finish_event.synchronize()
+        tree.loading_check()
+
+    def _match_tokens_for_chain(self, chain):
+        tokens: list[int] = []
+        for node in chain:
+            tokens.extend(node.key.token_ids)
+        return tokens
+
+    def _set_aux_host_tombstone(self, tree, node, component_type):
+        cd = node.component_data[component_type]
+        self.assertIsNotNone(cd.value)
+        if cd.host_value is None:
+            cd.host_value = cd.value.clone()
+        old_value = cd.value
+        cd.value = None
+        if component_type in tree.lru_lists and tree.lru_lists[component_type].in_list(
+            node
+        ):
+            tree.lru_lists[component_type].remove_node(node)
+        tree.host_lru_lists[component_type].insert_mru(node)
+        tree.component_evictable_size_[component_type] -= len(old_value)
+
+    def test_match_prefix_best_and_device_node_without_hicache(self):
+        tree, allocator, req_to_token_pool = build_fixture(self.cfg)
+        ps = self.cfg.page_size
+        min_tokens = 2 * ps
+        if self.cfg.has_swa:
+            min_tokens = max(min_tokens, self.cfg.sliding_window_size + ps)
+        seq = self._make_seq(1, (min_tokens + ps - 1) // ps)
+        self._insert(tree, allocator, req_to_token_pool, seq)
+
+        result = tree.match_prefix(MatchPrefixParams(key=RadixKey(seq)))
+
+        self.assertEqual(len(result.device_indices), len(seq))
+        self.assertIs(result.best_match_node, result.last_device_node)
+        self.assertIs(result.last_host_node, result.last_device_node)
+        self.assertEqual(result.host_hit_length, 0)
+
+    def test_hicache_mamba_host_best_match_keeps_device_anchor(self):
+        if not self.cfg.has_mamba or self.cfg.has_swa or self.cfg.page_size != 1:
+            self.skipTest("requires page_size=1 Full+Mamba")
+        tree, allocator, req_to_token_pool = self._build_hicache_fixture()
+        chain = self._build_chain_pages(tree, allocator, req_to_token_pool, 3)
+        if len(chain) < 3:
+            self.skipTest("chain too short")
+        leaf = chain[-1]
+        parent = chain[-2]
+        tokens = self._match_tokens_for_chain(chain)
+
+        self._backup_node(tree, leaf)
+        tree.evict(EvictParams(num_tokens=len(leaf.key)))
+        self.assertTrue(leaf.evicted)
+
+        result = tree.match_prefix(MatchPrefixParams(key=RadixKey(tokens)))
+
+        self.assertIs(result.best_match_node, leaf)
+        self.assertIs(result.last_device_node, parent)
+        self.assertEqual(len(result.device_indices), len(tokens) - len(leaf.key))
+        self.assertEqual(result.host_hit_length, len(leaf.key))
+
+    def test_hicache_swa_host_best_match_keeps_device_anchor(self):
+        if not self.cfg.has_swa or self.cfg.has_mamba or self.cfg.page_size != 1:
+            self.skipTest("requires page_size=1 Full+SWA")
+        tree, allocator, req_to_token_pool = self._build_hicache_fixture()
+        chain = self._build_chain_pages(tree, allocator, req_to_token_pool, 3)
+        if len(chain) < 3:
+            self.skipTest("chain too short")
+        leaf = chain[-1]
+        parent = chain[-2]
+        tokens = self._match_tokens_for_chain(chain)
+
+        self._backup_node(tree, leaf)
+        tree.evict(EvictParams(num_tokens=len(leaf.key)))
+        self.assertTrue(leaf.evicted)
+
+        result = tree.match_prefix(MatchPrefixParams(key=RadixKey(tokens)))
+
+        self.assertIs(result.best_match_node, leaf)
+        self.assertIs(result.last_device_node, parent)
+        self.assertEqual(len(result.device_indices), len(tokens) - len(leaf.key))
+        self.assertEqual(result.host_hit_length, 1)
+
+    def test_mamba_branching_seqlen_disabled_under_hicache(self):
+        if not self.cfg.has_mamba or self.cfg.has_swa or self.cfg.page_size != 1:
+            self.skipTest("requires page_size=1 Full+Mamba")
+        tree, allocator, req_to_token_pool = build_fixture(self.cfg)
+        chunk_size = get_global_server_args().mamba_cache_chunk_size
+        tokens = self._make_seq(1, chunk_size + 1)
+        self._insert(tree, allocator, req_to_token_pool, tokens)
+        leaf = tree.match_prefix(
+            MatchPrefixParams(key=RadixKey(tokens))
+        ).last_device_node
+
+        mamba_cd = leaf.component_data[ComponentType.MAMBA]
+        mamba_cd.value = None
+        no_hicache = tree.match_prefix(MatchPrefixParams(key=RadixKey(tokens)))
+        self.assertIs(no_hicache.best_match_node, tree.root_node)
+        self.assertIs(no_hicache.last_device_node, tree.root_node)
+        self.assertEqual(no_hicache.mamba_branching_seqlen, chunk_size)
+
+        tree_h, allocator_h, req_to_token_pool_h = self._build_hicache_fixture()
+        self._insert(tree_h, allocator_h, req_to_token_pool_h, tokens)
+        leaf_h = tree_h.match_prefix(
+            MatchPrefixParams(key=RadixKey(tokens))
+        ).last_device_node
+        self._backup_node(tree_h, leaf_h)
+        tree_h.evict(EvictParams(num_tokens=len(tokens)))
+        with_hicache = tree_h.match_prefix(MatchPrefixParams(key=RadixKey(tokens)))
+        self.assertIs(with_hicache.best_match_node, leaf_h)
+        self.assertIs(with_hicache.last_device_node, tree_h.root_node)
+        self.assertIsNone(with_hicache.mamba_branching_seqlen)
+
+    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:
+            self.skipTest("requires page_size=1 Full+Mamba")
+        tree, allocator, req_to_token_pool = self._build_hicache_fixture()
+        chain = self._build_chain_pages(tree, allocator, req_to_token_pool, 3)
+        if len(chain) < 3:
+            self.skipTest("chain too short")
+        leaf = chain[-1]
+        tokens = self._match_tokens_for_chain(chain)
+
+        self._backup_node(tree, leaf)
+        tree.evict(EvictParams(num_tokens=len(leaf.key)))
+        self.assertTrue(leaf.evicted)
+
+        req = self._make_req(req_to_token_pool)
+        match = tree.match_prefix(MatchPrefixParams(key=RadixKey(tokens), req=req))
+        req.prefix_indices = match.device_indices
+        req.last_node = match.last_device_node
+        req.best_match_node = match.best_match_node
+        req.host_hit_length = match.host_hit_length
+
+        new_indices, new_node = tree.init_load_back(
+            InitLoadBackParams(
+                best_match_node=req.best_match_node,
+                host_hit_length=req.host_hit_length,
+                req=req,
+            )
+        )
+
+        self.assertIs(new_node, leaf)
+        self.assertEqual(len(torch.cat([req.prefix_indices, new_indices])), len(tokens))
+        self.assertIsNotNone(leaf.component_data[ComponentType.MAMBA].value)
+        self._finish_pending_loads(tree)
+        self._release_ongoing_load_back_locks(tree)
+
+    def test_scheduler_hicache_aux_only_load_back_appends_full_device_indices(self):
+        if self.cfg.page_size != 1:
+            self.skipTest("page_size=1 keeps the expected suffix precise")
+        aux = None
+        if self.cfg.has_swa and not self.cfg.has_mamba:
+            aux = ComponentType.SWA
+        elif self.cfg.has_mamba and not self.cfg.has_swa:
+            aux = ComponentType.MAMBA
+        if aux is None:
+            self.skipTest("requires exactly one aux component")
+
+        tree, allocator, req_to_token_pool = self._build_hicache_fixture()
+        chain = self._build_chain_pages(tree, allocator, req_to_token_pool, 3)
+        if len(chain) < 3:
+            self.skipTest("chain too short")
+        leaf = chain[-1]
+        tokens = self._match_tokens_for_chain(chain)
+        leaf_full = leaf.component_data[ComponentType.FULL].value.clone()
+        self._backup_node(tree, leaf)
+        self._set_aux_host_tombstone(tree, leaf, aux)
+
+        req = self._make_req(req_to_token_pool)
+        match = tree.match_prefix(MatchPrefixParams(key=RadixKey(tokens), req=req))
+        req.prefix_indices = match.device_indices
+        req.last_node = match.last_device_node
+        req.best_match_node = match.best_match_node
+        req.host_hit_length = match.host_hit_length
+
+        new_indices, new_node = tree.init_load_back(
+            InitLoadBackParams(
+                best_match_node=req.best_match_node,
+                host_hit_length=req.host_hit_length,
+                req=req,
+            )
+        )
+
+        self.assertIs(new_node, leaf)
+        self.assertEqual(new_indices.tolist(), leaf_full.tolist())
+        self.assertEqual(len(torch.cat([req.prefix_indices, new_indices])), len(tokens))
+        self.assertEqual(
+            leaf.component_data[ComponentType.FULL].value.tolist(),
+            leaf_full.tolist(),
+        )
+        self.assertIsNotNone(leaf.component_data[aux].value)
+        self._finish_pending_loads(tree)
+        self._release_ongoing_load_back_locks(tree)
+
+    def test_scheduler_hicache_load_back_fallback_keeps_old_anchor(self):
+        if not self.cfg.has_mamba or self.cfg.has_swa or self.cfg.page_size != 1:
+            self.skipTest("requires page_size=1 Full+Mamba")
+        tree, allocator, req_to_token_pool = self._build_hicache_fixture()
+        chain = self._build_chain_pages(tree, allocator, req_to_token_pool, 3)
+        if len(chain) < 3:
+            self.skipTest("chain too short")
+        leaf = chain[-1]
+        tokens = self._match_tokens_for_chain(chain)
+
+        self._backup_node(tree, leaf)
+        tree.evict(EvictParams(num_tokens=len(leaf.key)))
+
+        req = self._make_req(req_to_token_pool)
+        match = tree.match_prefix(MatchPrefixParams(key=RadixKey(tokens), req=req))
+        req.prefix_indices = match.device_indices
+        req.last_node = match.last_device_node
+        req.best_match_node = match.best_match_node
+        req.host_hit_length = match.host_hit_length
+
+        new_indices, new_node = tree.init_load_back(
+            InitLoadBackParams(
+                best_match_node=req.best_match_node,
+                host_hit_length=req.host_hit_length,
+                req=req,
+                mem_quota=-1_000_000,
+            )
+        )
+
+        self.assertEqual(len(new_indices), 0)
+        self.assertIs(new_node, match.last_device_node)
+        self.assertIsNone(leaf.component_data[ComponentType.FULL].value)
+        self.assertIsNone(leaf.component_data[ComponentType.MAMBA].value)
+
     def test_hicache_swa_load_back_min_suffix(self):
         """LOAD_BACK collects only the suffix nodes needed to cover sliding_window_size."""
         if not self.cfg.has_swa:
@@ -1940,11 +2191,11 @@ class UnifiedRadixCacheSuite:
         if chain_pages * ps > self.cfg.kv_size // 2:
             self.skipTest("kv_size too small for the desired chain")
 
-        tree, allocator, req_to_token_pool = build_fixture(self.cfg)
+        tree, allocator, req_to_token_pool = self._build_hicache_fixture()
         chain = self._build_chain_pages(tree, allocator, req_to_token_pool, chain_pages)
         if len(chain) < chain_pages:
             self.skipTest("chain too short")
-        self._simulate_backup_tree(tree)
+        self._backup_tree(tree)
 
         x = chain[-1]
         y = chain[-window_pages]
@@ -1962,10 +2213,10 @@ class UnifiedRadixCacheSuite:
         return tree, chain, n, y, x, tokens
 
     def test_hicache_swa_match_prefix_picks_best_match_node_above_last_host(self):
-        tree, _, _, y, x, tokens = self._swa_anchor_setup()
+        tree, _, n, y, x, tokens = self._swa_anchor_setup()
         result = tree.match_prefix(MatchPrefixParams(key=RadixKey(tokens)))
         self.assertIs(result.best_match_node, x)
-        self.assertIs(result.last_device_node, x)
+        self.assertIs(result.last_device_node, n.parent)
         self.assertIs(result.last_host_node, y)
 
     def test_hicache_swa_load_back_anchored_on_best_match_node(self):