From bd7e54d7379e437cf5f027382d6ca214e046626b Mon Sep 17 00:00:00 2001 From: Zhiyao Jiang Date: Thu, 9 Jul 2026 04:30:51 -0400 Subject: [PATCH] [AMD] Fix AITER custom all-gather CUDA-graph capture crash under torch_memory_saver (#30557) Co-authored-by: Xinyu Jiang --- python/sglang/srt/distributed/parallel_state.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/python/sglang/srt/distributed/parallel_state.py b/python/sglang/srt/distributed/parallel_state.py index 66263c275..73a64ef9b 100644 --- a/python/sglang/srt/distributed/parallel_state.py +++ b/python/sglang/srt/distributed/parallel_state.py @@ -953,7 +953,8 @@ class GroupCoordinator: # 16B alignment, weak-contiguous, supported topology, and per-rank # size <= max_size/(world*2). # On a hit, writes directly into the caller's pre-allocated `output` via - # all_gather_reg during CUDA-graph capture and all_gather_unreg otherwise. + # all_gather_reg during CUDA-graph capture, and all_gather_unreg + # under torch_memory_saver and other paths. ca_comm = self.ca_comm if ( is_hip() @@ -966,7 +967,10 @@ class GroupCoordinator: ): if getattr(ca_comm, "_IS_CAPTURING", False): if torch.cuda.is_current_stream_capturing(): - ca_comm.all_gather_reg(input, out=output, dim=0) + if envs.SGLANG_MEMORY_SAVER_CUDA_GRAPH.get(): + ca_comm.all_gather_unreg(input, out=output, dim=0) + else: + ca_comm.all_gather_reg(input, out=output, dim=0) elif is_in_tc_piecewise_cuda_graph(): ca_comm.all_gather_unreg(input, out=output, dim=0) else: