[HiCache] Key storage prefetch by the request namespace (#36382)

python/sglang/srt/disaggregation/decode_hicache_mixin.py

@@ -124,7 +124,13 @@ class DecodeHiCachePreallocMixin:
                 else None
             )
             self.tree_cache.prefetch_from_storage(
-                req.rid, prefix_match.last_host_node, suffix, last_hash, prefix_keys
+                req.rid,
+                prefix_match.last_host_node,
+                suffix,
+                last_hash,
+                prefix_keys,
+                extra_key=req.extra_key,
+                cache_salt=req.cache_salt,
             )
             prefix_match.prefetch_registered = (
                 req.rid in self.tree_cache.ongoing_prefetch

python/sglang/srt/managers/scheduler.py

@@ -2845,6 +2845,8 @@ class Scheduler(
                     tree_cache.get_last_hash_value(last_host_node),
                     prefix_keys,
                     matched_prefix_tokens=req.full_untruncated_fill_ids[:matched_len],
+                    extra_key=req.extra_key,
+                    cache_salt=req.cache_salt,
                 )
 
     def _retry_missed_storage_prefetches(self):

python/sglang/srt/mem_cache/buffer_mode/pipeline.py

@@ -99,6 +99,7 @@ class _StagedPrefetch(msgspec.Struct):
     req_id: str
     key_tokens: list[int]
     extra_key: Optional[str]
+    cache_salt: Optional[str]
     matched_len: int
     num_tokens: int
     occupied_tokens: int
@@ -738,6 +739,7 @@ class BufferModePipeline:
             req_id=req_id,
             key_tokens=prefix_tokens + list(prefetch_key[:num_tokens].token_ids),
             extra_key=prefetch_key.extra_key,
+            cache_salt=prefetch_key.cache_salt,
             matched_len=len(prefix_tokens),
             num_tokens=num_tokens,
             occupied_tokens=occupied_tokens,
@@ -794,6 +796,19 @@ class BufferModePipeline:
             cc.prefetch_tokens_occupied -= f.occupied_tokens
             return unchanged
 
+        # A hold staged under a different namespace than the consuming request
+        # must never splice (wrong-namespace publish = duplicate slot
+        # ownership); unreachable while the prefetch key is request-derived.
+        if f.extra_key != req.extra_key or f.cache_salt != req.cache_salt:
+            logger.error(
+                "HiCache staged prefetch dropped req=%s reason=namespace "
+                "staged=%s req=%s",
+                req.rid,
+                (f.extra_key, f.cache_salt),
+                (req.extra_key, req.cache_salt),
+            )
+            return _drop()
+
         # Splice-validity: the span only fits if the device prefix still
         # ends exactly at the enqueue-time matched_len.
         if len(req.prefix_indices) != f.matched_len:
@@ -813,6 +828,7 @@ class BufferModePipeline:
             array("q", f.key_tokens),
             extra_key=f.extra_key,
             is_bigram=cache.tree_core.is_eagle,
+            cache_salt=f.cache_salt,
         ).page_aligned(cache.page_size)
         span_end = f.matched_len + f.num_tokens
 

python/sglang/srt/mem_cache/hiradix_cache.py

@@ -1780,6 +1780,10 @@ 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,

python/sglang/srt/mem_cache/unified_radix_cache.py

@@ -1660,12 +1660,25 @@ class UnifiedRadixCache(BasePrefixCache):
         last_hash: Optional[str] = None,
         prefix_keys: Optional[list[str]] = None,
         matched_prefix_tokens: Optional[list[int]] = None,
+        extra_key: Optional[str] = None,
+        cache_salt: Optional[str] = None,
     ) -> None:
         if not self.enable_storage or self.cache_controller is None:
             return
 
         buffer_mode = self.host_memory_mode == "buffer_only"
-        extra_key, cache_salt = self.tree_core.prefetch_anchor_info(last_host_node_id)
+        # Key the span by the request's namespace, not the anchor's (a root
+        # anchor has none): a span published under the wrong namespace gets
+        # re-owned by the request's own insert (double free).
+        anchor_extra_key, anchor_cache_salt = self.tree_core.prefetch_anchor_info(
+            last_host_node_id
+        )
+        assert (anchor_extra_key is None or anchor_extra_key == extra_key) and (
+            anchor_cache_salt is None or anchor_cache_salt == cache_salt
+        ), (
+            f"prefetch anchor namespace {(anchor_extra_key, anchor_cache_salt)} "
+            f"!= request namespace {(extra_key, cache_salt)}"
+        )
         prefetch_key = RadixKey(
             new_input_tokens,
             extra_key=extra_key,

test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py

@@ -2708,6 +2708,8 @@ class UnifiedRadixCacheSuite:
             prefix_len = f.matched_len
         req = mock.Mock()
         req.rid = req_id
+        req.extra_key = None
+        req.cache_salt = None
         if prefix_indices is not None:
             # Spliceable mid-anchor consumption publishes value=cat(prefix,
             # fill) — the real device prefix is required (zeros would insert
@@ -3342,6 +3344,8 @@ class UnifiedRadixCacheSuite:
         held = cons.buffer_pipeline.staged_prefetches[req_id]
         req = mock.Mock()
         req.rid = req_id
+        req.extra_key = None
+        req.cache_salt = None
         req.last_node = cons.root_node_handle()
         req.prefix_indices = torch.zeros(
             held.matched_len,
@@ -3561,6 +3565,8 @@ class UnifiedRadixCacheSuite:
         held = cons.buffer_pipeline.staged_prefetches[req_id]
         req = mock.Mock()
         req.rid = req_id
+        req.extra_key = None
+        req.cache_salt = None
         req.last_node = cons.root_node_handle()
         req.prefix_indices = torch.zeros(
             held.matched_len,
@@ -3763,6 +3769,8 @@ class UnifiedRadixCacheSuite:
         f = cons.buffer_pipeline.staged_prefetches[req_id]
         req = mock.Mock()
         req.rid = req_id
+        req.extra_key = None
+        req.cache_salt = None
         req.prefix_indices = torch.zeros(
             0,
             dtype=torch.int64,