diff --git a/python/sglang/srt/managers/cache_controller.py b/python/sglang/srt/managers/cache_controller.py index 0caa63cde..80174584e 100644 --- a/python/sglang/srt/managers/cache_controller.py +++ b/python/sglang/srt/managers/cache_controller.py @@ -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( diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 4ff5845f5..89e222109 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -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." diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index 09d395572..5a9e8715b 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/python/sglang/srt/mem_cache/hiradix_cache.py @@ -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)