perf: optimize fused AllReduce + RMSNorm (custom_all_reduce)

[ftyghome]

Motivation

This PR makes three small updates to the custom fused AllReduce + RMSNorm kernels:

• use the existing DPP-based block reduce on the one-stage decode path;
• simplify the two-stage stage-1 reduce-scatter path to avoid LDS and block-level sync;
• move the two-stage stage-2 RMSNorm reduction to the shared DPP wave-reduce primitive.

Technical Details

1stage

• decode path. ar_fusion_epilogue_block_reduce now uses multithread_reduce from hip_reduce.h instead of the generic warpReduce implementation. This switches the intra-wave reduction from the ds_bpermute XOR butterfly to the existing DPP tree reduce.

2stage

• reduce-scatter. reduce_scatter_cross_device_store is rewritten as a flat reduce-scatter. Each thread reduces one pack across all input ranks in fp32, downcasts the result to bf16, and writes it to every rank's temporary buffer directly. This removes the previous LDS staging and __syncthreads() from the warp-per-rank implementation. The reduce order over ptrs[0..ngpus-1] is preserved, so the result remains bit-identical to the previous path.

• stage-1 launch config. The stage-1 launch now uses block size 256 instead of 512. This gives small-m decode shapes more parallelism across CUs.

• stage-2 RMSNorm reduce. local_device_load_rmsnorm now uses wave_reduce from hip_reduce.h for the RMSNorm square-sum reduction. This keeps the reduction implementation consistent with the shared DPP primitive. In testing, this change alone is performance-neutral.

This PR does not change the one-stage/two-stage selection heuristic. The thresholds discussed in #3458 can therefore still be used as-is.

Test Plan

Test environment:

• GPU: MI355X
• ROCm: 7.2.3

Benchmark:

The benchmark script (test_fused_ar_rmsnorm_perf.py is inspired by the script posted by @TennyWang1223 in #3458, with a few changes to improve stability.

To reduce host launch overhead and cross-rank rendezvous jitter, 2000 op calls are captured into one CUDA graph, and a single replay is timed with CUDA events. The reported latency is the minimum over 80 replays, taking the max latency across ranks. GPU clocks were pinned with:

rocm-smi --setperfdeterminism 2400

Reproduce commands:

TP_LIST={4,8} AITER_AR_1STAGE={0,1} PYTHONPATH=$PWD \
  python <script_path>

Restore GPU clocks after benchmarking:

for g in 0 1 2 3 4 5 6 7; do
  rocm-smi --setperflevel auto -d $g
done

Test Results

For the CDF plots below, curves closer to the upper-left corner indicate better performance.

One-stage path

• TP4

• TP8

m	TP4 base_us	TP4 opt_us	TP4 speedup	TP8 base_us	TP8 opt_us	TP8 speedup
4	5.78	5.72	1.010x	7.40	7.28	1.016x
8	6.39	6.33	1.009x	7.99	7.91	1.010x
16	8.21	8.17	1.005x	9.83	9.73	1.010x
32	12.07	12.02	1.004x	14.64	14.54	1.007x
64	20.55	20.49	1.003x	25.68	25.59	1.004x

Two-stage path

• TP4

• TP8

m	TP4 base_us	TP4 opt_us	TP4 speedup	TP8 base_us	TP8 opt_us	TP8 speedup
4	8.73	8.62	1.013x	9.65	9.63	1.002x
8	9.15	9.04	1.012x	9.87	10.14	0.973x
16	9.92	9.96	0.996x	10.18	10.49	0.970x
32	11.93	12.01	0.993x	10.98	11.18	0.982x
64	16.32	16.31	1.001x	13.57	13.23	1.026x
128	26.21	25.72	1.019x	18.82	18.76	1.003x

Under the selection thresholds from #3458 (one-stage for m <= 32 at TP4 and m <= 16 at TP8), the latency-critical small-m decode shapes are served by the one-stage path, which this PR speeds up. The shapes where the two-stage path regresses largely fall in that one-stage range, so the practical impact is minimal; on the larger shapes that do use the two-stage path the change is a 2.5% net gain.

This PR does not change the 1stage/2stage selection heuristic, so the threshold from #3458 remains applicable.

[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 3817 --add-label <label>

[Copilot]

Pull request overview

This PR optimizes the fused custom AllReduce + RMSNorm implementation by switching several intra-wave and cross-rank reduction steps to shared DPP-based reduction primitives and by simplifying the 2-stage reduce-scatter path to remove LDS staging and block-level synchronization.

Changes:

• Switch 1-stage decode epilogue block reduction to multithread_reduce (DPP-based) from hip_reduce.h.
• Rewrite 2-stage stage-1 reduce_scatter_cross_device_store as a flat per-pack reduce (no LDS / __syncthreads()).
• Switch 2-stage stage-2 RMSNorm square-sum reduction to the shared wave_reduce primitive and adjust stage-1 launch config to 256 threads.

---

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

[ftyghome]

This is a follow-up to the previous closed PR: #3462.

That PR became outdated and cannot be reopened, so we created this new PR instead. Since #3458 has been merged, the 1stage improvements here can now be reflected in the actual model inference path. We also found a more stable way to measure the per-op AllReduce kernel latency, which should make the performance results easier to validate.

[ftyghome]

Hi all,

The failing CI seems can be resolved by rebase this branch with upstream.

Can anyone help rerun the CI? Thanks.

[ftyghome]

Hi all,

The CI is failing due to an env error. Can anyone help rerun this?

Thanks!

[ftyghome]

Closing this — we're not planning to pursue this further for now. Thanks for the reviews!