[HiCache] Merge HiCache event checks to reduce decode overhead (#30511)

Co-authored-by: alphabetc1 <47200617+alphabetc1@users.noreply.github.com>

python/sglang/srt/managers/scheduler.py

@@ -3226,11 +3226,6 @@ 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

python/sglang/srt/mem_cache/base_prefix_cache.py

@@ -364,14 +364,6 @@ 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

python/sglang/srt/mem_cache/hi_mamba_radix_cache.py

@@ -466,9 +466,6 @@ 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()

python/sglang/srt/mem_cache/hiradix_cache.py

@@ -202,7 +202,6 @@ class HiRadixCache(RadixCache):
             1 if server_args.hicache_write_policy == "write_through" else 2
         )
         self.load_back_threshold = 10
-
         # Detach storage backend automatically on process shutdown
         atexit.register(self.shutdown)
 
@@ -981,7 +980,42 @@ class HiRadixCache(RadixCache):
                 # write to host if the node is not backuped
                 self.write_backup(node)
 
-    def writing_check(self, write_back=False):
+    def _count_ready_acks(self, ack_queue) -> int:
+        ready_count = 0
+        for ack in ack_queue:
+            if not ack.finish_event.query():
+                break
+            ready_count += 1
+        return ready_count
+
+    def _sync_hicache_ready_counts(self) -> tuple[int, int, tuple[int, ...]]:
+        cache_controller = self.cache_controller
+        storage_queue_sizes = (
+            (
+                cache_controller.prefetch_revoke_queue.qsize(),
+                cache_controller.prefetch_hit_queue.qsize(),
+                cache_controller.ack_backup_queue.qsize(),
+                cache_controller.host_mem_release_queue.qsize(),
+            )
+            if self.enable_storage
+            else ()
+        )
+
+        ready_counts = torch.tensor(
+            [
+                self._count_ready_acks(cache_controller.ack_write_queue),
+                self._count_ready_acks(cache_controller.ack_load_queue),
+                *storage_queue_sizes,
+            ],
+            dtype=torch.int,
+            device="cpu",
+        )
+        self._all_reduce(ready_counts, torch.distributed.ReduceOp.MIN)
+
+        count_values = list(map(int, ready_counts.tolist()))
+        return count_values[0], count_values[1], tuple(count_values[2:])
+
+    def writing_check(self, write_back=False, finish_count: Optional[int] = None):
         if write_back:
             # blocking till all write back complete
             while len(self.ongoing_write_through) > 0:
@@ -993,17 +1027,19 @@ class HiRadixCache(RadixCache):
                 assert len(self.ongoing_write_through) == 0
             return
 
+        if finish_count is None:
             # Every rank must enter the all_reduce below; ongoing_write_through can
             # diverge across ranks (e.g. write_backup returning 0 on a subset under
             # host memory pressure), so a conditional skip desyncs the NCCL op
             # sequence and deadlocks under TP > 1. (Matches UnifiedRadixCache.)
             finish_count = 0
             if self.pp_rank == 0:
-            for ack in self.cache_controller.ack_write_queue:
+                finish_count = self._count_ready_acks(
-                if not ack.finish_event.query():
+                    self.cache_controller.ack_write_queue
-                    break
+                )
-                finish_count += 1
+            finish_count_tensor = torch.tensor(
-        finish_count_tensor = torch.tensor(finish_count, dtype=torch.int, device="cpu")
+                finish_count, dtype=torch.int, device="cpu"
+            )
             self._all_reduce(finish_count_tensor, torch.distributed.ReduceOp.MIN)
             finish_count = finish_count_tensor.item()
 
@@ -1016,14 +1052,16 @@ class HiRadixCache(RadixCache):
                 self._finish_write_through_ack(ack_id, release_lock=True)
             finish_count -= 1
 
-    def loading_check(self):
+    def loading_check(self, finish_count: Optional[int] = None):
+        if finish_count is None:
             finish_count = 0
             if self.pp_rank == 0:
-            for ack in self.cache_controller.ack_load_queue:
+                finish_count = self._count_ready_acks(
-                if not ack.finish_event.query():
+                    self.cache_controller.ack_load_queue
-                    break
+                )
-                finish_count += 1
+            finish_count_tensor = torch.tensor(
-        finish_count_tensor = torch.tensor(finish_count, dtype=torch.int, device="cpu")
+                finish_count, dtype=torch.int, device="cpu"
+            )
             self._all_reduce(finish_count_tensor, torch.distributed.ReduceOp.MIN)
             finish_count = finish_count_tensor.item()
 
@@ -1440,16 +1478,33 @@ class HiRadixCache(RadixCache):
         """
         return self.cache_controller.start_loading()
 
-    def flush_write_through_acks(self) -> None:
-        self.writing_check()
-
     def check_hicache_events(self):
         # Reap the previous round's PP-sync sends before issuing new ones.
         self._drain_async_work()
+
+        if self.pp_size != 1:
             self.writing_check()
             self.loading_check()
             if self.enable_storage:
                 self.drain_storage_control_queues()
+        else:
+            (
+                write_finish_count,
+                load_finish_count,
+                storage_queue_sizes,
+            ) = self._sync_hicache_ready_counts()
+            self.writing_check(finish_count=write_finish_count)
+            self.loading_check(finish_count=load_finish_count)
+
+            if self.enable_storage and storage_queue_sizes:
+                n_revoke, n_storage_hit, n_backup, n_release = storage_queue_sizes[:4]
+                self._drain_storage_control_queues_impl(
+                    n_revoke=n_revoke,
+                    n_storage_hit=n_storage_hit,
+                    n_backup=n_backup,
+                    n_release=n_release,
+                    log_metrics=True,
+                )
         if self.enable_storage_metrics:
             self.storage_metrics_collector.log_storage_metrics(
                 self.cache_controller.storage_backend.get_stats()

python/sglang/srt/mem_cache/storage/flexkv/README.md

@@ -379,8 +379,8 @@ Supported:
   here but requires `PoolTransfer` + `PoolHitPolicy` plumbing in
   `FlexKVConnector`.
 * Write-back acks are per-request (one `dec_lock_ref` per
-  `cache_finished_req`), not per-page like HiCache's
+  `cache_finished_req`), not per-page like HiCache's write-through
-  `flush_write_through_acks`.
+  ack queues.
 * `--radix-cache-backend=flexkv` and `--enable-flexkv` are
   mutually equivalent today; we don't yet emit a deprecation
   warning if both are set.

python/sglang/srt/mem_cache/unified_radix_cache.py

@@ -1746,7 +1746,64 @@ class UnifiedRadixCache(BasePrefixCache):
 
     # ---- HiCache: Async Event Management ----
 
-    def writing_check(self, write_back: bool = False) -> None:
+    def _count_ready_acks(self, ack_queue) -> int:
+        ready_count = 0
+        for ack in ack_queue:
+            if not ack.finish_event.query():
+                break
+            ready_count += 1
+        return ready_count
+
+    def _sync_hicache_ready_counts(
+        self,
+    ) -> tuple[int, int, tuple[int, ...], tuple[PoolName, ...]]:
+        cc = self.cache_controller
+        if cc is None:
+            write_acks = 0
+            load_acks = 0
+            storage_queue_sizes = ()
+            extra_pool_names = ()
+        else:
+            write_acks = self._count_ready_acks(cc.ack_write_queue)
+            load_acks = self._count_ready_acks(cc.ack_load_queue)
+            extra_release_queues = getattr(cc, "extra_host_mem_release_queues", {})
+            extra_pool_names = (
+                tuple(extra_release_queues) if self.enable_storage else ()
+            )
+            storage_queue_sizes = (
+                (
+                    cc.prefetch_revoke_queue.qsize(),
+                    cc.prefetch_hit_queue.qsize(),
+                    cc.ack_backup_queue.qsize(),
+                    cc.host_mem_release_queue.qsize(),
+                    *(extra_release_queues[name].qsize() for name in extra_pool_names),
+                )
+                if self.enable_storage
+                else ()
+            )
+
+        ready_counts = torch.tensor(
+            [
+                write_acks,
+                load_acks,
+                *storage_queue_sizes,
+            ],
+            dtype=torch.int,
+            device="cpu",
+        )
+        self._all_reduce(ready_counts, torch.distributed.ReduceOp.MIN)
+
+        count_values = list(map(int, ready_counts.tolist()))
+        return (
+            count_values[0],
+            count_values[1],
+            tuple(count_values[2:]),
+            extra_pool_names,
+        )
+
+    def writing_check(
+        self, write_back: bool = False, finish_count: Optional[int] = None
+    ) -> None:
         """Poll write-through completions."""
         cc = self.cache_controller
         if cc is None:
@@ -1764,16 +1821,15 @@ class UnifiedRadixCache(BasePrefixCache):
                 assert len(self.ongoing_write_through) == 0
             return
 
+        if finish_count is None:
             # Every rank must enter the all_reduce below; ongoing_write_through can
-        # diverge across ranks (e.g. a backup returning 0 on a subset).
+            # diverge across ranks (e.g. write_backup returning 0 on a subset).
             finish_count = 0
             if self.pp_rank == 0:
-            for ack in cc.ack_write_queue:
+                finish_count = self._count_ready_acks(cc.ack_write_queue)
-                if not ack.finish_event.query():
+            finish_count_tensor = torch.tensor(
-                    break
+                finish_count, dtype=torch.int, device="cpu"
-                finish_count += 1
+            )
-
-        finish_count_tensor = torch.tensor(finish_count, dtype=torch.int, device="cpu")
             self._all_reduce(finish_count_tensor, torch.distributed.ReduceOp.MIN)
             finish_count = finish_count_tensor.item()
 
@@ -1785,20 +1841,20 @@ class UnifiedRadixCache(BasePrefixCache):
                 self._finish_write_through_ack(ack_id)
             finish_count -= 1
 
-    def loading_check(self) -> None:
+    def loading_check(self, finish_count: Optional[int] = None) -> None:
         """Poll load-back completions."""
         cc = self.cache_controller
         if cc is None:
             return
+        if finish_count is None:
             # Every rank must enter the all_reduce below; ongoing_load_back can
             # diverge across ranks.
             finish_count = 0
             if self.pp_rank == 0:
-            for ack in cc.ack_load_queue:
+                finish_count = self._count_ready_acks(cc.ack_load_queue)
-                if not ack.finish_event.query():
+            finish_count_tensor = torch.tensor(
-                    break
+                finish_count, dtype=torch.int, device="cpu"
-                finish_count += 1
+            )
-        finish_count_tensor = torch.tensor(finish_count, dtype=torch.int, device="cpu")
             self._all_reduce(finish_count_tensor, torch.distributed.ReduceOp.MIN)
             finish_count = finish_count_tensor.item()
 
@@ -1867,19 +1923,44 @@ class UnifiedRadixCache(BasePrefixCache):
         """Called per scheduler step to poll async HiCache events."""
         # Reap the previous round's PP-sync sends before issuing new ones.
         self._drain_async_work()
+
+        if self.pp_size != 1:
             self.writing_check()
             self.loading_check()
             if self.enable_storage:
                 self.drain_storage_control_queues()
+        else:
+            (
+                write_finish_count,
+                load_finish_count,
+                storage_queue_sizes,
+                extra_pool_names,
+            ) = self._sync_hicache_ready_counts()
+            self.writing_check(finish_count=write_finish_count)
+            self.loading_check(finish_count=load_finish_count)
+
+            if self.enable_storage and storage_queue_sizes:
+                n_revoke, n_storage_hit, n_backup, n_release = storage_queue_sizes[:4]
+                extra_release_counts = {
+                    pool_name: count
+                    for pool_name, count in zip(
+                        extra_pool_names,
+                        storage_queue_sizes[4:],
+                    )
+                }
+                self._drain_storage_control_queues_impl(
+                    n_revoke=n_revoke,
+                    n_storage_hit=n_storage_hit,
+                    n_backup=n_backup,
+                    n_release=n_release,
+                    extra_release_counts=extra_release_counts,
+                    log_metrics=True,
+                )
         if self.enable_storage_metrics and self.storage_metrics_collector is not None:
             self.storage_metrics_collector.log_storage_metrics(
                 self.cache_controller.storage_backend.get_stats()
             )
 
-    def flush_write_through_acks(self) -> None:
-        """Flush pending write-through acknowledgements."""
-        self.writing_check()
-
     def ready_to_load_host_cache(self) -> int:
         """Notify the cache controller to start the KV cache loading."""
         if self.cache_controller is not None:

python/sglang/srt/session/streaming_session.py

@@ -602,9 +602,6 @@ class StreamingSession(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()
 

test/registered/unit/mem_cache/test_hicache_load_back_timing.py

@@ -66,16 +66,15 @@ class TestLoadBackDurationMetric(CustomTestCase):
             num_tokens=1024,
             timing_enabled=True,
         )
-        stub = SimpleNamespace(
+        stub = object.__new__(HiRadixCache)
-            cache_controller=SimpleNamespace(ack_load_queue=[ack]),
+        stub.cache_controller = SimpleNamespace(ack_load_queue=[ack])
-            ongoing_load_back={1: object(), 2: object()},
+        stub.ongoing_load_back = {1: object(), 2: object()}
-            dec_lock_ref=MagicMock(),
+        stub.dec_lock_ref = MagicMock()
-            metrics_collector=MagicMock(),
+        stub.metrics_collector = MagicMock()
-            pp_rank=0,
+        stub.pp_rank = 0
-            _all_reduce=MagicMock(),
+        stub._all_reduce = MagicMock()
-        )
 
-        HiRadixCache.loading_check(stub)
+        stub.loading_check()
 
         stub.metrics_collector.increment_load_back_num_tokens.assert_called_once_with(
             1024
@@ -102,16 +101,15 @@ class TestLoadBackDurationMetric(CustomTestCase):
             num_tokens=512,
             timing_enabled=False,
         )
-        stub = SimpleNamespace(
+        stub = object.__new__(HiRadixCache)
-            cache_controller=SimpleNamespace(ack_load_queue=[ack]),
+        stub.cache_controller = SimpleNamespace(ack_load_queue=[ack])
-            ongoing_load_back={7: object()},
+        stub.ongoing_load_back = {7: object()}
-            dec_lock_ref=MagicMock(),
+        stub.dec_lock_ref = MagicMock()
-            metrics_collector=MagicMock(),
+        stub.metrics_collector = MagicMock()
-            pp_rank=0,
+        stub.pp_rank = 0
-            _all_reduce=MagicMock(),
+        stub._all_reduce = MagicMock()
-        )
 
-        HiRadixCache.loading_check(stub)
+        stub.loading_check()
 
         stub.metrics_collector.increment_load_back_num_tokens.assert_called_once_with(
             512