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 2157bcc55..10e63c7f8 100644 --- a/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py +++ b/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py @@ -276,10 +276,14 @@ 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 + def write_backup(self, node: TreeNode, write_back=False) -> int: + # Backup invariant (for write-through mode): backed-up nodes must form a + # contiguous prefix from root — no gaps. Skip if parent isn't backed + # up yet; + if not write_back and ( + node.parent != self.root_node and not node.parent.backuped + ): + return 0 # 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: diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index fd75f859d..47b91fdab 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/python/sglang/srt/mem_cache/hiradix_cache.py @@ -607,10 +607,14 @@ class HiRadixCache(RadixCache): logger.warning("Hierarchical cache storage backend is not enabled.") 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 + def write_backup(self, node: TreeNode, write_back=False) -> int: + # Backup invariant (for write-through mode): backed-up nodes must form a + # contiguous prefix from root — no gaps. Skip if parent isn't backed + # up yet; + if not write_back and ( + node.parent != self.root_node and not node.parent.backuped + ): + return 0 host_indices = self.cache_controller.write( device_indices=node.value, @@ -800,8 +804,10 @@ class HiRadixCache(RadixCache): if not x.backuped: if self.cache_controller.write_policy == "write_back": # write to host if the node is not backuped - num_evicted += self.write_backup(x, write_back=True) - write_back_nodes.append(x) + written = self.write_backup(x, write_back=True) + num_evicted += written + if written > 0: + write_back_nodes.append(x) else: num_evicted += self._evict_regular(x) else: