diff --git a/python/sglang/srt/mem_cache/unified_cache_components/full_component.py b/python/sglang/srt/mem_cache/unified_cache_components/full_component.py index be506660f..0e62d7007 100644 --- a/python/sglang/srt/mem_cache/unified_cache_components/full_component.py +++ b/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) diff --git a/python/sglang/srt/mem_cache/unified_radix_cache.py b/python/sglang/srt/mem_cache/unified_radix_cache.py index 3694620b6..7caf33b45 100644 --- a/python/sglang/srt/mem_cache/unified_radix_cache.py +++ b/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: 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 3245b42e0..9c9247153 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 @@ -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):