diff --git a/python/sglang/srt/disaggregation/decode_hicache_mixin.py b/python/sglang/srt/disaggregation/decode_hicache_mixin.py index fbdc08d64..588056fa3 100644 --- a/python/sglang/srt/disaggregation/decode_hicache_mixin.py +++ b/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 diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 4250cbe93..d6f547328 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/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): diff --git a/python/sglang/srt/mem_cache/buffer_mode/pipeline.py b/python/sglang/srt/mem_cache/buffer_mode/pipeline.py index 91bcc9295..3986893fb 100644 --- a/python/sglang/srt/mem_cache/buffer_mode/pipeline.py +++ b/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 diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index 6f82bc810..292a980f0 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/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, diff --git a/python/sglang/srt/mem_cache/unified_radix_cache.py b/python/sglang/srt/mem_cache/unified_radix_cache.py index ed48bd724..72d891957 100644 --- a/python/sglang/srt/mem_cache/unified_radix_cache.py +++ b/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, diff --git a/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py b/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py index 6bf7a282d..fec8124de 100644 --- a/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py +++ b/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,