fix: handle empty HiCache match prefixes
This commit is contained in:
@@ -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)
|
||||
|
||||
@@ -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()
|
||||
|
||||
Reference in New Issue
Block a user