From f19d5c90fd8bbf678a6132bda87921d350480a30 Mon Sep 17 00:00:00 2001 From: alphabetc1 <2508695655@qq.com> Date: Thu, 17 Sep 2026 12:12:15 +0800 Subject: [PATCH] perf: pipeline the HiCache ack-count all_reduce across scheduler steps --- python/sglang/srt/environ.py | 5 + .../srt/mem_cache/unified_radix_cache.py | 107 +++++++++++++++--- .../mem_cache/test_hiradix_pp_sync_drain.py | 81 ++++++++++++- 3 files changed, 174 insertions(+), 19 deletions(-) diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index e2190ce77468..1ddc2265a9e3 100644 --- a/python/sglang/srt/environ.py +++ b/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) diff --git a/python/sglang/srt/mem_cache/unified_radix_cache.py b/python/sglang/srt/mem_cache/unified_radix_cache.py index dd3dc60aecb3..64d6d4bf7360 100644 --- a/python/sglang/srt/mem_cache/unified_radix_cache.py +++ b/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.""" diff --git a/test/registered/unit/mem_cache/test_hiradix_pp_sync_drain.py b/test/registered/unit/mem_cache/test_hiradix_pp_sync_drain.py index c3600af323eb..b7bb29788b29 100644 --- a/test/registered/unit/mem_cache/test_hiradix_pp_sync_drain.py +++ b/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])