HiCache Storage fix host memory leak (#9648)
This commit is contained in:
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user