fix: treat CP HiCache metadata as backup state

This commit is contained in:
2026-05-08 01:41:53 +08:00
parent 08a95a35a6
commit bebbbfd0e8
2 changed files with 48 additions and 3 deletions

View File

@@ -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)

View File

@@ -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()