From 6e1338dd1eb1546518bdffe3b69d2ea0e011ef28 Mon Sep 17 00:00:00 2001 From: Lianmin Zheng Date: Fri, 18 Sep 2026 22:46:17 -0700 Subject: [PATCH] Fix prefetch attempt cleanup on abort (#40262) Co-authored-by: cctry Co-authored-by: cctry --- python/sglang/srt/disaggregation/prefill.py | 6 ++---- python/sglang/srt/mem_cache/unified_radix_cache.py | 8 +++++--- 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index abb247f3d..e1db8cd18 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/python/sglang/srt/disaggregation/prefill.py @@ -1134,8 +1134,7 @@ class SchedulerDisaggregationPrefillMixin: req.update_finish_state() maybe_release_metadata_buffer(req, self.req_to_metadata_buffer_idx_allocator) req.pending_bootstrap = False - if self.enable_hicache_storage: - self.tree_cache.finish(req.cache_request_handle, CacheRequestOutcome.ABORT) + self.tree_cache.finish(req.cache_request_handle, CacheRequestOutcome.ABORT) if req.kv.holds_kv or req.kv.holds_mamba: release_kv_cache(req, self.tree_cache, is_insert=False) return True @@ -1166,8 +1165,7 @@ class SchedulerDisaggregationPrefillMixin: self.output_streamer.stream_output([req], req.return_logprob) if self.metrics_reporter.enable_metrics: self.metrics_collector.increment_bootstrap_failed_reqs() - if self.enable_hicache_storage: - self.tree_cache.finish(req.cache_request_handle, CacheRequestOutcome.ABORT) + self.tree_cache.finish(req.cache_request_handle, CacheRequestOutcome.ABORT) def handle_pending_bootstrap(self: Scheduler, req: Req, poll: KVPoll) -> bool: """Return True when bootstrap is finalized and KV transfer can proceed.""" diff --git a/python/sglang/srt/mem_cache/unified_radix_cache.py b/python/sglang/srt/mem_cache/unified_radix_cache.py index dd79204d3..b27536bb0 100644 --- a/python/sglang/srt/mem_cache/unified_radix_cache.py +++ b/python/sglang/srt/mem_cache/unified_radix_cache.py @@ -391,7 +391,9 @@ class UnifiedRadixCache(BasePrefixCache): # Rank-agreed L3-hit tokens not yet resolved as usable or unfulfilled. # Cache-mode entries survive L3->L2 until H2D succeeds or admission # fails; buffer-mode entries survive staging until the H2D ack. - self._storage_prefetch_hit_remaining_by_reqid: dict[str, int] = {} + self._storage_prefetch_hit_remaining_by_reqid: dict[ + CacheRequestHandle, int + ] = {} self.storage_prefetch_retries = StoragePrefetchRetries() self.ongoing_backup: dict[int, tuple[NodeId, DecLockRefParams]] = {} if self.buffer_pipeline is not None: @@ -2457,8 +2459,8 @@ class UnifiedRadixCache(BasePrefixCache): rid = request.rid if self.linker is not None: self.linker.release_request(rid) - self.prefetch_loaded_tokens_by_reqid.pop(rid, None) - self.prefetch_loaded_storage_start_by_reqid.pop(rid, None) + self.prefetch_loaded_tokens_by_reqid.pop(request, None) + self.prefetch_loaded_storage_start_by_reqid.pop(request, None) self.storage_prefetch_retries.cancel(rid) if ( self.buffer_pipeline is not None