[HiCache][perf]: pipeline the HiCache ack-count all_reduce across scheduler steps

python/sglang/srt/environ.py

@@ -750,6 +750,11 @@ class Envs:
     SGLANG_HICACHE_HF3FS_CONFIG_PATH = EnvStr(None)
     SGLANG_HICACHE_DECODE_OFFLOAD_STRIDE = EnvInt(None)
     SGLANG_HICACHE_SKIP_HOST_DUPLICATE_RECLAIM = EnvBool(False)
+    # Reduce the HiCache ack counts across TP ranks asynchronously: issue the
+    # all_reduce at the end of a step and read it on the next one, so the
+    # scheduler thread never blocks on the collective. Falls back to the
+    # blocking reduce under PP or when CP and TP both need a reduce.
+    SGLANG_ENABLE_HICACHE_ASYNC_ACK_SYNC = EnvBool(True)
     SGLANG_HICACHE_FILE_BACKEND_STORAGE_DIR = EnvStr(None)
     # File-backend LRU eviction (opt-in; sizes accept SI/IEC suffixes, "0" disables).
     SGLANG_HICACHE_FILE_BACKEND_MAX_SIZE = EnvStr(None)

python/sglang/srt/mem_cache/unified_radix_cache.py

@@ -256,6 +256,17 @@ def __init__(
         )
         self.pp_rank = params.pp_rank
         self.pp_size = params.pp_size
+        self._ready_counts_group = self._single_ready_counts_group()
+        # One collective per step is the precondition for pipelining it; CP+TP
+        # (two groups) and PP pipelines keep the blocking path.
+        self._hicache_async_ack_sync = (
+            envs.SGLANG_ENABLE_HICACHE_ASYNC_ACK_SYNC.get()
+            and self.pp_size == 1
+            and self._ready_counts_group is not None
+        )
+        self._pending_ready_counts: Optional[
+            tuple[torch.distributed.Work, torch.Tensor, tuple, int]
+        ] = None
         self.work_list: list[torch.distributed.Work] = []
 
         # HiCache D↔H defaults (overridden by init_hicache)
@@ -300,6 +311,20 @@ def __init__(
             f"Tree Core: {type(self.tree_core).__name__}"
         )
 
+    def _single_ready_counts_group(self):
+        """The one process group _all_reduce_attn_groups would reduce over, or
+        None when it would reduce over two groups or over nobody."""
+        groups = [
+            group
+            for group in (self.attn_cp_group, self.attn_tp_group)
+            if group is not None and torch.distributed.get_world_size(group=group) > 1
+        ]
+        if len(groups) == 1:
+            return groups[0]
+        if not groups and self.tp_world_size > 1:
+            return self.tp_group
+        return None
+
     def _all_reduce_attn_groups(self, tensor: torch.Tensor, op):
         reduced = False
         for group in (self.attn_cp_group, self.attn_tp_group):
@@ -3071,9 +3096,9 @@ def _count_ready_acks(self, ack_queue) -> int:
             ready_count += 1
         return ready_count
 
-    def _sync_hicache_ready_counts(
-        self,
-    ) -> tuple[int, int, tuple[int, ...], tuple[PoolName, ...]]:
+    def _ready_counts_tensor(self) -> tuple[torch.Tensor, tuple[PoolName, ...], int]:
+        """Local ack counts, storage queue sizes, and the reclaim digest, laid
+        out for a MIN all_reduce across ranks."""
         cc = self.cache_controller
         extra_release_queues = getattr(cc, "extra_host_mem_release_queues", {})
         extra_pool_names = tuple(extra_release_queues) if self.enable_storage else ()
@@ -3115,8 +3140,17 @@ def _sync_hicache_ready_counts(
             dtype=torch.int64,
             device="cpu",
         )
+        return ready_counts, extra_pool_names, digest
+
+    def _sync_hicache_ready_counts(
+        self,
+    ) -> tuple[int, int, tuple[int, ...], tuple[PoolName, ...]]:
+        ready_counts, extra_pool_names, digest = self._ready_counts_tensor()
         self._all_reduce(ready_counts, torch.distributed.ReduceOp.MIN)
+        return self._parse_ready_counts(ready_counts, extra_pool_names, digest)
 
+    @staticmethod
+    def _parse_ready_counts(ready_counts, extra_pool_names, digest):
         count_values = list(map(int, ready_counts.tolist()))
         assert digest == count_values[-2] and digest == -count_values[-1], (
             "write_back duplicate-reclaim victims diverged across PP/TP ranks"
@@ -3128,6 +3162,31 @@ def _sync_hicache_ready_counts(
             extra_pool_names,
         )
 
+    def _issue_async_ready_counts(self) -> None:
+        """Start this step's MIN all_reduce of the local ack counts; the result
+        is read one step later. Must run after this step's acks were popped,
+        or the popped acks would be counted a second time."""
+        ready_counts, extra_pool_names, digest = self._ready_counts_tensor()
+        work = torch.distributed.all_reduce(
+            ready_counts,
+            op=torch.distributed.ReduceOp.MIN,
+            group=self._ready_counts_group,
+            async_op=True,
+        )
+        self._pending_ready_counts = (work, ready_counts, extra_pool_names, digest)
+
+    def _consume_async_ready_counts(self):
+        """The previous step's reduced counts, or None on the first step. Every
+        count is a MIN over what each rank had ready then, so it is still safe
+        to pop now: acks only accumulate between two steps."""
+        pending = self._pending_ready_counts
+        if pending is None:
+            return None
+        self._pending_ready_counts = None
+        work, ready_counts, extra_pool_names, digest = pending
+        work.wait()
+        return self._parse_ready_counts(ready_counts, extra_pool_names, digest)
+
     def writing_check(
         self, write_back: bool = False, finish_count: Optional[int] = None
     ) -> None:
@@ -3318,12 +3377,33 @@ def check_hicache_events(self) -> None:
         # in get_next_batch_to_run, abort_request, and the PD prefill release.
         self.flush_pending_backups()
 
-        (
-            write_finish_count,
-            load_finish_count,
-            storage_queue_sizes,
-            extra_pool_names,
-        ) = self._sync_hicache_ready_counts()
+        if self._hicache_async_ack_sync:
+            ready = self._consume_async_ready_counts()
+        else:
+            ready = self._sync_hicache_ready_counts()
+        if ready is not None:
+            self._apply_ready_counts(*ready)
+        if self._hicache_async_ack_sync:
+            # After the pops above, so the next reduce counts only unpopped acks.
+            self._issue_async_ready_counts()
+        if self.buffer_pipeline is not None:
+            self.buffer_pipeline.flush_pending_writes()
+        if self.enable_storage_metrics and self.storage_metrics_collector is not None:
+            storage_metrics = self.cache_controller.storage_backend.get_stats()
+            if storage_metrics is None:
+                storage_metrics = StorageMetrics()
+            if not hasattr(storage_metrics, "prefetch_stats"):
+                storage_metrics.prefetch_stats = self.prefetch_outcome_stats_snapshot()
+            self.storage_metrics_collector.log_storage_metrics(storage_metrics)
+
+    def _apply_ready_counts(
+        self,
+        write_finish_count: int,
+        load_finish_count: int,
+        storage_queue_sizes: tuple[int, ...],
+        extra_pool_names: tuple[PoolName, ...],
+    ) -> None:
+        """Pop the acks and storage-queue entries every rank agreed are ready."""
         self.writing_check(finish_count=write_finish_count)
         self.loading_check(finish_count=load_finish_count)
 
@@ -3341,15 +3421,6 @@ def check_hicache_events(self) -> None:
                 extra_release_counts=extra_release_counts,
                 log_metrics=True,
             )
-        if self.buffer_pipeline is not None:
-            self.buffer_pipeline.flush_pending_writes()
-        if self.enable_storage_metrics and self.storage_metrics_collector is not None:
-            storage_metrics = self.cache_controller.storage_backend.get_stats()
-            if storage_metrics is None:
-                storage_metrics = StorageMetrics()
-            if not hasattr(storage_metrics, "prefetch_stats"):
-                storage_metrics.prefetch_stats = self.prefetch_outcome_stats_snapshot()
-            self.storage_metrics_collector.log_storage_metrics(storage_metrics)
 
     def flush_pending_backups(self) -> None:
         """Submit pending D2H backups as a merged operation."""

test/registered/unit/mem_cache/test_hiradix_pp_sync_drain.py

@@ -2,7 +2,7 @@
 
 import unittest
 from types import SimpleNamespace
-from unittest.mock import MagicMock
+from unittest.mock import MagicMock, patch
 
 import torch
 
@@ -63,6 +63,8 @@ def _make_cache(self, pp_rank, write_ready, load_ready):
         cache.storage_metrics_collector = None
         cache.buffer_pipeline = None
         cache.linker = None
+        cache._hicache_async_ack_sync = False
+        cache._pending_ready_counts = None
         cache._drain_async_work = MagicMock()
         cache._all_reduce = MagicMock()
         cache.writing_check = MagicMock()
@@ -113,3 +115,80 @@ def reduce_to_min(counts, _):
 
 if __name__ == "__main__":
     unittest.main()
+
+
+class TestUnifiedAsyncAckSync(unittest.TestCase):
+    """The pipelined ack-count reduce: nothing is popped on the first step, the
+    previous step's reduced counts are popped on the next one, and the reduce
+    issued after popping counts only the acks still queued."""
+
+    def _make_cache(self, write_ready):
+        cache = object.__new__(UnifiedRadixCache)
+        cache.tree_core = SimpleNamespace(
+            enable_storage=False,
+            write_back_duplicate_reclaim_digest=0,
+        )
+        cache.pp_rank = 0
+        cache.pp_size = 1
+        cache.enable_storage = False
+        cache.enable_storage_metrics = False
+        cache.storage_metrics_collector = None
+        cache.buffer_pipeline = None
+        cache.linker = None
+        cache._hicache_async_ack_sync = True
+        cache._pending_ready_counts = None
+        cache._ready_counts_group = object()
+        cache._drain_async_work = MagicMock()
+        cache.loading_check = MagicMock()
+        cache.cache_controller = SimpleNamespace(
+            start_writing=MagicMock(),
+            ack_write_queue=[
+                SimpleNamespace(
+                    finish_event=SimpleNamespace(query=MagicMock(return_value=ready))
+                )
+                for ready in write_ready
+            ],
+            ack_load_queue=[],
+        )
+        popped = []
+
+        def writing_check(finish_count):
+            popped.append(finish_count)
+            del cache.cache_controller.ack_write_queue[:finish_count]
+
+        cache.writing_check = writing_check
+        return cache, popped
+
+    def test_pops_previous_counts_then_reduces_the_remainder(self):
+        cache, popped = self._make_cache([True, True, False])
+        issued = []
+
+        def fake_all_reduce(tensor, op, group, async_op):
+            issued.append(tensor.clone())
+            return SimpleNamespace(wait=MagicMock())
+
+        with patch.object(torch.distributed, "all_reduce", fake_all_reduce):
+            cache.check_hicache_events()
+            # First step: nothing to consume, one reduce issued with 2 ready acks.
+            self.assertEqual(popped, [])
+            self.assertEqual(issued[0][0].item(), 2)
+
+            cache.check_hicache_events()
+            # Second step: pop the 2 acks from step one, then reduce the rest.
+            self.assertEqual(popped, [2])
+            self.assertEqual(len(cache.cache_controller.ack_write_queue), 1)
+            self.assertEqual(issued[1][0].item(), 0)
+
+    def test_reduce_result_is_waited_before_use(self):
+        cache, popped = self._make_cache([True])
+        waited = MagicMock()
+
+        def fake_all_reduce(tensor, op, group, async_op):
+            return SimpleNamespace(wait=waited)
+
+        with patch.object(torch.distributed, "all_reduce", fake_all_reduce):
+            cache.check_hicache_events()
+            waited.assert_not_called()
+            cache.check_hicache_events()
+            waited.assert_called_once()
+            self.assertEqual(popped, [1])