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 7b56634cc..381af58be 100644 --- a/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py +++ b/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py @@ -396,14 +396,15 @@ class HiMambaRadixCache(MambaRadixCache): assert len(self.ongoing_write_through) == 0 return - if len(self.ongoing_write_through) == 0: - return - + # Every rank must enter the all_reduce below; ongoing_write_through can + # diverge across ranks because loading_check processes DMA completions + # independently (no cross-rank sync). finish_count = 0 - for _, finish_event, ack_list in self.cache_controller.ack_write_queue: - if not finish_event.query(): - break - finish_count += 1 + if len(self.ongoing_write_through) > 0: + for _, finish_event, ack_list in self.cache_controller.ack_write_queue: + if not finish_event.query(): + break + finish_count += 1 queue_size = torch.tensor(finish_count, dtype=torch.int, device="cpu") if self.tp_world_size > 1: @@ -426,17 +427,30 @@ class HiMambaRadixCache(MambaRadixCache): finish_count -= 1 def loading_check(self): + # Every rank must enter the all_reduce below; ongoing_load_back can + # diverge across ranks. finish_count = 0 for _, finish_event, ack_list in self.cache_controller.ack_load_queue: if not finish_event.query(): - # the KV cache loading is still ongoing break finish_count += 1 + + queue_size = torch.tensor(finish_count, dtype=torch.int, device="cpu") + if self.tp_world_size > 1: + torch.distributed.all_reduce( + queue_size, + op=torch.distributed.ReduceOp.MIN, + group=self.tp_group, + ) + finish_count = int(queue_size.item()) + + while finish_count > 0: + _, finish_event, ack_list = self.cache_controller.ack_load_queue.pop(0) + finish_event.synchronize() for ack_id in ack_list: end_node = self.ongoing_load_back.pop(ack_id) self.dec_lock_ref(end_node) - - del self.cache_controller.ack_load_queue[:finish_count] + finish_count -= 1 def ready_to_load_host_cache(self) -> int: return self.cache_controller.start_loading() diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index 7fb56f943..1e8b838d9 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/python/sglang/srt/mem_cache/hiradix_cache.py @@ -916,10 +916,10 @@ class HiRadixCache(RadixCache): assert len(self.ongoing_write_through) == 0 return - # NOTE: all ranks has the same ongoing_write_through, can skip sync if empty - if 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 _, finish_event, ack_list in self.cache_controller.ack_write_queue: