fix(custom_all_reduce): prevent peer-read races in collectives and fused AR+RMSNorm

[Richardyu114]

Summary

Backport the applicable custom-collective completion barriers from ROCm/aiter#4082 and ROCm/aiter#4346 to this older AITER branch.

The ROCm#4082 backport adds a final cross-rank synchronization barrier to the four base peer-read kernels present in this tree:

• allgather_naive
• allgather_vec
• allgather_lastdim
• reduce_scatter_first_dim

The ROCm#4346 backport extends the same lifetime guarantee to the three 1-stage fused allreduce + RMSNorm/quant kernels present in this tree:

• allreduce_fusion_kernel_1stage
• allreduce_fusion_kernel_1stage_per_group
• allreduce_fusion_kernel_1stage_mxfp4

There are no Python API, operator signature, model, or tuning-table changes.

Root cause

These custom collectives read other ranks' inputs directly through registered IPC pointers. The existing start_sync<ngpus>() makes the inputs available before the peer-read phase starts, but it does not protect each input buffer's lifetime after a rank finishes its own local work.

Local kernel completion therefore does not imply that every peer has finished reading this rank's input:

rank A finishes locally and returns
    -> graph/overlapped execution reuses or overwrites rank A's input storage
    -> slower rank B is still reading rank A's registered IPC pointer
    -> rank B observes new data and produces a corrupted collective result

The issue is especially visible with CUDA/HIP Graph replay because graph-pool addresses are reused aggressively. TBO and EP load imbalance can amplify rank skew, but they are not the underlying bug.

Observed symptoms included nondeterministic collective corruption, NaN logits, invalid token IDs, and repeated invalid output.

Change

Base custom collectives: upstream ROCm#4082

Add the existing lifetime-only completion barrier before return from each affected kernel:

end_sync<ngpus, true>(sg, self_sg, rank);

This prevents any rank from releasing or reusing its input until all ranks have completed the remote-read phase.

Fused allreduce + RMSNorm: upstream ROCm#4346

The same race existed in the three 1-stage fused kernels. Each one called start_sync<ngpus>(), directly read every rank's registered input, produced the RMSNorm/quantized output, and then returned without a matching completion barrier. A faster rank could therefore start the next invocation and reuse its registered input or signal slot while another rank was still reading it.

Add the same final barrier after the local epilogue/output writes and before kernel return:

// Keep the branch-local epilogue/output code above this point.
end_sync<ngpus, true>(sg, self_sg, rank);

The upstream PR could not be cherry-picked without conflicts because this branch has independently changed the per-group and MXFP4 epilogue/template structure. The manual resolution deliberately keeps the branch-local epilogues and appends only the upstream synchronization semantics. Selecting the whole upstream side of either conflict would incorrectly remove the local epilogue/output operation.

Upstream provenance

• upstream PR: fix: synchronize custom collectives before return ROCm/aiter#4082
• upstream PR head: 4441c1ae078c3dc731f7e2da17e4c50bdf587cf3
• upstream merge commit: 0945298ef5bf2319d97d3326f7b7885ddd9e18d3
• first upstream release containing the fix: v0.1.19.post1

Additional fused-kernel backport:

• upstream PR: Fix: add missing end_sync barrier in fused allreduce+rmsnorm kernel ROCm/aiter#4346
• upstream PR head: 32c6e33cb33017195e5f6e97b537371fd17ef649
• upstream merge commit: 0e475758203b9f43144818357054c48a60f3c7e2
• upstream title: Fix: add missing end_sync barrier in fused allreduce+rmsnorm kernel
• upstream motivation: low-concurrency SGLang requests could produce silent accuracy corruption when fused allreduce + RMSNorm reused registered buffers before every peer had completed the prior invocation

The upstream ROCm#4082 patch covers eight synchronization sites. This fork predates the newer split last-/middle-dimension reduce-scatter variants, so only the four applicable hunks exist here. If those kernels are added later, their matching completion barriers must also be carried over.

Upstream ROCm#4346 changes only one file and seven lines. All three affected fused kernels exist in this branch, but two hunks require the manual context-preserving resolution described above.

Focused isolation

• registered AITER all-gather without the completion barrier reproduced corrupted/NaN output;
• disabling the registered AITER all-gather path was clean;
• forcing the unregistered staging-buffer path was clean;
• registered all-gather with the final barrier completed 20/20 batch-4 graph replay rounds cleanly.

Graph instrumentation observed 48 registered all-gathers per MiMo decode graph but only five unique input pointers. That reuse pattern makes a missing peer-read completion barrier directly hazardous.

End-to-end full-dataset validation

The following completed runs validate the earlier ROCm#4082 base-collective backport. They used enable_aiter_allreduce_fusion=False, so they do not constitute validation of the newly added ROCm#4346 fused-kernel barriers.

Tested with the companion SGLang MiMo MORI/TBO branch on 4 x MI355X, with custom all-reduce enabled, AITER attention/MoE, MORI BF16 dispatch/combine, full decode CUDA Graph, scheduler overlap, and MTP.

Configuration	Accuracy	Successful requests	API/model errors
TBO off	1263/1319 (95.7544%)	1319/1319	0
TBO on	1262/1319 (95.6785%)	1319/1319	0

Both runs completed without detected NaNs, NUL bytes, Unicode replacement characters, illegal control characters, or obvious repeated-output corruption.

Performance and risk

The completion barrier waits for the slowest peer reader and has a measurable micro-kernel cost. Upstream ROCm#4082 reported:

Path	Before	After	Delta
All-gather	6.910 us	8.994 us	+2.084 us / +30.16%
Reduce-scatter	7.069 us	8.363 us	+1.294 us / +18.30%

Every participating rank must launch and reach the same collective. The existing entry synchronization already requires matching rank behavior; this patch extends that requirement through completion.

The three ROCm#4346 barriers apply only when the AITER fused allreduce + RMSNorm or quantized variants are selected, for example with SGLang AITER allreduce fusion enabled. They do not alter MORI dispatch/combine or the RCCL reduce-scatter/all-gather path. Their latency impact has not yet been measured on this branch.

When validating a rebuilt tree, invalidate the stale JIT artifact first:

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

Scope

This patch fixes the AITER custom-collective peer-read lifetime race in both the base all-gather/reduce-scatter kernels and the 1-stage fused allreduce + RMSNorm/quant kernels. Separate SGLang-side MORI/TBO allocator-lifetime and scheduler ordering fixes are carried by the companion SGLang branch and are not part of this AITER change.