fix: handle empty HiCache match prefixes

This commit is contained in:
2026-05-08 02:47:25 +08:00
parent 829f3ebceb
commit 39f40c02b4
2 changed files with 51 additions and 0 deletions

View File

@@ -1521,6 +1521,14 @@ class HiRadixCache(RadixCache):
page_aligned_len = len(key) // self.page_size * self.page_size
key = key[:page_aligned_len]
if len(key) == 0:
return MatchResult(
device_indices=empty_value,
last_device_node=self.root_node,
last_host_node=self.root_node,
host_hit_length=0,
)
value, last_node = self._match_prefix_helper(self.root_node, key)
if value:
value = torch.cat(value)

View File

@@ -609,6 +609,49 @@ class TestHiRadixCacheCPLoadBack(CustomTestCase):
self.assertIs(result.last_host_node, node)
self.assertIs(result.last_device_node, root)
def test_cp_match_prefix_shorter_than_page_returns_empty_root_match(self):
cache = HiRadixCache.__new__(HiRadixCache)
cache._uses_cp_hicache = True
cache.device = "cpu"
cache.disable = False
cache.page_size = 4
cache.maybe_bigram_convert = lambda key: (key, None)
cache.get_child_key_fn = lambda key: key.token_ids[0]
cache.root_node = TreeNode()
result = cache.match_prefix(MatchPrefixParams(key=RadixKey([1, 2, 3])))
self.assertEqual(result.device_indices.tolist(), [])
self.assertIs(result.last_device_node, cache.root_node)
self.assertIs(result.last_host_node, cache.root_node)
self.assertEqual(result.host_hit_length, 0)
def test_cp_load_back_non_evicted_node_returns_none_without_loading(self):
cache = HiRadixCache.__new__(HiRadixCache)
cache._uses_cp_hicache = True
cache.load_back_threshold = 1
cache.evictable_size_ = 0
cache.metrics_collector = None
cache.ongoing_load_back = {}
cache.inc_lock_ref = lambda node: type("Result", (), {"delta": 0})()
cache.dec_lock_ref = lambda node: None
cache.cache_controller = type(
"Controller",
(),
{
"load_cp": lambda self, nodes, node_id=-1: (_ for _ in ()).throw(
AssertionError("must not load without evicted nodes")
)
},
)()
node = TreeNode()
node.value = torch.empty((0,), dtype=torch.int64)
loaded = cache.load_back(node)
self.assertIsNone(loaded)
self.assertEqual(cache.ongoing_load_back, {})
if __name__ == "__main__":
unittest.main()