From 6b33338242f8ce85cd66bac96a396a4ddf05da9e Mon Sep 17 00:00:00 2001 From: Shuwen Wang <47200617+alphabetc1@users.noreply.github.com> Date: Wed, 16 Sep 2026 14:05:42 +0800 Subject: [PATCH] [HiCache] Forward prefix metadata to v2 storage calls (#39567) Co-authored-by: Zhangheng --- python/sglang/srt/managers/cache_controller.py | 16 ++++++++-------- .../hybrid_cache/hybrid_cache_controller.py | 10 ++++++++-- 2 files changed, 16 insertions(+), 10 deletions(-) diff --git a/python/sglang/srt/managers/cache_controller.py b/python/sglang/srt/managers/cache_controller.py index b9a134e13..d38055c22 100644 --- a/python/sglang/srt/managers/cache_controller.py +++ b/python/sglang/srt/managers/cache_controller.py @@ -1100,8 +1100,8 @@ class HiCacheController: # Check termination if hit_pages != len(batch_hashes): all_success = False - if prefix_keys and len(prefix_keys) > 0: - prefix_keys += batch_hashes + if prefix_keys is not None: + prefix_keys = prefix_keys + batch_hashes completed_pages += hit_pages ack = PrefetchAck( rid=operation.request_id, @@ -1143,7 +1143,7 @@ class HiCacheController: for transfer in kv_derived_transfers ] sidecar_results = self.storage_backend.batch_get_v2( - current_kv_derived_transfers + current_kv_derived_transfers, extra_info=extra_info ) sidecar_hits = count_pool_hits(sidecar_results) @@ -1194,7 +1194,7 @@ class HiCacheController: def _storage_hit_query(self, operation) -> tuple[list[str], int]: last_hash = operation.last_hash tokens_to_fetch = operation.token_ids - prefix_keys = operation.prefix_keys.copy() if operation.prefix_keys else None + prefix_keys = operation.prefix_keys storage_query_count = 0 hash_value = [] @@ -1211,8 +1211,8 @@ class HiCacheController: storage_query_count += hit_page_num * self.page_size if hit_page_num < len(batch_hashes): break - if prefix_keys and len(prefix_keys) > 0: - prefix_keys += batch_hashes + if prefix_keys is not None: + prefix_keys = prefix_keys + batch_hashes return hash_value, storage_query_count @@ -1298,8 +1298,8 @@ class HiCacheController: ) break - if prefix_keys and len(prefix_keys) > 0: - prefix_keys += batch_hashes + if prefix_keys is not None: + prefix_keys = prefix_keys + batch_hashes operation.completed_tokens += self.page_size * len(batch_hashes) def backup_thread_func(self): diff --git a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py index 17b37747d..539b952e7 100644 --- a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py +++ b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py @@ -689,7 +689,10 @@ class HybridCacheController(BaseHiCacheController): ) self._sync_trailing_keys(transfers_nonkv, sidecar_hashes, sidecar_hit_pages) self._resolve_sidecar_nonkv_derived_pool_transfers(operation) - results = self.storage_backend.batch_get_v2(transfers_nonkv) + extra_info = HiCacheStorageExtraInfo(prefix_keys=operation.prefix_keys) + results = self.storage_backend.batch_get_v2( + transfers_nonkv, extra_info=extra_info + ) pool_hits = count_pool_hits(results) # Emit PrefetchAck to prefetch_sync_queue, even the operation has been canceled by the # scheduler thread. The prefetch sync thread expects the same number of PrefetchAck objects @@ -731,7 +734,10 @@ class HybridCacheController(BaseHiCacheController): if backup_transfers: self._resolve_sidecar_kv_derived_pool_transfers(operation) self._resolve_sidecar_nonkv_derived_pool_transfers(operation) - results = self.storage_backend.batch_set_v2(backup_transfers) + extra_info = HiCacheStorageExtraInfo(prefix_keys=operation.prefix_keys) + results = self.storage_backend.batch_set_v2( + backup_transfers, extra_info=extra_info + ) pool_hits = count_pool_hits(results) operation.pool_storage_result.update_extra_pool_hit_pages(pool_hits)