fix: address CP HiCache review regressions

test/registered/unit/mem_cache/test_cp_hicache_metadata.py

@@ -416,6 +416,9 @@ class TestHiRadixCacheCPSplitEvict(CustomTestCase):
         cache._clear_pin = lambda node: None
         cache._record_remove_event = lambda node: None
         cache._update_host_leaf_status = lambda node: None
+        cache._node_host_evict_indices = lambda node: torch.tensor(
+            [99], dtype=torch.int64
+        )
         freed = []
         cache.cache_controller = type(
             "Controller",
@@ -441,7 +444,7 @@ class TestHiRadixCacheCPSplitEvict(CustomTestCase):
         physical_freed = cache._evict_host_for_physical_slots(1)
 
         self.assertEqual(physical_freed, 1)
-        self.assertEqual(freed[0].tolist(), [70])
+        self.assertEqual(freed[0].tolist(), [99])
         self.assertEqual(node.host_len, 0)
         self.assertIsNone(node.cp_hicache)
 
@@ -567,7 +570,7 @@ class TestHiRadixCacheCPSplitEvict(CustomTestCase):
         cache = HiRadixCache.__new__(HiRadixCache)
         cache._uses_cp_hicache = True
         cache.tp_world_size = 2
-        cache.tp_group = None
+        cache.tp_group = object()
         cache.root_node = TreeNode()
         cache.root_node.key = RadixKey([])
         cache.evictable_host_leaves = set()
@@ -600,6 +603,7 @@ class TestHiRadixCacheCPSplitEvict(CustomTestCase):
         cache.evictable_host_leaves.add(node)
 
         def mark_not_done(done_tensor, op=None, group=None):
+            self.assertIs(group, cache.tp_group)
             done_tensor.fill_(0)
 
         with patch("torch.distributed.all_reduce", side_effect=mark_not_done):
@@ -613,6 +617,55 @@ class TestHiRadixCacheCPSplitEvict(CustomTestCase):
         self.assertEqual(node.host_len, 0)
         self.assertIsNone(node.cp_hicache)
 
+    def test_cp_host_eviction_skips_all_reduce_without_tp_group(self):
+        cache = HiRadixCache.__new__(HiRadixCache)
+        cache._uses_cp_hicache = True
+        cache.tp_world_size = 2
+        cache.tp_group = None
+        cache.root_node = TreeNode()
+        cache.root_node.key = RadixKey([])
+        cache.evictable_host_leaves = set()
+        cache.get_child_key_fn = lambda key: key.token_ids[0]
+        cache.eviction_strategy = type(
+            "Strategy", (), {"get_priority": lambda self, node: 0}
+        )()
+        cache._clear_pin = lambda node: None
+        cache._record_remove_event = lambda node: None
+        cache.cache_controller = type(
+            "Controller",
+            (),
+            {
+                "evict_host": lambda self, indices: (_ for _ in ()).throw(
+                    AssertionError("must not evict host slots")
+                )
+            },
+        )()
+        node = TreeNode()
+        node.parent = cache.root_node
+        node.key = RadixKey([1, 2, 3, 4])
+        node.value = None
+        node.host_len = 4
+        node.cp_hicache = CpHiCacheNodeMetadata(
+            logical_len=4,
+            owned_positions=torch.empty((0,), dtype=torch.int64),
+            host_indices=torch.empty((0,), dtype=torch.int64),
+        )
+        cache.root_node.children[1] = node
+        cache.evictable_host_leaves.add(node)
+
+        with patch(
+            "torch.distributed.all_reduce",
+            side_effect=AssertionError("must not all_reduce without tp_group"),
+        ):
+            physical_freed = cache._evict_host_for_physical_slots(
+                0, synchronize_across_ranks=True
+            )
+
+        self.assertEqual(physical_freed, 0)
+        self.assertIn(1, cache.root_node.children)
+        self.assertEqual(node.host_len, 4)
+        self.assertIsNotNone(node.cp_hicache)
+
 
 class TestHiRadixCacheCPLoadBack(CustomTestCase):
     def test_cp_load_back_uses_host_len_not_host_value(self):
@@ -742,6 +795,35 @@ class TestHiRadixCacheCPLoadBack(CustomTestCase):
         self.assertIs(result.last_host_node, cache.root_node)
         self.assertEqual(result.host_hit_length, 0)
 
+    def test_non_cp_match_prefix_uses_root_when_no_host_backup_exists(self):
+        cache = HiRadixCache.__new__(HiRadixCache)
+        cache._uses_cp_hicache = False
+        cache.device = "cpu"
+        cache.disable = False
+        cache.page_size = 1
+        cache.get_child_key_fn = lambda key: key.token_ids[0]
+        cache.key_match_fn = lambda child_key, key: sum(
+            1 for lhs, rhs in zip(child_key.token_ids, key.token_ids) if lhs == rhs
+        )
+        cache.maybe_bigram_convert = lambda key: (key, None)
+        root = TreeNode()
+        root.key = RadixKey([])
+        root.value = torch.empty((0,), dtype=torch.int64)
+        root.host_value = None
+        cache.root_node = root
+        node = TreeNode()
+        node.parent = root
+        node.key = RadixKey(list(range(4)))
+        node.value = torch.arange(4, dtype=torch.int64)
+        node.host_value = None
+        root.children[0] = node
+
+        result = cache.match_prefix(MatchPrefixParams(key=RadixKey(list(range(4)))))
+
+        self.assertEqual(result.host_hit_length, 0)
+        self.assertIs(result.last_host_node, root)
+        self.assertIs(result.last_device_node, node)
+
     def test_cp_load_back_non_evicted_node_returns_none_without_loading(self):
         cache = HiRadixCache.__new__(HiRadixCache)
         cache._uses_cp_hicache = True