[Unified] Fix UnifiedRadixCache write_backup issue in write-back mode(#27108)

Co-authored-by: hzh0425 <hzh0425@apache.org>
This commit is contained in:
Yongji Wu
2026-06-11 08:28:24 +08:00
committed by GitHub
co-authored by hzh0425
parent 740305e1d9
commit db061e97c0
2 changed files with 33 additions and 1 deletions
@@ -1424,7 +1424,9 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache):
self.cache_controller is not None
and self.cache_controller.write_policy == "write_back"
):
self.write_backup(node, write_back=True)
written = self.write_backup(node, write_back=True)
if written == 0:
return
self.writing_check(write_back=True)
self._evict_to_host(node, tracker)
return
@@ -2452,6 +2452,36 @@ 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."""
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")
ct = ComponentType.FULL
seq = self._make_seq(1, 2)
self._insert(tree, allocator, req_to_token_pool, seq)
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq))))
node = m.last_device_node
self.assertIsNot(node, tree.root_node)
self.assertFalse(node.backuped)
self.assertFalse(node.evicted)
tracker = {c: 0 for c in tree.tree_components}
with mock.patch.object(tree, "write_backup", return_value=0):
tree._evict_device_leaf(node, tracker)
self.assertFalse(node.evicted)
self.assertIsNotNone(node.component_data[ct].value)
self.assertIsNone(node.component_data[ct].host_value)
with self.assertRaises(AssertionError):
tree._evict_to_host(node, {c: 0 for c in tree.tree_components})
tree.sanity_check()
def test_hicache_node_states(self):
"""Verify device-only to device+host transition after real backup."""
if self._skip_unsupported_hicache_test():