[HiCache] Release buffer prefetch anchor locks during storage cleanup (#38483)

python/sglang/srt/mem_cache/unified_cache/storage_attachment.py

@@ -374,7 +374,13 @@ class StorageAttachment:
                     continue
                 completed_tokens, _ = controller.terminate_prefetch(info.operation)
                 del cache.ongoing_prefetch[req_id]
-                cache.dec_host_lock_ref(info.anchor_node_id, info.anchor_lock_params)
+                if info.anchor_lock_params is not None:
+                    cache.dec_host_lock_ref(
+                        info.anchor_node_id, info.anchor_lock_params
+                    )
+                if cache.buffer_pipeline is not None:
+                    cache.buffer_pipeline.pop_prefix_ctx(req_id)
+                    cache.buffer_pipeline.release_anchor_lock(req_id)
                 controller.append_host_mem_release(
                     host_indices=info.host_indices[:completed_tokens],
                     extra_pools=[
@@ -382,7 +388,11 @@ class StorageAttachment:
                     ],
                 )
                 controller.prefetch_tokens_occupied = max(
-                    0, controller.prefetch_tokens_occupied - len(info.prefetch_key)
+                    0,
+                    controller.prefetch_tokens_occupied
+                    - cache._prefetch_occupied_span(
+                        info.prefetch_key, info.host_indices
+                    ),
                 )
             except Exception:
                 logger.exception("Failed to release pending prefetch %s", req_id)

test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py

@@ -9535,6 +9535,43 @@ class TestAnchorLockOutcomePolicy(CustomTestCase):
         self.assertEqual(pipeline.anchor_locks, {})
         self.assertEqual(pipeline.anchor_locked_tokens_, 0)
 
+    def test_storage_cleanup_releases_buffer_prefetch_anchor(self):
+        cache = self._make_cache(live_match_len=len(self._PREFIX))
+        pipeline = self._make_pipeline(cache)
+        cache.buffer_pipeline = pipeline
+        self.assertEqual(pipeline.try_lock_anchor(self._REQ), "locked")
+        host_indices = torch.arange(4)
+        cache.ongoing_prefetch = {
+            self._REQ: _OngoingPrefetch(
+                anchor_node_id=99,
+                prefetch_key=RadixKey(array("q", range(16))),
+                host_indices=host_indices,
+                operation=mock.Mock(),
+                anchor_lock_params=None,
+                comp_xfers={},
+            )
+        }
+        cache.ongoing_backup = {}
+        cache.host_memory_mode = "buffer_only"
+        cache._prefetch_occupied_span.side_effect = lambda key, indices: (
+            UnifiedRadixCache._prefetch_occupied_span(cache, key, indices)
+        )
+        controller = cache.cache_controller
+        controller.terminate_prefetch.return_value = (4, None)
+        controller.prefetch_tokens_occupied = 12
+
+        StorageAttachment(cache)._release_pending_storage_ops()
+
+        self.assertEqual(cache.ongoing_prefetch, {})
+        self.assertEqual(pipeline.anchor_locks, {})
+        self.assertEqual(pipeline.anchor_locked_tokens_, 0)
+        self.assertNotIn(self._REQ, pipeline._prefetch_prefix_ctx)
+        cache.dec_lock_ref.assert_called_once_with(
+            99, cache.inc_lock_ref.return_value.to_dec_params.return_value
+        )
+        cache.dec_host_lock_ref.assert_not_called()
+        self.assertEqual(controller.prefetch_tokens_occupied, 8)
+
     def test_positive_hit_with_lost_anchor_is_reported_as_shrunk(self):
         cache = UnifiedRadixCache.__new__(UnifiedRadixCache)
         cache._storage_prefetch_missed_rids = set()