diff --git a/python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py b/python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py index 245e0eb16..dd42e2cac 100644 --- a/python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py +++ b/python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py @@ -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: diff --git a/test/registered/unit/mem_cache/test_mooncake_group_semantics.py b/test/registered/unit/mem_cache/test_mooncake_group_semantics.py index 456090c7a..394f78036 100644 --- a/test/registered/unit/mem_cache/test_mooncake_group_semantics.py +++ b/test/registered/unit/mem_cache/test_mooncake_group_semantics.py @@ -200,6 +200,7 @@ def _make_config( *, enable_group_semantics=True, extra_backend_tag=None, + model_name=None, is_mla_model=False, should_split_heads=False, tp_rank=0, @@ -225,7 +226,7 @@ def _make_config( is_mla_model=is_mla_model, enable_storage_metrics=False, is_page_first_layout=True, - model_name="test", + model_name=model_name, tp_lcm_size=tp_lcm_size, should_split_heads=should_split_heads, extra_config=extra_config, @@ -237,6 +238,7 @@ def _make_store( enable_group_semantics=True, replicate_config_cls=ReplicateConfigWithGroupIds, extra_backend_tag=None, + model_name=None, is_mla_model=False, should_split_heads=False, tp_rank=0, @@ -247,6 +249,7 @@ def _make_store( cfg = _make_config( enable_group_semantics=enable_group_semantics, extra_backend_tag=extra_backend_tag, + model_name=model_name, is_mla_model=is_mla_model, should_split_heads=should_split_heads, tp_rank=tp_rank, @@ -449,6 +452,25 @@ class TestMooncakeGroupSemantics(CustomTestCase): ["sglang-hicache:tag_page0", "sglang-hicache:tag_page1"], ) + def test_model_names_isolate_the_same_logical_key(self): + store_a, fake_store_a = _make_store( + enable_group_semantics=False, model_name="org/model-a" + ) + store_b, fake_store_b = _make_store( + enable_group_semantics=False, model_name="org/model-b" + ) + store_a.register_mem_pool_host(FakeHostKVCache(objects_per_page=2)) + store_b.register_mem_pool_host(FakeHostKVCache(objects_per_page=2)) + + self.assertEqual(store_a.batch_set_v1(["page0"], torch.tensor([0])), [True]) + self.assertEqual(store_b.batch_set_v1(["page0"], torch.tensor([0])), [True]) + + keys_a = fake_store_a.batch_put_calls[0]["keys"] + keys_b = fake_store_b.batch_put_calls[0]["keys"] + self.assertEqual(keys_a, ["org-model-a_page0_0_k", "org-model-a_page0_0_v"]) + self.assertEqual(keys_b, ["org-model-b_page0_0_k", "org-model-b_page0_0_v"]) + self.assertTrue(set(keys_a).isdisjoint(keys_b)) + if __name__ == "__main__": unittest.main(verbosity=3)