From f7e2e468843a6fe93012731ba1defedcafc7ba28 Mon Sep 17 00:00:00 2001 From: zovonoir Date: Wed, 3 Jun 2026 14:57:03 +0800 Subject: [PATCH] Fix: add missing end_sync in cross_device_reduce_1stage Without this barrier, a fast rank can exit the kernel while a slow rank is still reading peer input via IPC. The fast rank's caller can then overwrite the input slot (e.g. PyTorch graph_pool reuse in a captured CUDA graph), causing the slow rank to read garbage and produce NaN / Inf in the AR output. All other AR kernels in this file already call end_sync before exit. --- csrc/include/custom_all_reduce.cuh | 1 + 1 file changed, 1 insertion(+) diff --git a/csrc/include/custom_all_reduce.cuh b/csrc/include/custom_all_reduce.cuh index 6312190414..dec526d801 100644 --- a/csrc/include/custom_all_reduce.cuh +++ b/csrc/include/custom_all_reduce.cuh @@ -467,6 +467,7 @@ __global__ void __launch_bounds__(512, 1) cross_device_reduce_1stage(RankData* _ buf = next_buf; } + end_sync(sg, self_sg, rank); } template