From 51b276de7493cc0503a9325176c7f670e1431e3d Mon Sep 17 00:00:00 2001 From: Zhangheng Date: Sun, 5 Apr 2026 23:19:50 +0800 Subject: [PATCH] [BugFix][RadixTree]: Fix backup invariant violation in Hi-MambaRadixTree (#22062) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: 晟海 Co-authored-by: linjianyu77@foxmail.com --- python/sglang/srt/mem_cache/hi_mamba_radix_cache.py | 4 ++++ python/sglang/srt/mem_cache/hiradix_cache.py | 6 ++++++ 2 files changed, 10 insertions(+) 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 a08b2a33c..2157bcc55 100644 --- a/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py +++ b/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py @@ -277,6 +277,10 @@ class HiMambaRadixCache(MambaRadixCache): super().reset() def write_backup(self, node: TreeNode, write_back=False): + # Backup invariant: parent must be backed up before child. + if node.parent != self.root_node and not node.parent.backuped: + return + # If mamba host slot already exists, refresh its LRU position. if node.mamba_value is not None and node.mamba_host_value is not None: if self.mamba_host_lru_list.in_list(node): diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index 3c1e97daa..fd75f859d 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/python/sglang/srt/mem_cache/hiradix_cache.py @@ -608,6 +608,10 @@ class HiRadixCache(RadixCache): return False def write_backup(self, node: TreeNode, write_back=False): + # Backup invariant: parent must be backed up before child. + if node.parent != self.root_node and not node.parent.backuped: + return + host_indices = self.cache_controller.write( device_indices=node.value, node_id=node.id, @@ -836,6 +840,8 @@ class HiRadixCache(RadixCache): def _evict_regular(self, node: TreeNode): # evict a node not initiated write to host -- emit BlockRemoved + assert len(node.children) == 0, f"non-leaf, {node.id=}" + self._record_remove_event(node) self.cache_controller.mem_pool_device_allocator.free(node.value) num_evicted = len(node.value)