diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 892c20305..b212c9441 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/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 diff --git a/python/sglang/srt/mem_cache/base_prefix_cache.py b/python/sglang/srt/mem_cache/base_prefix_cache.py index f54a986b1..fe208c986 100644 --- a/python/sglang/srt/mem_cache/base_prefix_cache.py +++ b/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 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 81a6802c5..dbcec3547 100644 --- a/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py +++ b/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() diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index 8551a3882..86811cabf 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/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,19 +1027,21 @@ class HiRadixCache(RadixCache): assert len(self.ongoing_write_through) == 0 return - # 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: - if not ack.finish_event.query(): - break - 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() + 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: + finish_count = self._count_ready_acks( + self.cache_controller.ack_write_queue + ) + 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() if finish_count > 0: logger.debug(f"Process {finish_count} write back operations") @@ -1016,16 +1052,18 @@ class HiRadixCache(RadixCache): self._finish_write_through_ack(ack_id, release_lock=True) finish_count -= 1 - def loading_check(self): - finish_count = 0 - if self.pp_rank == 0: - for ack in self.cache_controller.ack_load_queue: - if not ack.finish_event.query(): - break - 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() + def loading_check(self, finish_count: Optional[int] = None): + if finish_count is None: + finish_count = 0 + if self.pp_rank == 0: + finish_count = self._count_ready_acks( + self.cache_controller.ack_load_queue + ) + 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() if finish_count > 0: logger.debug(f"Process {finish_count} load operations") @@ -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() - self.writing_check() - self.loading_check() - if self.enable_storage: - self.drain_storage_control_queues() + + 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() diff --git a/python/sglang/srt/mem_cache/storage/flexkv/README.md b/python/sglang/srt/mem_cache/storage/flexkv/README.md index 92e55a3d5..b3724540c 100644 --- a/python/sglang/srt/mem_cache/storage/flexkv/README.md +++ b/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 - `flush_write_through_acks`. + `cache_finished_req`), not per-page like HiCache's write-through + 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. diff --git a/python/sglang/srt/mem_cache/unified_radix_cache.py b/python/sglang/srt/mem_cache/unified_radix_cache.py index e10561f24..80cfd171a 100644 --- a/python/sglang/srt/mem_cache/unified_radix_cache.py +++ b/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,18 +1821,17 @@ class UnifiedRadixCache(BasePrefixCache): assert len(self.ongoing_write_through) == 0 return - # Every rank must enter the all_reduce below; ongoing_write_through can - # diverge across ranks (e.g. a backup returning 0 on a subset). - finish_count = 0 - if self.pp_rank == 0: - for ack in cc.ack_write_queue: - if not ack.finish_event.query(): - break - 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() + 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). + finish_count = 0 + if self.pp_rank == 0: + finish_count = self._count_ready_acks(cc.ack_write_queue) + 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() # Process completed acks while finish_count > 0: @@ -1785,22 +1841,22 @@ 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 - # 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: - if not ack.finish_event.query(): - break - 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() + 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: + finish_count = self._count_ready_acks(cc.ack_load_queue) + 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() while finish_count > 0: ack = cc.ack_load_queue.pop(0) @@ -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() - self.writing_check() - self.loading_check() - if self.enable_storage: - self.drain_storage_control_queues() + + 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: diff --git a/python/sglang/srt/session/streaming_session.py b/python/sglang/srt/session/streaming_session.py index 28cbd4cc6..95f45db47 100644 --- a/python/sglang/srt/session/streaming_session.py +++ b/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() diff --git a/test/registered/unit/mem_cache/test_hicache_load_back_timing.py b/test/registered/unit/mem_cache/test_hicache_load_back_timing.py index 75537503f..e056a07a5 100644 --- a/test/registered/unit/mem_cache/test_hicache_load_back_timing.py +++ b/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( - cache_controller=SimpleNamespace(ack_load_queue=[ack]), - ongoing_load_back={1: object(), 2: object()}, - dec_lock_ref=MagicMock(), - metrics_collector=MagicMock(), - pp_rank=0, - _all_reduce=MagicMock(), - ) + stub = object.__new__(HiRadixCache) + stub.cache_controller = SimpleNamespace(ack_load_queue=[ack]) + stub.ongoing_load_back = {1: object(), 2: object()} + stub.dec_lock_ref = MagicMock() + stub.metrics_collector = MagicMock() + stub.pp_rank = 0 + 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( - cache_controller=SimpleNamespace(ack_load_queue=[ack]), - ongoing_load_back={7: object()}, - dec_lock_ref=MagicMock(), - metrics_collector=MagicMock(), - pp_rank=0, - _all_reduce=MagicMock(), - ) + stub = object.__new__(HiRadixCache) + stub.cache_controller = SimpleNamespace(ack_load_queue=[ack]) + stub.ongoing_load_back = {7: object()} + stub.dec_lock_ref = MagicMock() + stub.metrics_collector = MagicMock() + stub.pp_rank = 0 + stub._all_reduce = MagicMock() - HiRadixCache.loading_check(stub) + stub.loading_check() stub.metrics_collector.increment_load_back_num_tokens.assert_called_once_with( 512