Write-back policy fix for unified tree (#31845)

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

@@ -182,7 +182,9 @@ class FullComponent(TreeComponent):
         # Only the last host node needs to be protected.
         if lock_host:
             cd = node.component_data[ct]
-            if cd.host_value is None:
+            # write_back mode: the anchor may be device-only (no host_value);
+            # pin it anyway so the host-pressure drop fallback cannot delete it.
+            if cd.host_value is None and not self.cache.is_write_back:
                 return result
             cd.host_lock_ref += 1
             self.cache._update_evictable_leaf_sets(node)
@@ -223,7 +225,10 @@ class FullComponent(TreeComponent):
         ct = self.component_type
         if lock_host:
             cd = node.component_data[ct]
-            if cd.host_value is None or cd.host_lock_ref == 0:
+            if cd.host_lock_ref == 0:
+                return
+            # Mirror of `acquire`. write_back uses a pure counter.
+            if cd.host_value is None and not self.cache.is_write_back:
                 return
             cd.host_lock_ref -= 1
             self.cache._update_evictable_leaf_sets(node)

python/sglang/srt/mem_cache/unified_radix_cache.py

@@ -617,6 +617,13 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache):
         result = self._insert_helper(self.root_node, key, value, params)
         return result
 
+    @property
+    def is_write_back(self) -> bool:
+        return (
+            self.cache_controller is not None
+            and self.cache_controller.write_policy == "write_back"
+        )
+
     def evict(self, params: EvictParams) -> EvictResult:
         if self.disable:
             return EvictResult()
@@ -1506,25 +1513,90 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache):
             ):
                 written = self.write_backup(node, write_back=True)
                 if written == 0:
+                    if self._drop_subtree_no_host(node, tracker):
+                        logger.warning(
+                            "write_back: KV subtree dropped without backup "
+                            "due to host memory pressure, root node %d",
+                            node.id,
+                        )
+                    else:
+                        logger.warning(
+                            "write_back: backup failed under host memory "
+                            "pressure but subtree drop declined (node "
+                            "locked); root node %d stays device-resident "
+                            "until host space frees",
+                            node.id,
+                        )
                     return
                 self.writing_check(write_back=True)
                 self._evict_to_host(node, tracker)
                 return
             else:
                 # Write-through: node has no backup, delete entirely.
-                self._record_remove_event(node, medium=StorageMedium.GPU)
-                for comp in self._components_tuple:
-                    self._evict_component_and_detach_lru(
-                        node, comp, target=EvictLayer.ALL, tracker=tracker
-                    )
-                self.evictable_device_leaves.discard(node)
-                parent = node.parent
-                self._remove_leaf_from_parent(node)
-                self._update_evictable_leaf_sets(parent)
-                self._iteratively_delete_tombstone_leaf(node, tracker)
+                self._delete_unbacked_device_leaf(node, tracker)
                 return
         self._evict_to_host(node, tracker)
 
+    def _drop_subtree_no_host(
+        self, node: UnifiedTreeNode, tracker: dict[ComponentType, int]
+    ) -> bool:
+        """Write-back fallback when a D-leaf's D->H backup fails under host
+        memory pressure: drop the subtree rooted at the unbacked leaf so
+        device eviction keeps making progress instead of leaving its KV
+        unevictable until host space frees up."""
+
+        assert self._is_device_leaf(node), f"node {node.id} is not a D-leaf"
+        # A failed backup never issues the D->H copy, so the subtree root has
+        # no host state and no in-flight DMA reading its device slots.
+        assert not node.backuped and node.write_through_pending_id is None
+        if any(cd.host_lock_ref > 0 for cd in node.component_data):
+            return False
+        descendants: list[UnifiedTreeNode] = []
+        stack = list(node.children.values())
+        while stack:
+            cur = stack.pop()
+            if any(
+                cd.lock_ref > 0 or cd.host_lock_ref > 0 for cd in cur.component_data
+            ):
+                return False
+            descendants.append(cur)
+            stack.extend(cur.children.values())
+        for desc in reversed(descendants):
+            # Host-only by construction: a device descendant would contradict
+            # this node being a D-leaf, and D-leaves evict before ancestors.
+            assert desc.evicted and desc.backuped, f"node {desc.id} not host-only"
+            assert desc.write_through_pending_id is None
+            self._release_all_component_layers(desc, StorageMedium.CPU, tracker)
+            self._remove_leaf_from_parent(desc)
+        self._delete_unbacked_device_leaf(node, tracker)
+        return True
+
+    def _release_all_component_layers(
+        self,
+        node: UnifiedTreeNode,
+        medium: StorageMedium,
+        tracker: dict[ComponentType, int],
+    ) -> None:
+        """Free every component layer on the node and detach it from the LRU
+        lists and evictable leaf sets."""
+        self._record_remove_event(node, medium=medium)
+        for comp in self._components_tuple:
+            self._evict_component_and_detach_lru(
+                node, comp, target=EvictLayer.ALL, tracker=tracker
+            )
+        self.evictable_device_leaves.discard(node)
+        self.evictable_host_leaves.discard(node)
+
+    def _delete_unbacked_device_leaf(
+        self, node: UnifiedTreeNode, tracker: dict[ComponentType, int]
+    ) -> None:
+        """Delete a device leaf that has no host backup, freeing all layers."""
+        self._release_all_component_layers(node, StorageMedium.GPU, tracker)
+        parent = node.parent
+        self._remove_leaf_from_parent(node)
+        self._update_evictable_leaf_sets(parent)
+        self._iteratively_delete_tombstone_leaf(node, tracker)
+
     def _evict_host_leaf(
         self, node: UnifiedTreeNode, tracker: dict[ComponentType, int]
     ) -> None:

test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py

@@ -388,6 +388,7 @@ class TestUnifiedRadixCacheEagleHiCacheStorageKey(CustomTestCase):
                 self.mem_pool_host = FakeHostPool()
                 self.prefetch_tokens_occupied = 0
                 self.prefetch_args = None
+                self.write_policy = "write_through"
 
             def prefetch_rate_limited(self):
                 return False
@@ -2734,34 +2735,94 @@ class UnifiedRadixCacheSuite:
             [conv[:, mamba_indices].float().cpu().clone() for conv in mamba_cache.conv],
         )
 
-    def test_hicache_evict_device_leaf_aborts_demote_when_backup_fails(self):
-        """when write_backup cannot allocate host pool,
-        _evict_device_leaf should not evict it to host."""
+    def test_hicache_write_back_evict_drops_unbacked_leaf_when_host_full(self):
+        """Write-back eviction will keep freeing device KV when the host pool
+        is exhausted and host eviction cannot free space to prevent OOM."""
         if self._skip_unsupported_hicache_test():
             return
         cache, allocator, req_to_token_pool = build_fixture(self.cfg)
         self._init_hicache(cache, write_policy="write_back")
-        ct = ComponentType.FULL
 
-        seq = self._make_seq(1, 2)
+        # Two-node chain: the drop must cascade parent-ward across eviction
+        # iterations, not just delete a single leaf.
+        seq_parent = self._make_seq(1, 2)
+        self._insert(cache, allocator, req_to_token_pool, seq_parent)
+        seq = seq_parent + self._make_seq(1000, 1)
         self._insert(cache, allocator, req_to_token_pool, seq)
         m = cache.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq))))
         node = m.last_device_node
         self.assertIsNot(node, cache.root_node)
         self.assertFalse(node.backuped)
-        self.assertFalse(node.evicted)
 
-        tracker = {c: 0 for c in cache.tree_components}
-        with mock.patch.object(cache, "write_backup", return_value=0):
-            cache._evict_device_leaf(node, tracker)
+        # Exhaust the KV host pool. The tree has no host leaves, so
+        # evict_host cannot free anything and every D->H backup fails.
+        host_pool = cache.cache_controller.mem_pool_host
+        self.assertIsNotNone(host_pool.alloc(host_pool.available_size()))
+        self.assertEqual(host_pool.available_size(), 0)
 
-        self.assertFalse(node.evicted)
-        self.assertIsNotNone(node.component_data[ct].value)
-        self.assertIsNone(node.component_data[ct].host_value)
+        result = cache.evict(EvictParams(num_tokens=len(seq)))
+        self.assertGreaterEqual(result.num_tokens_evicted, len(seq))
 
-        with self.assertRaises(AssertionError):
-            cache._evict_to_host(node, {c: 0 for c in cache.tree_components})
+        # The chain is gone entirely: no device hit, no host hit.
+        m = cache.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq))))
+        self.assertEqual(len(m.device_indices), 0)
+        self.assertEqual(m.host_hit_length, 0)
+        cache.sanity_check()
 
+    def test_hicache_write_back_drop_respects_pins_then_frees_subtree(self):
+        """The host-pressure drop fallback must decline while any node in the
+        unbacked subtree is host-pinned, then reclaim the whole subtree --
+        including a demoted child's host backup -- once unpinned."""
+        if self._skip_unsupported_hicache_test():
+            return
+        cache, allocator, req_to_token_pool = build_fixture(self.cfg)
+        self._init_hicache(cache, write_policy="write_back")
+        host_pool = cache.cache_controller.mem_pool_host
+        baseline_host = host_pool.available_size()
+
+        parent_seq = self._make_seq(1, 2)
+        self._insert(cache, allocator, req_to_token_pool, parent_seq)
+        child_seq = parent_seq + self._make_seq(1000, 1)
+        self._insert(cache, allocator, req_to_token_pool, child_seq)
+        m = cache.match_prefix(MatchPrefixParams(key=RadixKey(array("q", child_seq))))
+        child = m.last_device_node
+
+        # Evict only the child leaf -> real backup + demote, leaving it
+        # host-only under a still-unbacked device parent (write-back backs up
+        # single nodes leaf-first, so this is a normal intermediate state).
+        result = cache.evict(EvictParams(num_tokens=len(child.key)))
+        self.assertGreaterEqual(result.num_tokens_evicted, len(child.key))
+        self.assertTrue(child.evicted and child.backuped)
+        parent = child.parent
+        self.assertFalse(parent.backuped)
+        self.assertGreater(baseline_host - host_pool.available_size(), 0)
+
+        # From here every backup fails (controller.write returns None), so
+        # each evict() attempts the drop fallback on the parent.
+        with mock.patch.object(cache.cache_controller, "write", return_value=None):
+            # Pinned subtree root: drop declines, chain stays intact.
+            cache.inc_host_lock_ref(parent)
+            result = cache.evict(EvictParams(num_tokens=len(parent_seq)))
+            self.assertEqual(result.num_tokens_evicted, 0)
+            cache.dec_host_lock_ref(parent)
+
+            # Pinned host-only descendant: drop declines as well.
+            cache.inc_host_lock_ref(child)
+            result = cache.evict(EvictParams(num_tokens=len(parent_seq)))
+            self.assertEqual(result.num_tokens_evicted, 0)
+            m = cache.match_prefix(
+                MatchPrefixParams(key=RadixKey(array("q", parent_seq)))
+            )
+            self.assertEqual(len(m.device_indices), len(parent_seq))
+            cache.dec_host_lock_ref(child)
+
+            # Unpinned: the subtree drops and the child's host slots return.
+            result = cache.evict(EvictParams(num_tokens=len(parent_seq)))
+        self.assertGreaterEqual(result.num_tokens_evicted, len(parent_seq))
+        self.assertEqual(host_pool.available_size(), baseline_host)
+        m = cache.match_prefix(MatchPrefixParams(key=RadixKey(array("q", child_seq))))
+        self.assertEqual(len(m.device_indices), 0)
+        self.assertEqual(m.host_hit_length, 0)
         cache.sanity_check()
 
     def test_hicache_evict_to_host(self):