fix【hicache】fix the KV cache resource occupation and invalid loading from prefetch when pending requests are aborted. (#16369)
Signed-off-by: CLFutureX <chenyongqyl@163.com> Co-authored-by: Zhiqiang Xie <xiezhq@stanford.edu>
This commit is contained in:
@@ -726,7 +726,6 @@ class HiCacheController:
|
||||
operation = self.prefetch_queue.get(block=True, timeout=1)
|
||||
if operation is None:
|
||||
continue
|
||||
|
||||
hash_value, storage_hit_count = self._storage_hit_query(operation)
|
||||
if self.tp_world_size > 1:
|
||||
storage_hit_count_tensor = torch.tensor(
|
||||
|
||||
@@ -1652,6 +1652,8 @@ class Scheduler(
|
||||
direction * recv_req.priority < direction * candidate_req.priority
|
||||
)
|
||||
if abort_existing_req:
|
||||
if self.enable_hierarchical_cache:
|
||||
self.tree_cache.terminate_prefetch(candidate_req.rid)
|
||||
self.waiting_queue.pop(idx)
|
||||
req_to_abort = candidate_req
|
||||
message = "The request is aborted by a higher priority request."
|
||||
|
||||
@@ -683,6 +683,15 @@ class HiRadixCache(RadixCache):
|
||||
|
||||
return True
|
||||
|
||||
def terminate_prefetch(self, req_id: str):
|
||||
if req_id not in self.ongoing_prefetch:
|
||||
return
|
||||
|
||||
_, _, _, operation = self.ongoing_prefetch[req_id]
|
||||
if operation.host_indices is None:
|
||||
return
|
||||
operation.mark_terminate()
|
||||
|
||||
def match_prefix(self, key: RadixKey, **kwargs):
|
||||
empty_value = torch.empty((0,), dtype=torch.int64, device=self.device)
|
||||
key, _ = self.maybe_bigram_convert(key)
|
||||
|
||||
Reference in New Issue
Block a user