From 8b22deef5bea4ebd6e73a35f88ad316857213a21 Mon Sep 17 00:00:00 2001 From: PiteXChen <44110731+CLFutureX@users.noreply.github.com> Date: Fri, 16 Jan 2026 07:14:38 +0800 Subject: [PATCH] =?UTF-8?q?fix=E3=80=90hicache=E3=80=91fix=20the=20KV=20ca?= =?UTF-8?q?che=20resource=20occupation=20=20and=20invalid=20loading=20from?= =?UTF-8?q?=20prefetch=20when=20pending=20requests=20are=20aborted.=20(#16?= =?UTF-8?q?369)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: CLFutureX Co-authored-by: Zhiqiang Xie --- python/sglang/srt/managers/cache_controller.py | 1 - python/sglang/srt/managers/scheduler.py | 2 ++ python/sglang/srt/mem_cache/hiradix_cache.py | 9 +++++++++ 3 files changed, 11 insertions(+), 1 deletion(-) 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)