diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index 1e8b838d9..158685471 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/python/sglang/srt/mem_cache/hiradix_cache.py @@ -207,18 +207,17 @@ class HiRadixCache(RadixCache): if not waited and self.tp_world_size > 1: torch.distributed.barrier(group=self.tp_group) - def _reap_completed_async_work(self): + def _drain_async_work(self): """ - Poll outstanding async work and reap completed ones. + Block until all outstanding async sends are consumed, then clear. - Must be called in the scheduler thread. + Called at the start of each event round, so work_list holds the sends + accumulated since the last round. This bounds it and applies + backpressure when a downstream PP rank lags. Scheduler thread only. """ - count = 0 - while count < len(self.work_list) and self.work_list[count].is_completed(): - count += 1 - if count > 0: - logger.debug(f"Reap {count} completed async work") - self.work_list = self.work_list[count:] + for work in self.work_list: + work.wait() + self.work_list.clear() def _all_reduce(self, data: torch.Tensor, tp_reduce_op: torch.distributed.ReduceOp): """ @@ -1279,11 +1278,12 @@ class HiRadixCache(RadixCache): self.writing_check() def check_hicache_events(self): + # Reap the previous round's PP-sync sends before issuing new ones. + self._drain_async_work() self.writing_check() self.loading_check() if self.enable_storage: self.drain_storage_control_queues() - self._reap_completed_async_work() if self.enable_storage_metrics: self.storage_metrics_collector.log_storage_metrics( self.cache_controller.storage_backend.get_stats() diff --git a/python/sglang/srt/mem_cache/unified_radix_cache.py b/python/sglang/srt/mem_cache/unified_radix_cache.py index 9514172d8..fa26bbee5 100644 --- a/python/sglang/srt/mem_cache/unified_radix_cache.py +++ b/python/sglang/srt/mem_cache/unified_radix_cache.py @@ -395,18 +395,17 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache): if not waited and self.tp_world_size > 1: torch.distributed.barrier(group=self.tp_group) - def _reap_completed_async_work(self): + def _drain_async_work(self): """ - Poll outstanding async work and reap completed ones. + Block until all outstanding async sends are consumed, then clear. - Must be called in the scheduler thread. + Called at the start of each event round, so work_list holds the sends + accumulated since the last round. This bounds it and applies + backpressure when a downstream PP rank lags. Scheduler thread only. """ - count = 0 - while count < len(self.work_list) and self.work_list[count].is_completed(): - count += 1 - if count > 0: - logger.debug(f"Reap {count} completed async work") - self.work_list = self.work_list[count:] + for work in self.work_list: + work.wait() + self.work_list.clear() def _all_reduce(self, data: torch.Tensor, tp_reduce_op: torch.distributed.ReduceOp): """ @@ -2463,11 +2462,12 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache): def check_hicache_events(self) -> None: """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() self.writing_check() self.loading_check() if self.enable_storage: self.drain_storage_control_queues() - self._reap_completed_async_work() 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() 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 new file mode 100644 index 000000000..6e1ba43fe --- /dev/null +++ b/test/registered/unit/mem_cache/test_hiradix_pp_sync_drain.py @@ -0,0 +1,50 @@ +"""Unit test for HiRadixCache._drain_async_work PP-sync backpressure.""" + +import unittest + +from sglang.srt.mem_cache.hiradix_cache import HiRadixCache +from sglang.srt.mem_cache.unified_radix_cache import UnifiedRadixCache +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=1, suite="base-a-test-cpu") + + +class _FakeWork: + def __init__(self): + self.waited = False + + def wait(self): + self.waited = True + + +class _Holder: + """Minimal carrier exposing only what _drain_async_work touches.""" + + +class TestPPSyncDrain(unittest.TestCase): + def _drain_fns(self): + return (HiRadixCache._drain_async_work, UnifiedRadixCache._drain_async_work) + + def test_drain_waits_all_and_clears(self): + for drain in self._drain_fns(): + holder = _Holder() + works = [_FakeWork(), _FakeWork(), _FakeWork()] + holder.work_list = list(works) + + drain(holder) + + self.assertTrue(all(w.waited for w in works)) + self.assertEqual(holder.work_list, []) + + def test_drain_empty_is_noop(self): + for drain in self._drain_fns(): + holder = _Holder() + holder.work_list = [] + + drain(holder) + + self.assertEqual(holder.work_list, []) + + +if __name__ == "__main__": + unittest.main()