[HiCache] Configurable and Dynamic Prefetch Timeout (#10512)

Co-authored-by: Zhiqiang Xie <xiezhq@stanford.edu>
This commit is contained in:
ykwd
2025-10-01 21:44:10 +08:00
committed by GitHub
parent 86cb4db058
commit bfa274380b
2 changed files with 75 additions and 21 deletions

View File

@@ -250,7 +250,7 @@ class HiCacheController:
storage_backend: Optional[str] = None,
prefetch_threshold: int = 256,
model_name: Optional[str] = None,
storage_backend_extra_config: Optional[str] = None,
storage_backend_extra_config: Optional[dict] = None,
):
self.mem_pool_device_allocator = token_to_kv_pool_allocator
self.mem_pool_device = token_to_kv_pool_allocator.get_kvcache()
@@ -361,7 +361,7 @@ class HiCacheController:
def _generate_storage_config(
self,
model_name: Optional[str] = None,
storage_backend_extra_config: Optional[str] = None,
storage_backend_extra_config: Optional[dict] = None,
):
if is_dp_attention_enabled():
@@ -376,23 +376,13 @@ class HiCacheController:
# Currently, AscendMLAPagedTokenToKVPool is the subclass of MLATokenToKVPool.
is_mla_backend = isinstance(self.mem_pool_device, MLATokenToKVPool)
# Parse extra config JSON if provided
extra_config = None
if storage_backend_extra_config:
try:
import json
extra_config = json.loads(storage_backend_extra_config)
except Exception as e:
logger.error(f"Invalid backend extra config JSON: {e}")
return HiCacheStorageConfig(
tp_rank=self.tp_rank,
tp_size=self.tp_size,
is_mla_model=is_mla_backend,
is_page_first_layout=self.mem_pool_host.layout == "page_first",
model_name=model_name,
extra_config=extra_config,
extra_config=storage_backend_extra_config,
)
def reset(self):