[AMD] Fix AITER custom reduce-scatter CUDA-graph capture crash under torch_memory_saver (#34203)

2026-08-10 15:36:15 -07:00
parent c80a38edcd
commit 2c72323d90
1 changed files with 4 additions and 1 deletions
@@ -1080,7 +1080,10 @@ class GroupCoordinator:
            return False
        if getattr(ca_comm, "_IS_CAPTURING", False):
            if torch.cuda.is_current_stream_capturing():
-                ca_comm.reduce_scatter(input, output, registered=True)
+                if envs.SGLANG_MEMORY_SAVER_CUDA_GRAPH.get():
+                    ca_comm.reduce_scatter(input, output, registered=False)
+                else:
+                    ca_comm.reduce_scatter(input, output, registered=True)
            elif is_in_tc_piecewise_cuda_graph():
                ca_comm.reduce_scatter(input, output, registered=False)
            else: