fix: treat CP HiCache metadata as backup state
This commit is contained in:
@@ -1006,7 +1006,7 @@ class HiRadixCache(RadixCache):
|
||||
|
||||
if self._is_pinned(x):
|
||||
# Still active: demote to host if possible
|
||||
if x.backuped:
|
||||
if self._node_backuped(x):
|
||||
num_evicted += self._evict_backuped(x)
|
||||
continue
|
||||
written = self.write_backup(x, write_back=True)
|
||||
@@ -1024,7 +1024,7 @@ class HiRadixCache(RadixCache):
|
||||
# Expired pin: clear and fall through to normal eviction
|
||||
self._clear_pin(x)
|
||||
|
||||
if not x.backuped:
|
||||
if not self._node_backuped(x):
|
||||
if self.cache_controller.write_policy == "write_back":
|
||||
# write to host if the node is not backuped
|
||||
num_evicted += self.write_backup(x, write_back=True)
|
||||
@@ -1047,7 +1047,7 @@ class HiRadixCache(RadixCache):
|
||||
if self.cache_controller.write_policy == "write_back":
|
||||
self.writing_check(write_back=True)
|
||||
for node in write_back_nodes:
|
||||
assert node.backuped
|
||||
assert self._node_backuped(node)
|
||||
self._evict_backuped(node)
|
||||
|
||||
self.update_eviction_metrics(num_evicted, start_time)
|
||||
|
||||
@@ -10,6 +10,7 @@ for _mod in ("sgl_kernel", "sgl_kernel.kvcacheio"):
|
||||
if _mod not in sys.modules:
|
||||
sys.modules[_mod] = MagicMock()
|
||||
|
||||
from sglang.srt.mem_cache.base_prefix_cache import EvictParams
|
||||
from sglang.srt.mem_cache.hiradix_cache import CpHiCacheNodeMetadata, HiRadixCache
|
||||
from sglang.srt.mem_cache.radix_cache import RadixKey, TreeNode
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
@@ -183,6 +184,13 @@ class FakeEvictionStrategy:
|
||||
return 0
|
||||
|
||||
|
||||
class FakeEvictDeviceController:
|
||||
write_policy = "write_through"
|
||||
|
||||
def evict_device(self, device_indices):
|
||||
return len(device_indices)
|
||||
|
||||
|
||||
class TestHiRadixCacheCPBackup(CustomTestCase):
|
||||
def test_node_backuped_uses_cp_metadata(self):
|
||||
cache = HiRadixCache.__new__(HiRadixCache)
|
||||
@@ -260,6 +268,43 @@ class TestHiRadixCacheCPBackup(CustomTestCase):
|
||||
self.assertNotIn(1, root.children)
|
||||
self.assertEqual(node.host_len, 16)
|
||||
|
||||
def test_evict_demotes_cp_backed_node_without_deleting_radix_child(self):
|
||||
cache = HiRadixCache.__new__(HiRadixCache)
|
||||
cache._uses_cp_hicache = True
|
||||
cache.cache_controller = FakeEvictDeviceController()
|
||||
cache.evictable_leaves = set()
|
||||
cache.evictable_host_leaves = set()
|
||||
cache.eviction_strategy = FakeEvictionStrategy()
|
||||
cache.evictable_size_ = 4
|
||||
cache.protected_size_ = 0
|
||||
cache._record_remove_event = lambda node: (_ for _ in ()).throw(
|
||||
AssertionError("must not delete backed radix child")
|
||||
)
|
||||
|
||||
root = TreeNode()
|
||||
root.key = RadixKey(token_ids=[], extra_key=None)
|
||||
root.value = []
|
||||
cache.root_node = root
|
||||
node = TreeNode()
|
||||
node.parent = root
|
||||
node.key = RadixKey(token_ids=[1, 2, 3, 4], extra_key=None)
|
||||
node.value = torch.arange(4, dtype=torch.int64)
|
||||
node.host_len = 4
|
||||
node.cp_hicache = CpHiCacheNodeMetadata(
|
||||
logical_len=4,
|
||||
owned_positions=torch.tensor([0], dtype=torch.int64),
|
||||
host_indices=torch.tensor([55], dtype=torch.int64),
|
||||
)
|
||||
root.children[1] = node
|
||||
cache.evictable_leaves.add(node)
|
||||
|
||||
cache.evict(EvictParams(num_tokens=4))
|
||||
|
||||
self.assertIn(1, root.children)
|
||||
self.assertIsNone(node.value)
|
||||
self.assertIsNotNone(node.cp_hicache)
|
||||
self.assertEqual(node.cp_hicache.host_indices.tolist(), [55])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
Reference in New Issue
Block a user