[bugfix][decode hicache] _storage_hit_query use HybridPrefetchOperation (#28422)
Co-authored-by: Zhangheng <hzh0425@apache.org>
This commit is contained in:
@@ -35,6 +35,9 @@ from sglang.srt.mem_cache.hicache_storage import (
|
|||||||
from sglang.srt.mem_cache.hybrid_cache.hybrid_cache_controller import (
|
from sglang.srt.mem_cache.hybrid_cache.hybrid_cache_controller import (
|
||||||
HybridCacheController,
|
HybridCacheController,
|
||||||
)
|
)
|
||||||
|
from sglang.srt.mem_cache.hybrid_cache.hybrid_cache_controller import (
|
||||||
|
PrefetchOperation as HybridPrefetchOperation,
|
||||||
|
)
|
||||||
from sglang.srt.mem_cache.hybrid_cache.hybrid_pool_assembler import (
|
from sglang.srt.mem_cache.hybrid_cache.hybrid_pool_assembler import (
|
||||||
attach_hybrid_dsa_pool_to_hiradix_cache,
|
attach_hybrid_dsa_pool_to_hiradix_cache,
|
||||||
)
|
)
|
||||||
@@ -1249,12 +1252,21 @@ class HiRadixCache(RadixCache):
|
|||||||
if len(prefetch_key) < self.prefetch_threshold:
|
if len(prefetch_key) < self.prefetch_threshold:
|
||||||
return 0
|
return 0
|
||||||
|
|
||||||
operation = PrefetchOperation(
|
prefetch_op_cls = (
|
||||||
|
HybridPrefetchOperation
|
||||||
|
if isinstance(self.cache_controller, HybridCacheController)
|
||||||
|
else PrefetchOperation
|
||||||
|
)
|
||||||
|
extra_kwargs = {}
|
||||||
|
if prefetch_op_cls is HybridPrefetchOperation:
|
||||||
|
extra_kwargs["pool_transfers"] = self._get_extra_pools().get("extra_pools")
|
||||||
|
operation = prefetch_op_cls(
|
||||||
"__storage_hit_query__",
|
"__storage_hit_query__",
|
||||||
self.cache_controller.mem_pool_host.get_dummy_flat_data_page()[:0],
|
self.cache_controller.mem_pool_host.get_dummy_flat_data_page()[:0],
|
||||||
prefetch_key,
|
prefetch_key,
|
||||||
last_hash,
|
last_hash,
|
||||||
prefix_keys,
|
prefix_keys,
|
||||||
|
**extra_kwargs,
|
||||||
)
|
)
|
||||||
hash_values, storage_hit_count = self.cache_controller._storage_hit_query(
|
hash_values, storage_hit_count = self.cache_controller._storage_hit_query(
|
||||||
operation
|
operation
|
||||||
|
|||||||
Reference in New Issue
Block a user