fix: synchronize custom collectives before return

[jpy794]

Motivation

This PR fixes a race in custom collective kernels where a collective may return before all peer reads from the input buffer have completed.

In CUDA graph / overlapped execution scenarios, the caller may reuse or overwrite the input tensor immediately after the custom collective returns. Without an end-of-kernel synchronization, another rank can still be reading that input buffer, which can lead to corrupted collective outputs. The repro script demonstrates this by filling the input with NaNs right after the collective and checking whether NaNs leak into the output.

cc @ColorsWind

Technical Details

This change adds end_sync<ngpus, true>(sg, self_sg, rank) before returning from custom collective kernels that directly read peer input buffers.

Covered paths in csrc/include/custom_all_reduce.cuh include:

• all-gather kernels:

  • allgather_naive
  • allgather_vec
  • allgather_lastdim

• reduce-scatter kernels:

  • reduce_scatter_split_first_dim
  • reduce_scatter_split_lastdim_naive
  • reduce_scatter_split_lastdim
  • reduce_scatter_split_middim_naive
  • reduce_scatter_split_middim

• 1-stage fused all-reduce + RMSNorm kernels:

  • standard variant
  • per-group quant variant
  • MXFP4 variant

The synchronization ensures that all ranks complete their remote reads before any rank is allowed to proceed and potentially mutate or recycle its input buffer.

Test Plan

Use the custom collective race repro script here https://gist.github.com/jpy794/a2849ad6d3d2a3f0d1a03c6c1ec977ac

rm -rf aiter/aiter/jit/module_custom_all_reduce.so \
       aiter/aiter/jit/build/module_custom_all_reduce
python custom_collective_race_repro.py

The repro captures a CUDA graph that repeatedly:

1. resets the collective input to finite values,
2. runs the custom collective,
3. overwrites rank 0's input with NaNs immediately after the collective,
4. checks the collective output on device for leaked NaNs.
   Expected result after this fix: all ranks report bad_nan_rows=0.

Test Result

Branch	AG	RS	AR 1-stage	Exit
main	FAIL, 1477 bad rows	FAIL, 676 bad rows	FAIL, 59 bad rows	1
fix-agrs-race	PASS, 0 bad rows	PASS, 0 bad rows	PASS, 0 bad rows	0

Submission Checklist

• Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

[github-actions]

🏷️ CI Guide

Runs automatically on every PR:

• ✅ Pre-checks (submodule verification, code formatting)
• ✅ Aiter op tests (gfx942 + gfx950)
• ✅ Triton tests on MI35X (only when aiter/ops/triton/** or related paths are changed)

Extended tests (opt-in via labels):

Label	Tests
ci:triton-300x	Run an additional Triton test job on MI300X in PRs; main branch always runs both MI35X and MI300X
ci:sglang	SGLang integration tests: DeepSeek-R1-MXFP4 accuracy, Qwen 3.5 accuracy
ci:atom	ATOM benchmark: DeepSeek-R1-0528, GPT-OSS-120B
ci:atom_full	ATOM accuracy suite for PR and main models from ATOM models_accuracy.json
ci:vllm	vLLM benchmark: GPT-OSS-120B, DeepSeek-R1-0528, Kimi-K2.5
ci:all	All standard extended tests (excludes ci:atom_full)

Only add ci:atom_full for FlyDSL or Triton upgrades.
Add labels via the sidebar or gh pr edit 4082 --add-label <label>

[Copilot]

Pull request overview

This PR adds an explicit end-of-kernel cross-rank synchronization to custom collective GPU kernels that directly read peer input buffers, preventing the kernel from returning before all peer reads complete (avoiding races when the caller immediately reuses/overwrites the input, e.g., under CUDA graph / overlapped execution).

Changes:

• Add end_sync<ngpus, true>(sg, self_sg, rank) to all-gather kernels after completing peer-input reads.
• Add end_sync<ngpus, true>(sg, self_sg, rank) to reduce-scatter kernels after completing peer-input reads/reductions.
• Add end_sync<ngpus, true>(sg, self_sg, rank) to the 1-stage fused all-reduce + RMSNorm kernels after the main compute loop.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[jpy794]

Performance Test

benchmark script: https://gist.github.com/jpy794/d1ffc4a89c9ce6f36e57374a41a004ae

python tools/custom_collective_microbench.py --graph-iters 100 --label fix-agrs-race-graph100 --output-json results/custom_collective_microbench/fix-agrs-race-graph100-rerun.json
AITER_ROOT=/tmp/up-aiter-main python tools/custom_collective_microbench.py --graph-iters 100 --label main-graph100 --output-json results/custom_collective_microbench/main-graph100-rerun.json --baseline-json results/custom_collective_microbench/fix-agrs-race-graph100-rerun.json

case	main	current fix branch	delta
ag	6.910 us ±2.6%	8.994 us ±7.5%	+2.084 us (+30.16%)
rs	7.069 us ±4.2%	8.363 us ±1.5%	+1.294 us (+18.30%)
ar1stage	9.087 us ±6.0%	10.216 us ±2.0%	+1.129 us (+12.42%)

[ColorsWind]

Related to #3514, which fixed this same race for cross_device_reduce_1stage by adding the missing end_sync barrier. This PR extends the same fix to the remaining collectives (all-gather, reduce-scatter, and the 1-stage fused all-reduce + RMSNorm variants).