minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 13 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 # Check termination
if hit_pages != len(batch_hashes): if hit_pages != len(batch_hashes):
all_success = False all_success = False
if prefix_keys and len(prefix_keys) > 0: if prefix_keys is not None:
prefix_keys += batch_hashes prefix_keys = prefix_keys + batch_hashes
completed_pages += hit_pages completed_pages += hit_pages
ack = PrefetchAck( ack = PrefetchAck(
rid=operation.request_id, rid=operation.request_id,
@@ -1143,7 +1143,7 @@ class HiCacheController:
for transfer in kv_derived_transfers for transfer in kv_derived_transfers
] ]
sidecar_results = self.storage_backend.batch_get_v2( 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) sidecar_hits = count_pool_hits(sidecar_results)
@@ -1194,7 +1194,7 @@ class HiCacheController:
def _storage_hit_query(self, operation) -> tuple[list[str], int]: def _storage_hit_query(self, operation) -> tuple[list[str], int]:
last_hash = operation.last_hash last_hash = operation.last_hash
tokens_to_fetch = operation.token_ids 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 storage_query_count = 0
hash_value = [] hash_value = []
@@ -1211,8 +1211,8 @@ class HiCacheController:
storage_query_count += hit_page_num * self.page_size storage_query_count += hit_page_num * self.page_size
if hit_page_num < len(batch_hashes): if hit_page_num < len(batch_hashes):
break break
if prefix_keys and len(prefix_keys) > 0: if prefix_keys is not None:
prefix_keys += batch_hashes prefix_keys = prefix_keys + batch_hashes
return hash_value, storage_query_count return hash_value, storage_query_count
@@ -1298,8 +1298,8 @@ class HiCacheController:
) )
break break
if prefix_keys and len(prefix_keys) > 0: if prefix_keys is not None:
prefix_keys += batch_hashes prefix_keys = prefix_keys + batch_hashes
operation.completed_tokens += self.page_size * len(batch_hashes) operation.completed_tokens += self.page_size * len(batch_hashes)
def backup_thread_func(self): 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._sync_trailing_keys(transfers_nonkv, sidecar_hashes, sidecar_hit_pages)
self._resolve_sidecar_nonkv_derived_pool_transfers(operation) 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) pool_hits = count_pool_hits(results)
# Emit PrefetchAck to prefetch_sync_queue, even the operation has been canceled by the # 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 # scheduler thread. The prefetch sync thread expects the same number of PrefetchAck objects
@@ -731,7 +734,10 @@ class HybridCacheController(BaseHiCacheController):
if backup_transfers: if backup_transfers:
self._resolve_sidecar_kv_derived_pool_transfers(operation) self._resolve_sidecar_kv_derived_pool_transfers(operation)
self._resolve_sidecar_nonkv_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) pool_hits = count_pool_hits(results)
operation.pool_storage_result.update_extra_pool_hit_pages(pool_hits) operation.pool_storage_result.update_extra_pool_hit_pages(pool_hits)