fix(unified radix cache w/ hicache): backup ancestor nodes before leaf in write_back eviction (#24585)

Co-authored-by: Zhangheng <hzh0425@apache.org>

python/sglang/srt/mem_cache/unified_radix_cache.py

@@ -390,6 +390,12 @@ class UnifiedRadixCache(BasePrefixCache):
         for component in self._components_tuple:
             component.drive_eviction(params=params, tracker=tracker)
 
+        if (
+            self.cache_controller is not None
+            and self.cache_controller.write_policy == "write_back"
+        ):
+            self.writing_check(write_back=True)
+
         self.update_eviction_metrics(sum(tracker.values()), start_time)
         return EvictResult(
             num_tokens_evicted=tracker[BASE_COMPONENT_TYPE],
@@ -1653,6 +1659,11 @@ class UnifiedRadixCache(BasePrefixCache):
         if self.session.any_holding_kv():
             return
 
+        write_back = (
+            self.cache_controller is not None
+            and self.cache_controller.write_policy == "write_back"
+        )
+
         errors: list[str] = []
         E = errors.append
         all_nodes = self._collect_all_nodes()
@@ -1709,7 +1720,7 @@ class UnifiedRadixCache(BasePrefixCache):
                 p_hst = node.parent.component_data[FCT].host_value is not None
                 if full_dev and not p_dev:
                     E(f"node {nid} device present but parent {node.parent.id} evicted")
-                if full_hst and not p_hst:
+                if full_hst and not p_hst and not write_back:
                     E(f"node {nid} backed up but parent {node.parent.id} not backed up")
 
             # Lock hierarchy and counters must stay sane.

test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py

@@ -1291,7 +1291,7 @@ class UnifiedRadixCacheSuite:
                 self._simulate_backup(tree, node)
             stack.extend(node.children.values())
 
-    def _init_hicache(self, tree):
+    def _init_hicache(self, tree, *, write_policy: str = "write_through"):
         import sglang.srt.mem_cache.hybrid_cache.hybrid_pool_assembler as assembler
 
         orig_kv_host_pool = assembler.MHATokenToKVPoolHost
@@ -1325,7 +1325,7 @@ class UnifiedRadixCacheSuite:
             model_path="dummy",
             page_size=self.cfg.page_size,
             hicache_io_backend="direct",
-            hicache_write_policy="write_through",
+            hicache_write_policy=write_policy,
         )
         set_global_server_args_for_scheduler(server_args)
         tree.init_hicache(server_args, tree.cache_init_params)
@@ -2431,6 +2431,46 @@ class UnifiedRadixCacheSuite:
         overlap = tree.evictable_device_leaves & tree.evictable_host_leaves
         self.assertEqual(len(overlap), 0)
 
+    def test_hicache_write_back_leaf_backup(self):
+        """write_back: evicting a device leaf backs it up to host"""
+        if self._skip_unsupported_hicache_test():
+            return
+        tree, allocator, req_to_token_pool = build_fixture(self.cfg)
+        self._init_hicache(tree, write_policy="write_back")
+
+        base = self._make_seq(1, 2)
+        leaf_seq = base + self._make_seq(500, 2)
+        self._insert(tree, allocator, req_to_token_pool, base)
+        self._insert(tree, allocator, req_to_token_pool, leaf_seq)
+
+        m = tree.match_prefix(MatchPrefixParams(key=RadixKey(leaf_seq)))
+        leaf = m.last_device_node
+        parent = leaf.parent
+        self.assertIsNot(parent, tree.root_node)
+
+        self.assertFalse(leaf.backuped)
+        self.assertFalse(parent.backuped)
+
+        lr = tree.inc_lock_ref(parent)
+        try:
+            evict_tokens = len(leaf_seq) - len(base)
+            tree.evict(EvictParams(num_tokens=evict_tokens))
+        finally:
+            tree.dec_lock_ref(
+                parent,
+                DecLockRefParams(
+                    swa_uuid_for_lock=getattr(lr, "swa_uuid_for_lock", None)
+                ),
+            )
+
+        self.assertTrue(leaf.evicted, "leaf should be demoted to host")
+        self.assertTrue(leaf.backuped, "write_back must back up the leaf on eviction")
+        self.assertFalse(
+            parent.backuped, "parent must NOT be backed up under write_back"
+        )
+
+        tree.sanity_check()
+
 
 _CONFIGS: list[CacheConfig] = [
     CacheConfig(page_size=1, components=(ComponentType.FULL,)),