EngineX-Hygon/sglang
5
0
Fork 0
Code Issues Pull Requests Actions 7 Projects Releases Wiki Activity

[HiCache] feat: add multi tenant with prefix tag (#9256)

Browse Source
This commit is contained in:
Teng Ma
2025-10-11 00:23:28 +08:00
committed by GitHub
parent 3b9d97f335
commit 9082a7d323
1 changed files with 16 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py
View File

@@ -151,6 +151,12 @@ class MooncakeStore(HiCacheStorage):
) )
per_tp_local_buffer_size = self.config.local_buffer_size // tp_scale_factor per_tp_local_buffer_size = self.config.local_buffer_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}")
ret_code = self.store.setup( ret_code = self.store.setup(
self.config.local_hostname, self.config.local_hostname,
self.config.metadata_server, self.config.metadata_server,
@@ -257,6 +263,11 @@ class MooncakeStore(HiCacheStorage):
host_indices: torch.Tensor, host_indices: torch.Tensor,
extra_info: Optional[HiCacheStorageExtraInfo] = None, extra_info: Optional[HiCacheStorageExtraInfo] = None,
) -> List[bool]: ) -> List[bool]:
# Apply extra_backend_tag prefix if available
if self.extra_backend_tag is not None:
prefix = self.extra_backend_tag
keys = [f"{prefix}_{key}" for key in keys]
key_strs, buffer_ptrs, buffer_sizes = self._batch_preprocess(keys, host_indices) key_strs, buffer_ptrs, buffer_sizes = self._batch_preprocess(keys, host_indices)
get_results = self._get_batch_zero_copy_impl( get_results = self._get_batch_zero_copy_impl(
key_strs, buffer_ptrs, buffer_sizes key_strs, buffer_ptrs, buffer_sizes
@@ -269,6 +280,11 @@ class MooncakeStore(HiCacheStorage):
host_indices: torch.Tensor, host_indices: torch.Tensor,
extra_info: Optional[HiCacheStorageExtraInfo] = None, extra_info: Optional[HiCacheStorageExtraInfo] = None,
) -> List[bool]: ) -> List[bool]:
# Apply extra_backend_tag prefix if available
if self.extra_backend_tag is not None:
prefix = self.extra_backend_tag
keys = [f"{prefix}_{key}" for key in keys]
key_strs, buffer_ptrs, buffer_sizes = self._batch_preprocess(keys, host_indices) key_strs, buffer_ptrs, buffer_sizes = self._batch_preprocess(keys, host_indices)
exist_result = self._batch_exist(key_strs) exist_result = self._batch_exist(key_strs)