[UnifiedTree] fix: allow partial match on evicted+backuped nodes (#24943)

This commit is contained in:
shuwenn
2026-05-13 11:43:12 +08:00
committed by GitHub
parent 6140e45ef3
commit 5ed9a494d0
2 changed files with 43 additions and 3 deletions
@@ -657,10 +657,9 @@ class UnifiedRadixCache(BasePrefixCache):
prefix_len = child.key.match(key, page_size=self.page_size)
if prefix_len < len(child.key):
if child.evicted:
break
node = self._split_node(child.key, child, prefix_len)
value.append(node.component_data[BASE_COMPONENT_TYPE].value)
if not node.evicted:
value.append(node.component_data[BASE_COMPONENT_TYPE].value)
_update_best_if_valid(node)
break
@@ -1472,6 +1472,47 @@ class UnifiedRadixCacheSuite:
self.assertGreaterEqual(len(m.device_indices), len(base))
tree.sanity_check()
def test_hicache_partial_match_splits_evicted_backed_up_node(self):
"""Partial matches on host-only nodes must keep the host prefix usable."""
tree, allocator, req_to_token_pool = build_fixture(self.cfg)
ps = self.cfg.page_size
seq = self._make_seq(1, 4)
expected_prefix = seq[: 2 * ps]
expected_suffix = seq[len(expected_prefix) :]
query = expected_prefix + self._make_seq(9000, 1)
self._insert(tree, allocator, req_to_token_pool, seq)
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(seq)))
node = m.last_device_node
self._simulate_backup(tree, node)
tree.evict(EvictParams(num_tokens=len(seq)))
self.assertTrue(node.evicted)
self.assertTrue(node.backuped)
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(query)))
self.assertEqual(len(m.device_indices), 0)
self.assertIs(m.last_device_node, tree.root_node)
split_parent = node.parent
self.assertIsNot(split_parent, tree.root_node)
self.assertTrue(split_parent.evicted)
self.assertTrue(split_parent.backuped)
self.assertEqual(split_parent.key.token_ids, expected_prefix)
self.assertEqual(node.key.token_ids, expected_suffix)
if self.cfg.has_mamba:
self.assertEqual(m.host_hit_length, 0)
self.assertIs(m.last_host_node, tree.root_node)
self.assertIsNone(
split_parent.component_data[ComponentType.MAMBA].host_value
)
else:
self.assertEqual(m.host_hit_length, len(expected_prefix))
self.assertIs(m.last_host_node, split_parent)
tree.sanity_check()
def test_hicache_d_leaf_h_leaf_mutual_exclusion(self):
"""D-leaf and H-leaf sets are always disjoint."""
if self._skip_unsupported_hicache_test():