From 42f18fe560886ac83caada75a4e38e7d7bedbb2d Mon Sep 17 00:00:00 2001
From: shuwenn <47200617+alphabetc1@users.noreply.github.com>
Date: Mon, 16 Mar 2026 14:49:31 +0800
Subject: [PATCH] [HiCache] fix: release write-through lock_ref during decode
 (#20049)

---
 python/sglang/srt/managers/scheduler.py             | 5 +++++
 python/sglang/srt/mem_cache/base_prefix_cache.py    | 8 ++++++++
 python/sglang/srt/mem_cache/hi_mamba_radix_cache.py | 3 +++
 python/sglang/srt/mem_cache/hiradix_cache.py        | 3 +++
 python/sglang/srt/mem_cache/session_aware_cache.py  | 3 +++
 5 files changed, 22 insertions(+)

diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py
index 5fdcb4074..1ed4fd07b 100644
--- a/python/sglang/srt/managers/scheduler.py
+++ b/python/sglang/srt/managers/scheduler.py
@@ -2362,6 +2362,11 @@ class Scheduler(
             batch.batch_is_full = False
             return batch
 
+        # Eagerly release lock_ref on completed write-through nodes so they
+        # become evictable, improving batch scheduling headroom.
+        if self.enable_hierarchical_cache:
+            self.tree_cache.flush_write_through_acks()
+
         # Check if decode out of memory
         if (kv_full_retract_flag := not batch.check_decode_mem()) or (
             TEST_RETRACT and self.forward_ct % TEST_RETRACT_INTERVAL == 0
diff --git a/python/sglang/srt/mem_cache/base_prefix_cache.py b/python/sglang/srt/mem_cache/base_prefix_cache.py
index 4342ba82f..f0df9e989 100644
--- a/python/sglang/srt/mem_cache/base_prefix_cache.py
+++ b/python/sglang/srt/mem_cache/base_prefix_cache.py
@@ -229,6 +229,14 @@ class BasePrefixCache(ABC, PrefixCacheTrait):
         """
         raise NotImplementedError()
 
+    def flush_write_through_acks(self) -> None:
+        """Release lock_ref on radix-tree nodes whose write-through has completed.
+
+        Lightweight operation that only processes finished write acks.
+        No-op for caches without hierarchical write-through support.
+        """
+        pass
+
     def check_hicache_events(self) -> Any:
         """
         Check HiCache related activities to update radix tree and synchronize across TP workers if needed
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 2f738508c..7ce972a85 100644
--- a/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py
+++ b/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py
@@ -343,6 +343,9 @@ class HiMambaRadixCache(MambaRadixCache):
     def ready_to_load_host_cache(self) -> int:
         return self.cache_controller.start_loading()
 
+    def flush_write_through_acks(self) -> None:
+        self.writing_check()
+
     def check_hicache_events(self):
         self.writing_check()
         self.loading_check()
diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py
index 6f9190399..d1e9cf1ae 100644
--- a/python/sglang/srt/mem_cache/hiradix_cache.py
+++ b/python/sglang/srt/mem_cache/hiradix_cache.py
@@ -1094,6 +1094,9 @@ class HiRadixCache(RadixCache):
         """
         return self.cache_controller.start_loading()
 
+    def flush_write_through_acks(self) -> None:
+        self.writing_check()
+
     def check_hicache_events(self):
         self.writing_check()
         self.loading_check()
diff --git a/python/sglang/srt/mem_cache/session_aware_cache.py b/python/sglang/srt/mem_cache/session_aware_cache.py
index 3405918bb..54664387c 100644
--- a/python/sglang/srt/mem_cache/session_aware_cache.py
+++ b/python/sglang/srt/mem_cache/session_aware_cache.py
@@ -325,6 +325,9 @@ class SessionAwareCache(BasePrefixCache):
     def ready_to_load_host_cache(self):
         return self.inner.ready_to_load_host_cache()
 
+    def flush_write_through_acks(self) -> None:
+        return self.inner.flush_write_through_acks()
+
     def check_hicache_events(self):
         return self.inner.check_hicache_events()