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

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:

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)