[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
|
self.cache_controller is not None
|
||||||
and self.cache_controller.write_policy == "write_back"
|
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.writing_check(write_back=True)
|
||||||
self._evict_to_host(node, tracker)
|
self._evict_to_host(node, tracker)
|
||||||
return
|
return
|
||||||
|
|||||||
@@ -2452,6 +2452,36 @@ class UnifiedRadixCacheSuite:
|
|||||||
[conv[:, mamba_indices].float().cpu().clone() for conv in mamba_cache.conv],
|
[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):
|
def test_hicache_node_states(self):
|
||||||
"""Verify device-only to device+host transition after real backup."""
|
"""Verify device-only to device+host transition after real backup."""
|
||||||
if self._skip_unsupported_hicache_test():
|
if self._skip_unsupported_hicache_test():
|
||||||
|
|||||||
Reference in New Issue
Block a user