[HiCache][LoRA] Isolate storage pages by extra key (#38577)

Co-authored-by: Shuwen Wang <47200617+alphabetc1@users.noreply.github.com>

python/sglang/srt/disaggregation/decode_hicache_mixin.py

@@ -90,6 +90,8 @@ class DecodeHiCachePreallocMixin:
                     suffix_tokens,
                     last_hash,
                     prefix_keys,
+                    extra_key=req.extra_key,
+                    cache_salt=req.cache_salt,
                 )
 
         return DecodePrefixMatch(

python/sglang/srt/disaggregation/decode_kvcache_offload_manager.py

@@ -22,6 +22,7 @@ from sglang.srt.mem_cache.memory_pool import (
     MLATokenToKVPool,
     ReqToTokenPool,
 )
+from sglang.srt.mem_cache.utils import storage_namespace_seed
 from sglang.srt.runtime_context import (
     get_memory,
     get_schedule,
@@ -138,7 +139,9 @@ class DecodeKVCacheOffloadManager:
         state = self.offloaded_state.get(req)
         if state is None:
             prefill_hashes = self._compute_prefix_hash(
-                req.origin_input_ids[:prefill_offloaded_len]
+                req.origin_input_ids[:prefill_offloaded_len],
+                extra_key=req.extra_key,
+                cache_salt=req.cache_salt,
             )
             last_prefill_hash = (
                 prefill_hashes[-1] if prefill_offloaded_len > 0 else None
@@ -271,7 +274,12 @@ class DecodeKVCacheOffloadManager:
         self, req, host_indices, incremental_tokens, start_time, prior_hash
     ):
         """Trigger async backup from host to storage."""
-        page_hashes = self._compute_prefix_hash(incremental_tokens, prior_hash)
+        page_hashes = self._compute_prefix_hash(
+            incremental_tokens,
+            prior_hash,
+            extra_key=req.extra_key,
+            cache_salt=req.cache_salt,
+        )
         ack_id = self.cache_controller.write_storage(
             host_indices,
             incremental_tokens,
@@ -280,9 +288,12 @@ class DecodeKVCacheOffloadManager:
         self.ongoing_backup[ack_id] = (req.rid, host_indices, start_time)
         return page_hashes[-1] if len(page_hashes) > 0 else prior_hash
 
-    def _compute_prefix_hash(self, tokens, prior_hash=""):
+    def _compute_prefix_hash(
+        self, tokens, prior_hash="", extra_key=None, cache_salt=None
+    ):
+        """Match prefill storage hashes."""
         page_hashes = []
-        last_hash = prior_hash
+        last_hash = prior_hash or storage_namespace_seed(extra_key, cache_salt)
         for offset in range(0, len(tokens), self.page_size):
             page_tokens = tokens[offset : offset + self.page_size]
             last_hash = self.cache_controller.get_hash_str(page_tokens, last_hash)

python/sglang/srt/managers/cache_controller.py

@@ -42,6 +42,7 @@ from sglang.srt.layers.dp_attention import (
 )
 from sglang.srt.mem_cache.l2_transfer import L2Transfer, L2TransferEngine
 from sglang.srt.mem_cache.memory_pool import MLATokenToKVPool
+from sglang.srt.mem_cache.utils import get_storage_hash_str
 from sglang.srt.runtime_context import get_parallel
 from sglang.srt.utils import get_device_module
 
@@ -1197,7 +1198,7 @@ class HiCacheController:
 
         storage_query_count = 0
         hash_value = []
-        page_hashes = self.get_hash_str(
+        page_hashes = get_storage_hash_str(
             tokens_to_fetch, last_hash, page_size=self.page_size
         )
         operation.all_hash_values = page_hashes

python/sglang/srt/mem_cache/events.py

@@ -89,9 +89,9 @@ class KVCacheEventRecorder:
         """Hash values to publish for ``node``, computing them if not yet set."""
         if node.hash_value is None:
             node.hash_value = compute_node_hash_values(node, self.page_size)
-        if node.key.cache_salt is not None:
-            return compute_node_event_hash_values(node, self.page_size)
-        return node.hash_value
+        if node.key.extra_key is None and node.key.cache_salt is None:
+            return node.hash_value
+        return compute_node_event_hash_values(node, self.page_size)
 
     def _parent_block_hash(self, node: Any) -> Optional[int]:
         """The hash the first page of ``node`` links back to.
@@ -103,7 +103,7 @@ class KVCacheEventRecorder:
         parent = node.parent
         if parent is None or parent.parent is None:
             return None
-        if node.key.cache_salt is not None:
+        if node.key.extra_key is not None or node.key.cache_salt is not None:
             parent_hash_values = parent.event_hash_value
             assert parent_hash_values is not None
         else:

python/sglang/srt/mem_cache/hiradix_cache.py

@@ -1488,15 +1488,17 @@ class HiRadixCache(RadixCache):
         new_input_tokens: List[int],
         last_hash: Optional[str] = None,
         prefix_keys: Optional[List[str]] = None,
+        extra_key: Optional[str] = None,
+        cache_salt: Optional[str] = None,
     ) -> int:
         if not self.enable_storage or self.cache_controller.prefetch_rate_limited():
             return 0
 
         prefetch_key = RadixKey(
             new_input_tokens,
-            extra_key=last_host_node.key.extra_key,
+            extra_key=extra_key,
             is_bigram=self.is_eagle,
-            cache_salt=last_host_node.key.cache_salt,
+            cache_salt=cache_salt,
         ).page_aligned(self.page_size)
         if len(prefetch_key) < self.prefetch_threshold:
             return 0
@@ -1773,16 +1775,14 @@ class HiRadixCache(RadixCache):
         prefix_keys: Optional[List[str]] = None,
         # Scheduler-call parity with UnifiedRadixCache; unused in cache mode.
         matched_prefix_tokens: Optional[List[int]] = None,
-        # Cache mode write-through keeps the anchor on the request's own path,
-        # so the namespace is already carried by ``last_host_node.key``.
         extra_key: Optional[str] = None,
         cache_salt: Optional[str] = None,
     ):
         prefetch_key = RadixKey(
             new_input_tokens,
-            extra_key=last_host_node.key.extra_key,
+            extra_key=extra_key,
             is_bigram=self.is_eagle,
-            cache_salt=last_host_node.key.cache_salt,
+            cache_salt=cache_salt,
         )
         # align the number of fetching tokens to the page size
         prefetch_key = prefetch_key.page_aligned(self.page_size)

python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py

@@ -39,6 +39,8 @@ from sglang.srt.mem_cache.pool_host.mha import MHATokenToKVPoolHost
 if TYPE_CHECKING:
     from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator
 
+from sglang.srt.mem_cache.utils import get_storage_hash_str
+
 logger = logging.getLogger(__name__)
 
 
@@ -576,7 +578,7 @@ class HybridCacheController(BaseHiCacheController):
         return operation.id
 
     def _storage_hit_query(self, operation) -> tuple[list[str], int]:
-        hash_value = self.get_hash_str(
+        hash_value = get_storage_hash_str(
             operation.token_ids, operation.last_hash, page_size=self.page_size
         )
         operation.all_hash_values = hash_value

python/sglang/srt/mem_cache/radix_cache.py

@@ -71,12 +71,9 @@ class RadixKey:
     ):
         # token ids sequence (raw ints in both modes)
         self.token_ids = token_ids
-        # Extra key for caller-defined cache classification.
+        # Namespaces the tree and storage; omitted from KV events.
         self.extra_key = extra_key
-        # Cache salt is kept distinct so it cannot collide with extra_key.
-        # It namespaces the in-process radix tree and external KV events;
-        # external L3/remote storage keys remain token-only and are outside
-        # this contract.
+        # Namespaces the tree, storage and KV events.
         self.cache_salt = cache_salt or None
         # bigram view over token_ids: length = max(0, len(token_ids) - 1)
         self.is_bigram = is_bigram

python/sglang/srt/mem_cache/unified_cache/unified_cache_linker.py

@@ -42,7 +42,7 @@ from sglang.srt.mem_cache.unified_cache.components import (
     LinkerTransferPhase,
     TreeComponent,
 )
-from sglang.srt.mem_cache.utils import get_hash_str
+from sglang.srt.mem_cache.utils import get_storage_hash_str
 
 if TYPE_CHECKING:
     from sglang.srt.managers.schedule_batch import Req
@@ -279,7 +279,7 @@ class UnifiedCacheLinkerWrapper:
         tail_len = (len(key) - device_hit_len) // page * page
         if tail_len == 0:
             return []
-        return get_hash_str(
+        return get_storage_hash_str(
             key[device_hit_len : device_hit_len + tail_len],
             last_hash,
             page_size=page,

python/sglang/srt/mem_cache/unified_radix_cache.py

@@ -1829,8 +1829,10 @@ class UnifiedRadixCache(BasePrefixCache):
         new_input_tokens: list[int],
         last_hash: Optional[str] = None,
         prefix_keys: Optional[list[str]] = None,
+        extra_key: Optional[str] = None,
+        cache_salt: Optional[str] = None,
     ) -> int:
-        """Synchronously probe L3 storage for the reusable prefix length."""
+        """Probe L3 with the request namespace."""
         if (
             not self.enable_storage
             or self.cache_controller is None
@@ -1838,7 +1840,6 @@ class UnifiedRadixCache(BasePrefixCache):
         ):
             return 0
 
-        extra_key, cache_salt = self.tree_core.prefetch_anchor_info(last_host_node_id)
         prefetch_key = RadixKey(
             new_input_tokens,
             extra_key=extra_key,

python/sglang/srt/mem_cache/utils.py

@@ -120,6 +120,36 @@ def get_hash_str(
     return get_native_hash(token_ids, prior_digest, page_size)
 
 
+def storage_namespace_seed(
+    extra_key: Optional[str], cache_salt: Optional[str]
+) -> Optional[str]:
+    """Seed storage chains; preserve unnamespaced keys and Rust byte parity."""
+    if extra_key is None and cache_salt is None:
+        return None
+    digest = hashlib.sha256(b"sglang-cache-namespace-v1")
+    # Presence and UTF-8 byte length distinguish absent, empty and joined parts.
+    for part in (extra_key, cache_salt):
+        if part is None:
+            digest.update(b"\x00")
+            continue
+        encoded = part.encode("utf-8")
+        digest.update(b"\x01" + len(encoded).to_bytes(8, "little") + encoded)
+    return digest.hexdigest()
+
+
+def get_storage_hash_str(
+    key: Any,
+    prior_hash: Optional[str] = None,
+    page_size: Optional[int] = None,
+) -> str | List[str]:
+    """Seed new storage chains with the request namespace."""
+    if prior_hash is None:
+        prior_hash = storage_namespace_seed(
+            getattr(key, "extra_key", None), getattr(key, "cache_salt", None)
+        )
+    return get_hash_str(key, prior_hash, page_size=page_size)
+
+
 def hash_str_to_int64(hash_str: str) -> int:
     """Convert SHA256 hex string to signed 64-bit integer for events.
 
@@ -138,15 +168,15 @@ def compute_node_hash_values(node: Any, page_size: int) -> List[str]:
         if len(node.parent.key) > 0 and len(node.parent.hash_value) > 0:
             parent_hash = node.parent.hash_value[-1]
 
-    hash_values = get_hash_str(node.key, parent_hash, page_size=page_size)
+    hash_values = get_storage_hash_str(node.key, parent_hash, page_size=page_size)
     assert isinstance(hash_values, list)
     return hash_values
 
 
 def compute_node_event_hash_values(node: Any, page_size: int) -> List[str]:
-    """Compute and memoize namespace-aware external KV-event hashes."""
-    cache_salt = node.key.cache_salt
-    if cache_salt is None:
+    """Hash tokens with the legacy salt seed; omit extra_key."""
+    namespace = (node.key.extra_key, node.key.cache_salt)
+    if namespace == (None, None):
         return compute_node_hash_values(node, page_size)
 
     if node.event_hash_value is not None:
@@ -154,31 +184,22 @@ def compute_node_event_hash_values(node: Any, page_size: int) -> List[str]:
 
     missing_nodes = []
     current = node
-    while (
-        current is not None
-        and current.key is not None
-        and len(current.key) > 0
-        and current.event_hash_value is None
-    ):
-        if current.key.cache_salt != cache_salt:
-            raise ValueError("Radix path contains mismatched cache_salt values")
+    while current is not None and current.key is not None and len(current.key) > 0:
+        if (current.key.extra_key, current.key.cache_salt) != namespace:
+            raise ValueError("Radix path contains mismatched cache namespaces")
+        if current.event_hash_value is not None:
+            break
         missing_nodes.append(current)
         current = current.parent
 
-    if (
-        current is not None
-        and current.key is not None
-        and len(current.key) > 0
-        and current.key.cache_salt != cache_salt
-    ):
-        raise ValueError("Radix path contains mismatched cache_salt values")
-
     if current is not None and current.event_hash_value:
         parent_hash = current.event_hash_value[-1]
-    else:
+    elif node.key.cache_salt is not None:
         parent_hash = hashlib.sha256(
-            b"sglang-cache-salt-v1\0" + cache_salt.encode("utf-8")
+            b"sglang-cache-salt-v1\0" + node.key.cache_salt.encode("utf-8")
         ).hexdigest()
+    else:
+        parent_hash = None
 
     for missing_node in reversed(missing_nodes):
         hash_values = get_hash_str(missing_node.key, parent_hash, page_size=page_size)