Fix prefetch attempt cleanup on abort (#40262)

Co-authored-by: cctry <cctry@meta.com>
Co-authored-by: cctry <cctry@fb.com>

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."""

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