From 79c82c5c4288b195ad32bafbd78e7e09c5e2714d Mon Sep 17 00:00:00 2001
From: Vladislav Nosivskoy <vladnosiv@gmail.com>
Date: Wed, 8 Apr 2026 11:42:57 +0300
Subject: [PATCH] [HiCache] Fix write_backup return type when parent not backed
 up (#22185)

Signed-off-by: Vladislav Nosivskoy <vladnosiv@gmail.com>
Co-authored-by: hzh0425 <hzh0425@apache.org>
---
 .../srt/mem_cache/hi_mamba_radix_cache.py      | 12 ++++++++----
 python/sglang/srt/mem_cache/hiradix_cache.py   | 18 ++++++++++++------
 2 files changed, 20 insertions(+), 10 deletions(-)

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: