[Unified] Fix UnifiedRadixCache write_backup issue in write-back mode(#27108)
Co-authored-by: hzh0425 <hzh0425@apache.org>
This commit is contained in:
@@ -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():
|
||||
|
||||
Reference in New Issue
Block a user