perf: eliminate end_sync in custom allreduce by delaying input tensor release until next sync

[jpy794]

Motivation

Follow-up to #4082.

PR #4082 fixed a custom collective race by adding an end-of-kernel cross-rank end_sync, ensuring peer ranks finish reading input buffers before the caller can reuse or overwrite them. That fix is correct, but the extra device-side synchronization adds measurable latency to small custom collective paths.

This PR keeps the race fix while providing a performance path that avoids the final end_sync. Instead of synchronizing at the end of each collective, Python keeps the input tensor alive until the next custom collective launch. The next collective's existing start synchronization proves all ranks have advanced past the previous collective, so the previous input can be safely released at that point.

cc @ColorsWind

Technical Details

This change adds an optional end_sync flag through the custom allreduce Python/C++/HIP stack.

When hold_inputs=False, behavior stays unchanged and kernels still run the final end_sync.

When hold_inputs=True:

• custom collective launches pass end_sync=False
• kernels skip the final end_sync<ngpus, true>(...)
• CustomAllreduce records the input tensor in Python
• the recorded input is replaced on the next collective launch, after the next collective's start sync has established that all ranks are done with the previous input
• recorded inputs are cleared on communicator close

The feature is controlled through:

AITER_CUSTOM_AR_HOLD_INPUTS=1

Covered paths include:

• all-gather
• reduce-scatter
• fused all-reduce + RMSNorm
• fused all-reduce + RMSNorm quant variants

Test Plan

You can find reproduce scripts here https://gist.github.com/jpy794/b21fd3c7b9a9ac1d0e87db62824cf78b

Race repro

rm -rf aiter/aiter/jit/module_custom_all_reduce.so \
       aiter/aiter/jit/build/module_custom_all_reduce

AITER_CUSTOM_AR_HOLD_INPUTS=1 python tools/custom_collective_race_repro.py

Performance microbench

AITER_CUSTOM_AR_HOLD_INPUTS=1 \
python tools/custom_collective_microbench.py \
  --graph-iters 100 \
  --label hold-inputs-graph100 \
  --output-json results/custom_collective_microbench/hold-inputs-graph100.json

Test Result

Correctness

for case in ag rs ar1stage; do
  AITER_ROOT=/tmp/up-aiter-main python tools/custom_collective_race_repro.py --cases "$case" --lifetime-reuse --timeout-sec 420
  AITER_CUSTOM_AR_HOLD_INPUTS=0 python tools/custom_collective_race_repro.py --cases "$case" --lifetime-reuse --timeout-sec 420
  AITER_CUSTOM_AR_HOLD_INPUTS=1 python tools/custom_collective_race_repro.py --cases "$case" --lifetime-reuse --timeout-sec 420
done

• ptr denotes input tensor reuse after custom ar.
• hold denotes value of AITER_CUSTOM_AR_HOLD_INPUTS

Case	main	current + hold=0	current + hold=1
ag	FAIL, 1342 / 22937600, ptr 360 / 400	PASS, 0 / 22937600, ptr 360 / 400	PASS, 0 / 22937600, ptr 0 / 400
rs	FAIL, 871 / 5734400, ptr 400 / 400	PASS, 0 / 5734400, ptr 400 / 400	PASS, 0 / 5734400, ptr 0 / 400
ar1stage	FAIL, 16099 / 5734400, ptr 400 / 400	PASS, 0 / 5734400, ptr 400 / 400	PASS, 0 / 5734400, ptr 0 / 400

Performance

python tools/custom_collective_microbench.py --graph-iters 100

Case	main	current + hold=0	current + hold=1
ag	6.928 us ±2.2%	8.621 us ±4.3%	7.321 us ±6.3%
rs	7.546 us ±10.5%	8.639 us ±4.4%	6.896 us ±1.9%
ar1stage	8.852 us ±2.8%	10.247 us ±2.1%	8.879 us ±2.6%

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 4084 --add-label <label>

[Copilot]

Pull request overview

This PR introduces an opt-in performance mode for the custom collective stack by making the final cross-rank end_sync optional and (when enabled) attempting to preserve correctness by keeping inputs alive across launches from Python.

Changes:

• Thread an end_sync boolean from Python → pybind → C++ dispatch → HIP kernels for all-gather, reduce-scatter, and fused AR+RMSNorm variants.
• Add a hold_inputs mode to CustomAllreduce (controlled via AITER_CUSTOM_AR_HOLD_INPUTS) to disable end_sync and retain inputs on the Python side.
• Wire the new behavior into CudaCommunicator construction via an environment variable.

Reviewed changes

Copilot reviewed 5 out of 7 changed files in this pull request and generated no comments.

Show a summary per file

File	Description
csrc/kernels/custom_all_reduce.cu	Add end_sync parameter plumbing into C++ entrypoints and dispatchers.
csrc/include/rocm_ops.hpp	Expose end_sync as a new pybind argument (defaulting to true) for affected ops.
csrc/include/custom_all_reduce.h	Update C++ declarations with end_sync default parameters.
csrc/include/custom_all_reduce.cuh	Make kernel-side end_sync conditional for affected kernels/launchers.
aiter/ops/custom_all_reduce.py	Update Python op signatures/stubs to include end_sync.
aiter/dist/device_communicators/custom_all_reduce.py	Add hold_inputs + input recording, and pass end_sync based on hold_inputs (but currently with correctness gaps).
aiter/dist/device_communicators/communicator_cuda.py	Add env-controlled wiring for hold_inputs into CustomAllreduce.

---

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

[TennyWang1223]

LGTM