From e4378ff37f61891cfc70fefceb7336fc8acfe021 Mon Sep 17 00:00:00 2001 From: Ke Bao Date: Thu, 14 May 2026 13:18:12 +0800 Subject: [PATCH] [UnifiedRadixCache] Fix HiCache load back start node (#25088) --- python/sglang/srt/managers/schedule_batch.py | 4 + python/sglang/srt/managers/schedule_policy.py | 4 +- .../sglang/srt/mem_cache/base_prefix_cache.py | 13 +- python/sglang/srt/mem_cache/chunk_cache.py | 1 + .../srt/mem_cache/hi_mamba_radix_cache.py | 5 +- python/sglang/srt/mem_cache/hiradix_cache.py | 4 +- .../sglang/srt/mem_cache/mamba_radix_cache.py | 2 + python/sglang/srt/mem_cache/radix_cache.py | 2 + .../sglang/srt/mem_cache/radix_cache_cpp.py | 1 + .../storage/lmcache/lmc_radix_cache.py | 1 + .../sglang/srt/mem_cache/swa_radix_cache.py | 2 + .../full_component.py | 29 +++- .../mamba_component.py | 5 +- .../unified_cache_components/swa_component.py | 14 +- .../srt/mem_cache/unified_radix_cache.py | 80 ++++----- .../sglang/srt/session/streaming_session.py | 1 + .../unit/mem_cache/test_radix_force_miss.py | 3 + .../mem_cache/test_streaming_session_unit.py | 1 + .../test_unified_radix_cache_unittest.py | 160 +++++++++++++++++- 19 files changed, 268 insertions(+), 64 deletions(-) diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 75ed09458..ab8cbc78f 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -740,8 +740,10 @@ class Req(ReqDllmMixin): self.extend_input_len = 0 # The relative logprob_start_len in an extend batch self.extend_logprob_start_len = 0 + # TODO(ispobock): rename to last_device_node self.last_node: Any = None self.last_host_node: Any = None + self.best_match_node: Any = None self.host_hit_length = 0 # Tokens loaded from storage backend (L3) during prefetch for this request self.storage_hit_length = 0 @@ -1041,12 +1043,14 @@ class Req(ReqDllmMixin): self.prefix_indices, self.last_node, self.last_host_node, + self.best_match_node, self.host_hit_length, self.mamba_branching_seqlen, ) = ( match_result.device_indices, match_result.last_device_node, match_result.last_host_node, + match_result.best_match_node, match_result.host_hit_length, match_result.mamba_branching_seqlen, ) diff --git a/python/sglang/srt/managers/schedule_policy.py b/python/sglang/srt/managers/schedule_policy.py index 29b90038a..2ee02b2e7 100644 --- a/python/sglang/srt/managers/schedule_policy.py +++ b/python/sglang/srt/managers/schedule_policy.py @@ -105,11 +105,13 @@ def match_prefix_for_req( req.prefix_indices, req.last_node, req.last_host_node, + req.best_match_node, req.host_hit_length, ) = ( match_result.device_indices, match_result.last_device_node, match_result.last_host_node, + match_result.best_match_node, match_result.host_hit_length, ) if match_result.mamba_branching_seqlen is not None: @@ -877,7 +879,7 @@ class PrefillAdder: if req.host_hit_length > 0: new_indices, req.last_node = self.tree_cache.init_load_back( InitLoadBackParams( - last_host_node=req.last_host_node, + best_match_node=req.best_match_node, host_hit_length=req.host_hit_length, req=req, ) diff --git a/python/sglang/srt/mem_cache/base_prefix_cache.py b/python/sglang/srt/mem_cache/base_prefix_cache.py index ae9df719e..802cf8c66 100644 --- a/python/sglang/srt/mem_cache/base_prefix_cache.py +++ b/python/sglang/srt/mem_cache/base_prefix_cache.py @@ -134,9 +134,9 @@ class DecLockRefResult: @dataclasses.dataclass class InitLoadBackParams: - """Unified parameters for init_load_back across different cache types""" + """Unified parameters for init_load_back across different cache types.""" - last_host_node: Any + best_match_node: Any host_hit_length: int mem_quota: Optional[int] = None req: Optional[Req] = None @@ -151,6 +151,13 @@ class MatchResult(NamedTuple): last_host_node : The last TreeNode on the host that was matched. Note that if HiCache is not enabled, this **must** be the same as `last_device_node`. + Reserved for L3 storage prefetch anchoring; L2 load_back + uses `best_match_node` instead. + best_match_node : Deepest node accepted by all component validators + during match_prefix. Anchor for every L2 host->device + load_back walk (FULL / SWA / ...). For legacy caches + that don't run multi-component validation, set this + equal to `last_host_node`. host_hit_length : Length of the host cache hit. For pure-KV caches this is the number of evicted KV tokens on CPU. For hybrid Mamba models this is max(kv_host_tokens, 1-if-mamba-on-host) so that a mamba-only @@ -164,6 +171,7 @@ class MatchResult(NamedTuple): device_indices: torch.Tensor last_device_node: Any last_host_node: Any + best_match_node: Any host_hit_length: int = 0 mamba_branching_seqlen: Optional[int] = None cache_protected_len: Optional[int] = None @@ -180,6 +188,7 @@ def zero_match_result(tree_cache, match_result: "MatchResult") -> "MatchResult": device_indices=match_result.device_indices[:0], last_device_node=root, last_host_node=root, + best_match_node=root, host_hit_length=0, ) diff --git a/python/sglang/srt/mem_cache/chunk_cache.py b/python/sglang/srt/mem_cache/chunk_cache.py index 6d34a3aa1..350a2ec8b 100644 --- a/python/sglang/srt/mem_cache/chunk_cache.py +++ b/python/sglang/srt/mem_cache/chunk_cache.py @@ -69,6 +69,7 @@ class ChunkCache(BasePrefixCache): device_indices=torch.empty((0,), dtype=torch.int64), last_device_node=None, last_host_node=None, + best_match_node=None, ) def insert(self, params: InsertParams) -> InsertResult: diff --git a/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py b/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py index d95cf3791..0475b3515 100644 --- a/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py +++ b/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py @@ -345,7 +345,7 @@ class HiMambaRadixCache(MambaRadixCache): self, params: InitLoadBackParams, ): - last_node = params.last_host_node + last_node = params.best_match_node mem_quota = params.mem_quota req = params.req if last_node.evicted or (last_node.mamba_evicted and last_node.mamba_backuped): @@ -932,6 +932,7 @@ class HiMambaRadixCache(MambaRadixCache): device_indices=torch.empty((0,), dtype=torch.int64, device=self.device), last_device_node=self.root_node, last_host_node=self.root_node, + best_match_node=self.root_node, host_hit_length=0, ) @@ -1066,6 +1067,8 @@ class HiMambaRadixCache(MambaRadixCache): device_indices=value, last_device_node=last_device_node, last_host_node=last_host_node, + # TODO(ispobock): use best_match_node as start node for load_back + best_match_node=last_host_node, host_hit_length=host_hit_length, mamba_branching_seqlen=mamba_branching_seqlen, ) diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index 4a75e965a..56e1d52ba 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/python/sglang/srt/mem_cache/hiradix_cache.py @@ -1046,7 +1046,7 @@ class HiRadixCache(RadixCache): self, params: InitLoadBackParams, ): - last_node = params.last_host_node + last_node = params.best_match_node mem_quota = params.mem_quota if last_node.evicted: loading_values = self.load_back(last_node, mem_quota) @@ -1251,6 +1251,8 @@ class HiRadixCache(RadixCache): device_indices=value, last_device_node=last_node, last_host_node=last_host_node, + # TODO(ispobock): use best_match_node as start node for load_back + best_match_node=last_host_node, host_hit_length=host_hit_length, ) diff --git a/python/sglang/srt/mem_cache/mamba_radix_cache.py b/python/sglang/srt/mem_cache/mamba_radix_cache.py index 55ee7983e..2462af688 100644 --- a/python/sglang/srt/mem_cache/mamba_radix_cache.py +++ b/python/sglang/srt/mem_cache/mamba_radix_cache.py @@ -488,6 +488,7 @@ class MambaRadixCache(KVCacheEventMixin, BasePrefixCache): ), last_device_node=self.root_node, last_host_node=self.root_node, + best_match_node=self.root_node, ) value, last_node, best_value_len = self._match_prefix_helper(key) @@ -1072,6 +1073,7 @@ class MambaRadixCache(KVCacheEventMixin, BasePrefixCache): device_indices=value, last_device_node=last_node, last_host_node=last_node, + best_match_node=last_node, mamba_branching_seqlen=mamba_branching_seqlen, ) diff --git a/python/sglang/srt/mem_cache/radix_cache.py b/python/sglang/srt/mem_cache/radix_cache.py index 8a24c5e15..5f8a256f6 100644 --- a/python/sglang/srt/mem_cache/radix_cache.py +++ b/python/sglang/srt/mem_cache/radix_cache.py @@ -353,6 +353,7 @@ class RadixCache(KVCacheEventMixin, BasePrefixCache): ), last_device_node=self.root_node, last_host_node=self.root_node, + best_match_node=self.root_node, ) self._record_all_cleared_event() @@ -413,6 +414,7 @@ class RadixCache(KVCacheEventMixin, BasePrefixCache): device_indices=value, last_device_node=last_node, last_host_node=last_node, + best_match_node=last_node, ) def insert(self, params: InsertParams) -> InsertResult: diff --git a/python/sglang/srt/mem_cache/radix_cache_cpp.py b/python/sglang/srt/mem_cache/radix_cache_cpp.py index 66f9fad96..97bf02835 100644 --- a/python/sglang/srt/mem_cache/radix_cache_cpp.py +++ b/python/sglang/srt/mem_cache/radix_cache_cpp.py @@ -107,6 +107,7 @@ class RadixCacheCpp(BasePrefixCache): device_indices=self._merge_tensor(device_indices_vec), last_device_node=node_gpu, last_host_node=node_cpu, + best_match_node=node_cpu, host_hit_length=host_indices_length, ) diff --git a/python/sglang/srt/mem_cache/storage/lmcache/lmc_radix_cache.py b/python/sglang/srt/mem_cache/storage/lmcache/lmc_radix_cache.py index da0d4cf17..1df5ec639 100644 --- a/python/sglang/srt/mem_cache/storage/lmcache/lmc_radix_cache.py +++ b/python/sglang/srt/mem_cache/storage/lmcache/lmc_radix_cache.py @@ -207,6 +207,7 @@ class LMCRadixCache(RadixCache): device_indices=value, last_device_node=last_node, last_host_node=last_node, + best_match_node=last_node, ) return base_res diff --git a/python/sglang/srt/mem_cache/swa_radix_cache.py b/python/sglang/srt/mem_cache/swa_radix_cache.py index af2d99e96..4aa5edc0e 100644 --- a/python/sglang/srt/mem_cache/swa_radix_cache.py +++ b/python/sglang/srt/mem_cache/swa_radix_cache.py @@ -405,6 +405,7 @@ class SWARadixCache(KVCacheEventMixin, BasePrefixCache): ), last_device_node=self.root_node, last_host_node=self.root_node, + best_match_node=self.root_node, ) value, last_node, best_value_len = self._match_prefix_helper(key) @@ -960,6 +961,7 @@ class SWARadixCache(KVCacheEventMixin, BasePrefixCache): device_indices=value, last_device_node=last_node, last_host_node=last_node, + best_match_node=last_node, ) def _compact_single_child_chain(self, node: TreeNode) -> None: diff --git a/python/sglang/srt/mem_cache/unified_cache_components/full_component.py b/python/sglang/srt/mem_cache/unified_cache_components/full_component.py index 805009dbb..3263cfbd1 100644 --- a/python/sglang/srt/mem_cache/unified_cache_components/full_component.py +++ b/python/sglang/srt/mem_cache/unified_cache_components/full_component.py @@ -152,12 +152,20 @@ class FullComponent(TreeComponent): ) -> IncLockRefResult: ct = self.component_type root = self.cache.root_node - delta = 0 cur = node - while cur != root: - cd = cur.component_data[ct] - assert cd.value is not None + # Skip the bottom evicted segment + while cur is not root and cur.component_data[ct].value is None: + result.skip_lock_node_ids.setdefault(ct, set()).add(cur.id) + cur = cur.parent + + # Lock the device-on segment up to root + delta = 0 + while cur is not root: + cd = cur.component_data[ct] + assert ( + cd.value is not None + ), f"FULL invariant broken: evicted ancestor {cur.id} above device-on segment" if cd.lock_ref == 0: key_len = len(cd.value) self.cache.component_evictable_size_[ct] -= key_len @@ -174,8 +182,12 @@ class FullComponent(TreeComponent): ) -> None: ct = self.component_type root = self.cache.root_node + skip_lock_node_ids = params.skip_lock_node_ids.get(ct, ()) if params else () cur = node while cur != root: + if cur.id in skip_lock_node_ids: + cur = cur.parent + continue cd = cur.component_data[ct] assert cd.value is not None assert cd.lock_ref > 0 @@ -203,15 +215,16 @@ class FullComponent(TreeComponent): return None if phase == CacheTransferPhase.LOAD_BACK: - # Walk evicted chain, collect host_values and nodes + # `node` is best_match_node. FULL device evict only from leaves, + # so once we hit a device-on node, everything above is also device-on backed_up: list[torch.Tensor] = [] nodes: list = [] cur = node while cur.evicted: cd = cur.component_data[ct] - if cd.host_value is not None: - backed_up.append(cd.host_value) - nodes.append(cur) + assert cd.host_value is not None + backed_up.append(cd.host_value) + nodes.append(cur) cur = cur.parent backed_up.reverse() nodes.reverse() diff --git a/python/sglang/srt/mem_cache/unified_cache_components/mamba_component.py b/python/sglang/srt/mem_cache/unified_cache_components/mamba_component.py index 45397c656..c1ef99b88 100644 --- a/python/sglang/srt/mem_cache/unified_cache_components/mamba_component.py +++ b/python/sglang/srt/mem_cache/unified_cache_components/mamba_component.py @@ -67,7 +67,7 @@ class MambaComponent(TreeComponent): ) -> MatchResult: cow_mamba = params.cow_mamba req = params.req - last_node = result.last_device_node + last_node = result.best_match_node if len(value_chunks) > best_value_len: chunk_size = get_global_server_args().mamba_cache_chunk_size @@ -101,8 +101,7 @@ class MambaComponent(TreeComponent): # HiCache: if mamba was evicted from device but has host backup, # ensure host_hit_length >= 1 so load_back is triggered. - host_node = result.last_host_node - cd = host_node.component_data[self.component_type] + cd = last_node.component_data[self.component_type] if cd.value is None and cd.host_value is not None: result = result._replace(host_hit_length=max(result.host_hit_length, 1)) diff --git a/python/sglang/srt/mem_cache/unified_cache_components/swa_component.py b/python/sglang/srt/mem_cache/unified_cache_components/swa_component.py index 8c27e0c94..8a0d6bac2 100644 --- a/python/sglang/srt/mem_cache/unified_cache_components/swa_component.py +++ b/python/sglang/srt/mem_cache/unified_cache_components/swa_component.py @@ -95,11 +95,12 @@ class SWAComponent(TreeComponent): ) -> MatchResult: ct = self.component_type n_swa = 0 - node = result.last_host_node + node = result.best_match_node root = self.cache.root_node while node is not root and n_swa < self.sliding_window_size: cd = node.component_data[ct] if cd.value is None and cd.host_value is not None: + # TODO(ispobock): refactor host_hit_length usage return result._replace(host_hit_length=max(result.host_hit_length, 1)) if cd.value is not None: n_swa += len(cd.value) @@ -440,11 +441,14 @@ class SWAComponent(TreeComponent): ] if phase == CacheTransferPhase.LOAD_BACK: + # `node` is best_match_node; the SWA validator guarantees every + # ancestor within `sliding_window_size` has value or host_value. n_swa = 0 backed_up: list[torch.Tensor] = [] nodes: list = [] - while node is not self.cache.root_node and n_swa < self.sliding_window_size: - cd = node.component_data[ct] + cur = node + while cur is not self.cache.root_node and n_swa < self.sliding_window_size: + cd = cur.component_data[ct] assert cd.host_value is not None or cd.value is not None if cd.value is not None: # device exists, skip it @@ -452,9 +456,9 @@ class SWAComponent(TreeComponent): else: # host only, collect it backed_up.append(cd.host_value) - nodes.append(node) + nodes.append(cur) n_swa += len(cd.host_value) - node = node.parent + cur = cur.parent if not backed_up: return None diff --git a/python/sglang/srt/mem_cache/unified_radix_cache.py b/python/sglang/srt/mem_cache/unified_radix_cache.py index 80a3da5bb..5b264307a 100644 --- a/python/sglang/srt/mem_cache/unified_radix_cache.py +++ b/python/sglang/srt/mem_cache/unified_radix_cache.py @@ -290,6 +290,7 @@ class UnifiedRadixCache(BasePrefixCache): ), last_device_node=self.root_node, last_host_node=self.root_node, + best_match_node=self.root_node, ) def init_hicache(self, server_args: ServerArgs, params: CacheInitParams) -> None: @@ -352,8 +353,10 @@ class UnifiedRadixCache(BasePrefixCache): if len(key) == 0: return self._empty_match_result - value, last_node, best_value_len = self._match_prefix_helper(key) - return self._match_post_processor(params, value, last_node, best_value_len) + value, best_match_node, best_value_len = self._match_prefix_helper(key) + return self._match_post_processor( + params, value, best_match_node, best_value_len + ) def insert(self, params: InsertParams) -> InsertResult: if self.disable: @@ -595,16 +598,16 @@ class UnifiedRadixCache(BasePrefixCache): child_key = key.child_key(self.page_size) value: list[torch.Tensor] = [] best_value_len = 0 - best_node = node + best_match_node = node validators = tuple( comp.create_match_validator() for comp in self._components_tuple ) def _update_best_if_valid(node): - nonlocal best_value_len, best_node + nonlocal best_value_len, best_match_node if all(v(node) for v in validators): best_value_len = len(value) - best_node = node + best_match_node = node while len(key) > 0 and child_key in node.children: child = node.children[child_key] @@ -625,7 +628,7 @@ class UnifiedRadixCache(BasePrefixCache): key = key[prefix_len:] if len(key): child_key = key.child_key(self.page_size) - return value, best_node, best_value_len + return value, best_match_node, best_value_len def _match_prefix_helper( self, key: RadixKey @@ -634,16 +637,16 @@ class UnifiedRadixCache(BasePrefixCache): child_key = key.child_key(self.page_size) value: list[torch.Tensor] = [] best_value_len = 0 - best_node = node + best_match_node = node validators = tuple( comp.create_match_validator() for comp in self._components_tuple ) def _update_best_if_valid(node): - nonlocal best_value_len, best_node + nonlocal best_value_len, best_match_node if all(v(node) for v in validators): best_value_len = len(value) - best_node = node + best_match_node = node while len(key) > 0 and child_key in node.children: child = node.children[child_key] @@ -667,16 +670,16 @@ class UnifiedRadixCache(BasePrefixCache): key = key[prefix_len:] if len(key): child_key = key.child_key(self.page_size) - return value, best_node, best_value_len + return value, best_match_node, best_value_len def _match_post_processor( self, params: MatchPrefixParams, value: list[torch.Tensor], - last_node: UnifiedTreeNode, + best_match_node: UnifiedTreeNode, best_value_len: int, ) -> MatchResult: - node_update = last_node + node_update = best_match_node for comp in self._components_tuple: if comp.component_type == BASE_COMPONENT_TYPE: continue # Full uses last_access_time, not LRU @@ -691,12 +694,12 @@ class UnifiedRadixCache(BasePrefixCache): node_update = node_update.parent # Walk up to find last_device_node - last_device_node = last_node + last_device_node = best_match_node while last_device_node is not self.root_node and last_device_node.evicted: last_device_node = last_device_node.parent # Walk up to find last_host_node - last_host_node = last_node + last_host_node = best_match_node while last_host_node is not self.root_node and not last_host_node.backuped: last_host_node = last_host_node.parent @@ -708,6 +711,7 @@ class UnifiedRadixCache(BasePrefixCache): device_indices=device_indices, last_device_node=last_device_node, last_host_node=last_host_node, + best_match_node=best_match_node, host_hit_length=0, ) @@ -1213,7 +1217,7 @@ class UnifiedRadixCache(BasePrefixCache): def load_back( self, - node: UnifiedTreeNode, + best_match_node: UnifiedTreeNode, mem_quota: Optional[int] = None, req=None, ) -> Optional[torch.Tensor]: @@ -1222,15 +1226,12 @@ class UnifiedRadixCache(BasePrefixCache): return None # Build KV transfer - last_hit_node = node kv_xfer = self.components[BASE_COMPONENT_TYPE].build_hicache_transfers( - last_hit_node, CacheTransferPhase.LOAD_BACK + best_match_node, CacheTransferPhase.LOAD_BACK )[0] # Lock path & pre-evict if device pool is insufficient - nodes_to_load = kv_xfer.nodes_to_load - ancestor_node = nodes_to_load[0].parent if nodes_to_load else last_hit_node - result = self.inc_lock_ref(ancestor_node) + result = self.inc_lock_ref(best_match_node) ancestor_lock_params = result.to_dec_params() kv_tokens = len(kv_xfer.host_indices) @@ -1240,7 +1241,7 @@ class UnifiedRadixCache(BasePrefixCache): if comp.component_type == BASE_COMPONENT_TYPE: continue t = comp.build_hicache_transfers( - last_hit_node, CacheTransferPhase.LOAD_BACK, req=req + best_match_node, CacheTransferPhase.LOAD_BACK, req=req ) if t: comp_xfers[comp.component_type] = t @@ -1255,7 +1256,7 @@ class UnifiedRadixCache(BasePrefixCache): if (kv_tokens < self.load_back_threshold and not comp_xfers) or ( mem_quota is not None and kv_tokens > mem_quota + result.delta ): - self.dec_lock_ref(ancestor_node, ancestor_lock_params) + self.dec_lock_ref(best_match_node, ancestor_lock_params) return None avail = self.token_to_kv_pool_allocator.available_size() @@ -1263,7 +1264,7 @@ class UnifiedRadixCache(BasePrefixCache): needed = kv_tokens - avail result = self.evict(EvictParams(num_tokens=needed)) if result.num_tokens_evicted < needed: - self.dec_lock_ref(ancestor_node, ancestor_lock_params) + self.dec_lock_ref(best_match_node, ancestor_lock_params) return None # Load H→D @@ -1271,32 +1272,32 @@ class UnifiedRadixCache(BasePrefixCache): aux_xfers.extend(anchor_kv_shared_indices_xfers) device_indices = self.cache_controller.load( host_indices=kv_xfer.host_indices, - node_id=last_hit_node.id, + node_id=best_match_node.id, extra_pools=aux_xfers or None, ) - self.dec_lock_ref(ancestor_node, ancestor_lock_params) + self.dec_lock_ref(best_match_node, ancestor_lock_params) if device_indices is None: return None # Commit: each component gets only its own transfers kv_xfer.device_indices = device_indices self.components[BASE_COMPONENT_TYPE].commit_hicache_transfer( - last_hit_node, + best_match_node, CacheTransferPhase.LOAD_BACK, [kv_xfer], ) for ct, xfers in comp_xfers.items(): self.components[ct].commit_hicache_transfer( - last_hit_node, + best_match_node, CacheTransferPhase.LOAD_BACK, xfers, ) - self._update_evictable_leaf_sets(ancestor_node) - self.ongoing_load_back[last_hit_node.id] = ( - last_hit_node, - self.inc_lock_ref(last_hit_node).to_dec_params(), + self._update_evictable_leaf_sets(best_match_node) + self.ongoing_load_back[best_match_node.id] = ( + best_match_node, + self.inc_lock_ref(best_match_node).to_dec_params(), ) return device_indices @@ -1384,27 +1385,28 @@ class UnifiedRadixCache(BasePrefixCache): ) -> tuple[torch.Tensor, UnifiedTreeNode]: """Prepare KV cache loading from host to device. Returns (device_indices, last_node) tuple.""" - last_node = params.last_host_node + best_match_node = params.best_match_node mem_quota = params.mem_quota req = params.req - if last_node.evicted or params.host_hit_length > 0: - loading_values = self.load_back(last_node, mem_quota, req=req) + if best_match_node.evicted or params.host_hit_length > 0: + loading_values = self.load_back(best_match_node, mem_quota, req=req) if loading_values is not None: logger.debug( "init_load_back success: loaded %d tokens for node %d", len(loading_values), - last_node.id, + best_match_node.id, ) - return loading_values, last_node + return loading_values, best_match_node # Fallback: walk up to non-evicted ancestor - while last_node is not self.root_node and last_node.evicted: - last_node = last_node.parent + # TODO(ispobock): The fallback path is not correct. The last_device_node should consider all the components. + while best_match_node is not self.root_node and best_match_node.evicted: + best_match_node = best_match_node.parent return ( self._empty_match_result.device_indices, - last_node, + best_match_node, ) def check_hicache_events(self) -> None: diff --git a/python/sglang/srt/session/streaming_session.py b/python/sglang/srt/session/streaming_session.py index a60b3376c..17602c3b3 100644 --- a/python/sglang/srt/session/streaming_session.py +++ b/python/sglang/srt/session/streaming_session.py @@ -259,6 +259,7 @@ class StreamingSession(BasePrefixCache): device_indices=device_indices, last_device_node=slot.virtual_node, last_host_node=slot.virtual_node, + best_match_node=slot.virtual_node, cache_protected_len=slot.cache_protected_len, ) diff --git a/test/registered/unit/mem_cache/test_radix_force_miss.py b/test/registered/unit/mem_cache/test_radix_force_miss.py index df8046cdd..24d5f345b 100644 --- a/test/registered/unit/mem_cache/test_radix_force_miss.py +++ b/test/registered/unit/mem_cache/test_radix_force_miss.py @@ -33,6 +33,7 @@ class _StubReq: self.prefix_indices = None self.last_node = None self.last_host_node = None + self.best_match_node = None self.host_hit_length = None self.mamba_branching_seqlen = None self.cache_protected_len = None @@ -50,6 +51,7 @@ class TestZeroMatchResult(unittest.TestCase): self.assertEqual(int(zeroed.device_indices.numel()), 0) self.assertIs(zeroed.last_device_node, tree.root_node) self.assertIs(zeroed.last_host_node, tree.root_node) + self.assertIs(zeroed.best_match_node, tree.root_node) self.assertEqual(zeroed.host_hit_length, 0) # dtype/device preserved (slice-not-allocate). self.assertEqual(zeroed.device_indices.dtype, match.device_indices.dtype) @@ -64,6 +66,7 @@ class TestZeroMatchResult(unittest.TestCase): device_indices=torch.empty((0,), dtype=torch.int64), last_device_node=None, last_host_node=None, + best_match_node=None, host_hit_length=0, ) self.assertIs(zero_match_result(_StubChunkCache(), original), original) diff --git a/test/registered/unit/mem_cache/test_streaming_session_unit.py b/test/registered/unit/mem_cache/test_streaming_session_unit.py index 7960470bb..f833872fb 100644 --- a/test/registered/unit/mem_cache/test_streaming_session_unit.py +++ b/test/registered/unit/mem_cache/test_streaming_session_unit.py @@ -104,6 +104,7 @@ def test_preabort_detaches_session_and_preserves_slot(): device_indices=torch.tensor([], dtype=torch.int64), last_device_node=None, last_host_node=None, + best_match_node=None, ) ], ) 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 415ce22e5..a96c8a82a 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 @@ -922,19 +922,19 @@ class UnifiedRadixCacheSuite: self.assertEqual(node_count_before, 2) tree._match_prefix_helper(RadixKey([1, 2])) - value, best_node, best_value_len = tree._match_prefix_helper( + value, best_match_node, best_value_len = tree._match_prefix_helper( RadixKey([1, 2, 3, 4]) ) self.assertEqual(best_value_len, 2) - self.assertEqual(best_node.key.token_ids, [3, 4]) + self.assertEqual(best_match_node.key.token_ids, [3, 4]) node_count_after_regular = count_nodes(tree.root_node) self.assertEqual(node_count_after_regular, node_count_before + 2) - value, best_node, best_value_len = tree._match_prefix_helper_readonly( + value, best_match_node, best_value_len = tree._match_prefix_helper_readonly( RadixKey([1, 2, 3]) ) self.assertEqual(best_value_len, 1) - self.assertEqual(best_node.key.token_ids, [1, 2]) + self.assertEqual(best_match_node.key.token_ids, [1, 2]) node_count_after_readonly = count_nodes(tree.root_node) self.assertEqual(node_count_after_readonly, node_count_after_regular) @@ -1821,6 +1821,7 @@ class UnifiedRadixCacheSuite: ), last_device_node=leaf, last_host_node=leaf, + best_match_node=leaf, host_hit_length=0, ) result = swa_comp.finalize_match_result( @@ -1905,6 +1906,102 @@ class UnifiedRadixCacheSuite: n_swa, ) + def _swa_anchor_chain_tokens(self, num_pages: int) -> list[int]: + """Reproduce the token sequence used by _build_chain_pages.""" + tokens: list[int] = [] + for i in range(num_pages): + tokens.extend(self._make_seq(1000 * (i + 1), 1)) + return tokens + + def _swa_anchor_setup(self): + """Chain layout (root-to-leaf): + + ... top padding ... + chain[-(window_pages + 1)] = N : SWA tombstone + chain[-(window_pages)] = Y : SWA host-only, FULL host-backed + chain[-(window_pages - 1)..-2] : SWA device, FULL.host=None + chain[-1] = X : SWA device, FULL.host=None + + Anchor=X has window_pages of device pages below N, so the SWA + load_back walker accumulates n_swa >= sliding_window_size and + exits before reaching N. Anchor=Y has only 1 page and walks into + N's tombstone. + """ + if not self.cfg.has_swa: + self.skipTest("requires SWA") + if self.cfg.has_mamba: + self.skipTest("SWA-only path keeps the chain construction simple") + ps = self.cfg.page_size + sw = self.cfg.sliding_window_size + if ps >= sw: + self.skipTest("test scenario requires ps < sw") + window_pages = (sw + ps - 1) // ps + chain_pages = window_pages + 3 + if chain_pages * ps > self.cfg.kv_size // 2: + self.skipTest("kv_size too small for the desired chain") + + tree, allocator, req_to_token_pool = build_fixture(self.cfg) + chain = self._build_chain_pages(tree, allocator, req_to_token_pool, chain_pages) + if len(chain) < chain_pages: + self.skipTest("chain too short") + self._simulate_backup_tree(tree) + + x = chain[-1] + y = chain[-window_pages] + n = chain[-(window_pages + 1)] + + n.component_data[ComponentType.SWA].value = None + n.component_data[ComponentType.SWA].host_value = None + y.component_data[ComponentType.SWA].value = None + # Strip FULL.host on X + intermediates so last_host_node walks past + # them to Y. Y.FULL untouched preserves the leaf-up evict invariant. + for node in chain[-(window_pages - 1) :]: + node.component_data[ComponentType.FULL].host_value = None + + tokens = self._swa_anchor_chain_tokens(len(chain)) + return tree, chain, n, y, x, tokens + + def test_hicache_swa_match_prefix_picks_best_match_node_above_last_host(self): + tree, _, _, y, x, tokens = self._swa_anchor_setup() + result = tree.match_prefix(MatchPrefixParams(key=RadixKey(tokens))) + self.assertIs(result.best_match_node, x) + self.assertIs(result.last_device_node, x) + self.assertIs(result.last_host_node, y) + + def test_hicache_swa_load_back_anchored_on_best_match_node(self): + tree, _, _, y, x, _ = self._swa_anchor_setup() + ps = self.cfg.page_size + swa_comp = tree.components[ComponentType.SWA] + + transfers = swa_comp.build_hicache_transfers(x, CacheTransferPhase.LOAD_BACK) + self.assertEqual(len(transfers), 1) + xfer = transfers[0] + self.assertEqual(xfer.name, PoolName.SWA) + self.assertEqual(xfer.nodes_to_load, [y]) + self.assertEqual(int(xfer.host_indices.numel()), ps) + + with self.assertRaises(AssertionError): + swa_comp.build_hicache_transfers(y, CacheTransferPhase.LOAD_BACK) + + def test_hicache_swa_finalize_anchored_on_best_match_node(self): + tree, _, _, y, x, _ = self._swa_anchor_setup() + swa_comp = tree.components[ComponentType.SWA] + + base = MatchResult( + device_indices=torch.empty((0,), dtype=torch.int64, device=tree.device), + last_device_node=x, + last_host_node=y, + best_match_node=x, + host_hit_length=0, + ) + result = swa_comp.finalize_match_result( + result=base, + params=MatchPrefixParams(key=RadixKey(self._make_seq(1, 1))), + value_chunks=[], + best_value_len=0, + ) + self.assertEqual(result.host_hit_length, 1) + def test_hicache_swa_temp_lock_does_not_release_restored_tombstone(self): """A temporary scheduler lock that skipped a SWA tombstone must not release later load-back/request locks after the tombstone is restored. @@ -1950,6 +2047,61 @@ class UnifiedRadixCacheSuite: tree.dec_lock_ref(leaf, request_lock.to_dec_params()) self.assertEqual(cd.lock_ref, 0) + def test_hicache_full_temp_lock_skips_evicted_anchor_and_mirrors_on_release( + self, + ): + """Acquire records the evicted anchor in skip_lock_node_ids (phase 1) + and locks device-on ancestors only (phase 2). After load_back + restores the anchor, a second acquire covers it; releasing the + first must mirror the skip so the anchor's lock_ref is not + decremented twice. + """ + if self._skip_unsupported_hicache_test(): + return + ps = self.cfg.page_size + if 3 * ps > self.cfg.kv_size // 2: + self.skipTest("kv_size too small") + tree, allocator, req_to_token_pool = self._build_hicache_fixture() + chain = self._build_chain_pages(tree, allocator, req_to_token_pool, 3) + if len(chain) < 3: + self.skipTest("chain too short") + a, y, anchor = chain + self._simulate_backup_tree(tree) + + cd_anchor = anchor.component_data[ComponentType.FULL] + cd_a = a.component_data[ComponentType.FULL] + cd_y = y.component_data[ComponentType.FULL] + anchor_value = cd_anchor.value + cd_anchor.value = None + + self.assertEqual(cd_anchor.lock_ref, 0) + self.assertEqual(cd_y.lock_ref, 0) + self.assertEqual(cd_a.lock_ref, 0) + + temp_lock = tree.inc_lock_ref(anchor) + self.assertEqual(cd_anchor.lock_ref, 0) + self.assertEqual(cd_y.lock_ref, 1) + self.assertEqual(cd_a.lock_ref, 1) + self.assertIn(ComponentType.FULL, temp_lock.skip_lock_node_ids) + self.assertIn(anchor.id, temp_lock.skip_lock_node_ids[ComponentType.FULL]) + + cd_anchor.value = anchor_value + + second_lock = tree.inc_lock_ref(anchor) + self.assertEqual(cd_anchor.lock_ref, 1) + self.assertEqual(cd_y.lock_ref, 2) + self.assertEqual(cd_a.lock_ref, 2) + + tree.dec_lock_ref(anchor, temp_lock.to_dec_params()) + self.assertEqual(cd_anchor.lock_ref, 1) + self.assertEqual(cd_y.lock_ref, 1) + self.assertEqual(cd_a.lock_ref, 1) + + tree.dec_lock_ref(anchor, second_lock.to_dec_params()) + self.assertEqual(cd_anchor.lock_ref, 0) + self.assertEqual(cd_y.lock_ref, 0) + self.assertEqual(cd_a.lock_ref, 0) + def test_hicache_mamba_temp_lock_does_not_release_restored_tombstone(self): """A temporary scheduler lock that skipped a Mamba tombstone must not release later load-back/request locks after the tombstone is restored.