HiCache Storage fix host memory leak (#9648)

This commit is contained in:
Zhiqiang Xie
2025-08-26 10:49:40 -07:00
committed by GitHub
parent 79ce3688bb
commit 43de1d7304
3 changed files with 8 additions and 4 deletions

View File

@@ -1296,10 +1296,11 @@ class Scheduler(
def _prefetch_kvcache(self, req: Req):
if self.enable_hicache_storage:
req.init_next_round_input(self.tree_cache)
last_hash = req.last_host_node.get_last_hash_value()
matched_len = len(req.prefix_indices) + req.host_hit_length
# todo, free-form fetching, calculating hash keys on the fly
if (matched_len > 0 and last_hash is not None) or matched_len == 0:
if req.last_node.backuped:
# only to initiate the prefetch if the last node is backuped
# otherwise, the allocated GPU memory must be locked for integrity
last_hash = req.last_host_node.get_last_hash_value()
matched_len = len(req.prefix_indices) + req.host_hit_length
new_input_tokens = req.fill_ids[matched_len:]
self.tree_cache.prefetch_from_storage(
req.rid, req.last_host_node, new_input_tokens, last_hash