minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[HiCache] Forward prefix metadata to v2 storage calls (#39567)

Browse Source 
Co-authored-by: Zhangheng <hzh0425@apache.org>
This commit is contained in:
Shuwen Wang
2026-09-16 14:05:42 +08:00
committed by GitHub
co-authored by Zhangheng
parent 91f691c490
commit 6b33338242
2 changed files with 16 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/managers/cache_controller.py
+8 -8
View File
@@ -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):
python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py
+8 -2
View File
@@ -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)