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

[HiCache] Add model-aware key isolation to Mooncake Store (#31920)

Browse Source
This commit is contained in:
huangtingwei
2026-07-23 10:20:13 +08:00
committed by GitHub
parent 60dea26077
commit 86e1bb584d
2 changed files with 40 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py
+17 -10
View File
@@ -400,11 +400,18 @@ class MooncakeStore(HiCacheStorage, MooncakeBaseStore):
self.config.global_segment_size // tp_scale_factor
)
# Check if extra_backend_tag should be passed to MooncakeDistributedStore
self.extra_backend_tag = None
if extra_config and "extra_backend_tag" in extra_config:
self.extra_backend_tag = extra_config["extra_backend_tag"]
logger.info(f"Using extra_backend_tag: {self.extra_backend_tag}")
# Use the backend tag and model name as a prefix to isolate tenants
# and models sharing one store.
self.config_prefix = None
config_prefix_parts = []
if extra_config and extra_config.get("extra_backend_tag") is not None:
config_prefix_parts.append(str(extra_config["extra_backend_tag"]))
if storage_config is not None and storage_config.model_name:
model_name = "-".join(storage_config.model_name.split("/"))
config_prefix_parts.append(model_name)
if config_prefix_parts:
self.config_prefix = "_".join(config_prefix_parts)
logger.info(f"Using Mooncake config prefix: {self.config_prefix}")
# Check server status
if self.config.check_server:
@@ -681,9 +688,9 @@ class MooncakeStore(HiCacheStorage, MooncakeBaseStore):
super().register_buffer(buf)
def _tag_keys(self, keys: List[str]) -> List[str]:
if self.extra_backend_tag is None:
if self.config_prefix is None:
return keys
return [f"{self.extra_backend_tag}_{key}" for key in keys]
return [f"{self.config_prefix}_{key}" for key in keys]
def _can_use_group_semantics(self) -> bool:
return self._use_group_semantics
@@ -1005,7 +1012,7 @@ class MooncakeStore(HiCacheStorage, MooncakeBaseStore):
# DeepSeek V4's KV anchor is logical only; v2 side pools carry data.
return [True] * len(keys)
# Apply extra_backend_tag prefix if available
# Apply config prefix if available.
keys = self._tag_keys(keys)
key_strs, buffer_ptrs, buffer_sizes = self._batch_preprocess(keys, host_indices)
@@ -1034,7 +1041,7 @@ class MooncakeStore(HiCacheStorage, MooncakeBaseStore):
# DeepSeek V4's KV anchor is logical only; v2 side pools carry data.
return [True] * len(keys)
# Apply extra_backend_tag prefix if available
# Apply config prefix if available.
keys = self._tag_keys(keys)
key_strs, buffer_ptrs, buffer_sizes = self._batch_preprocess(keys, host_indices)
@@ -1208,7 +1215,7 @@ class MooncakeStore(HiCacheStorage, MooncakeBaseStore):
def batch_exists(
self, keys, extra_info: Optional[HiCacheStorageExtraInfo] = None
) -> int:
# Apply extra_backend_tag prefix if available
# Apply config prefix if available.
keys = self._tag_keys(keys)
if self.is_mla_backend: