[Bugfix] Fix IPC publication ordering in custom all-reduce

[Fangzhou-Ai]

Summary

• Publish eager fixed-buffer copies and intermediate all-reduce writes with system-scoped release stores.
• Acquire system-scoped visibility after all peer ready signals are observed, before reading peer IPC buffers.
• Apply the same ordering to the generic and gfx1250 synchronization helpers while keeping final read-completion barriers relaxed.
• Add a TP8 stress test for repeated eager-buffer reuse and graph replay under rotating rank skew.

Root cause

The eager path copies input into an IPC-registered fixed buffer immediately before launching the collective. The ready flag was system-scoped but relaxed, and peers polled it with device scope. That did not establish a system-scoped happens-before edge from the producer copy to peer buffer reads. The same gap existed for intermediate writes consumed after a non-final barrier.

Duplicate-work check

I reviewed #4082 and #4084. Those changes protect buffer lifetime after peer reads by adding terminal synchronization or retaining caller tensors. This PR instead publishes data before peer reads. It intentionally excludes their terminal-barrier changes, so the fixes are complementary rather than duplicate. Searches for custom all-reduce IPC publication/order fixes found no other open PR.

Validation

Tested on an 8x gfx950 host at current main commit dedc19d32bb31ec134ad0388f878e4cd7a83af16.

HIP_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 PYTHONPATH=$PWD CK_DIR=$PWD/3rdparty/composable_kernel \
  /home/amd/vllm/.venv/bin/python \
  op_tests/multigpu_tests/test_custom_allreduce_ipc_ordering.py \
  --iters 100 --tp-size 8

Passed on all eight ranks, including graph capture/replay. This is a stress regression guard rather than a deterministic reproducer; an unpatched-main 20-iteration control also completed.

The fused all-reduce/RMS integration sweep passed 11,200/11,200 comparisons across M=1,4,8,16,32,64,4096 (100 iterations each), with maximum absolute error 0.125. A graph-free DeepSeek-V4-Pro GSM8K run completed 1,319 samples at 95.906% strict / 95.830% flexible accuracy; its server handled 1,351 successful HTTP responses with no GPU fault.

The stock TP8 graph benchmark was also run on main and this patch:

BF16 shape	main min-max (us)	patched min-max (us)
(2, 7168)	5.45-6.69	6.25-7.50
(128, 8192)	14.53-15.11	15.74-16.29

HIP_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 PYTHONPATH=$PWD CK_DIR=$PWD/3rdparty/composable_kernel \
  /home/amd/vllm/.venv/bin/python \
  op_tests/multigpu_tests/test_custom_allreduce.py \
  -d bf16 -s 2,7168 -g true
HIP_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 PYTHONPATH=$PWD CK_DIR=$PWD/3rdparty/composable_kernel \
  /home/amd/vllm/.venv/bin/python \
  op_tests/multigpu_tests/test_custom_allreduce.py \
  -d bf16 -s 128,8192 -g true

Static checks:

uvx --from black==26.3.0 black --check op_tests/multigpu_tests/test_custom_allreduce_ipc_ordering.py
uvx --from ruff==0.15.7 ruff check op_tests/multigpu_tests/test_custom_allreduce_ipc_ordering.py
/home/amd/vllm/.venv/bin/python -m py_compile op_tests/multigpu_tests/test_custom_allreduce_ipc_ordering.py
uvx --from clang-format==18.1.8 clang-format --dry-run --Werror --lines=154:236 csrc/include/custom_all_reduce.cuh
uvx --from clang-format==18.1.8 clang-format --dry-run --Werror --lines=308:360 csrc/include/custom_all_reduce_gfx1250.cuh
git diff --check origin/main...HEAD

The gfx1250 helper mirrors the same memory-order contract but was not runtime-tested on gfx1250 hardware.

AI assistance

OpenAI Codex was used to audit related PRs, prepare the patch and regression test, run validation, and draft this description. The human submitter is responsible for reviewing every changed line and defending the change end-to-end.

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

[TennyWang1223]

PR Review

Recommendation: do not merge. Based on an ISA-level assembly comparison and cycle-accurate microbenchmarks, the cross-rank publication problem this PR claims to fix is not actually addressed at the instruction level, and the added test cannot demonstrate that the current main has the problem. I suggest keeping main as-is.

Full reasoning below.

1. Verdict on the three synchronization functions (gfx950, isolated compile)

Function	PR change	Assembly evidence	Verdict
end_sync<final=true>	none	byte-for-byte identical to main	no effect, meaningless
start_sync	store relaxed→release, load + acquire fence	main has no buffer_wbl2/buffer_inv here; PR adds both	redundant (see §2)
end_sync<final=false>	buffer_inv moved inside→outside loop, scope sc1→sc0 sc1	buffer_wbl2 present on both sides; only the inv position and scope differ	unnecessary (see §3)

2. The start_sync change is redundant

start_sync sits at the very beginning of the collective. Its execution flow is:
input copy → kernel launch → start_sync → remote load → compute → end_sync<false>.

• Producer side release→buffer_wbl2: this instruction writes back the current grid's own dirty L2 lines. But the eager input copy is a separate, already-retired memcpy grid — it is not in this grid's writeback set. Moreover, there are no memory writes in this kernel before start_sync. So this writeback flushes nothing meaningful and cannot publish the copy it claims to publish.
• Consumer side acquire fence→buffer_inv: the previous iteration's end_sync<false> buffer_inv already invalidated the local cache. There are no IPC-buffer-related dirty lines that need to be invalidated again.

A probe confirms buffer_wbl2 is triggered entirely by __ATOMIC_RELEASE (independent of whether prior dirty writes exist), and main's start_sync uses relaxed so it has no such instruction. Conclusion: the two cache instructions this PR adds to start_sync are superfluous.

3. end_sync<final=false>: sc1 (device scope) is enough, sc0 sc1 (system scope) is not needed

The consumer side only polls the flag in its own local tmp buffer. Reading the correct flag already means the peer has finished writing; from that point on the consumer proceeds locally and its behavior does not depend on any other rank. Establishing system-scoped ordering here therefore buys nothing — device scope (sc1) is sufficient, and sc0 sc1 only widens the invalidation scope, which can only lengthen latency, never shorten it.

Moving the buffer_inv from inside the poll loop to outside it makes no measurable difference to performance (confirmed by cycle-accurate microbenchmark on gfx942). So from the HIP source there is no justification either to bump the scope to system or to move the fence outside the while loop — sc1 inside the poll, as in main, is both correct and adequate.

4. The accompanying test script cannot prove main is broken

test_custom_allreduce_ipc_ordering.py is a regression guard, not a deterministic reproducer:

1. The author states in the description that the "unpatched-main 20-iteration control also completed" — main passes too. A test that cannot make main fail cannot prove main has a bug.
2. The assertion reads output for comparison only after torch.cuda.synchronize(). Weak-memory-order / visibility bugs are transient; a full device sync collapses any intermediate state, so the assertion window opens after the only moment the bug could manifest.
3. The time.sleep(skew_s) introduces host-side launch skew, which targets the "read before peer launch" scenario that the start_sync barrier is responsible for blocking — a different mechanism from "data publication / cache visibility," and the skew is absorbed by the barrier anyway.
4. The fill values are rotating small-integer linear combinations; a stale value and the expected value can numerically coincide, and the test lacks a "unique per-iteration sentinel + active clobber recheck" to sharpen stale-read detection.

5. Note on gfx1250

The assembly above was generated on gfx9. There is no need to modify start_sync / end_sync for gfx1250: the kernel implementation there will be reworked and will no longer use these two device-side synchronization functions. That update is coming soon, so I did not bother comparing the gfx1250 assembly.

Summary

The change does not hit the problem it claims to address (release targets the wrong object; there is nothing to publish before start_sync), and the test cannot reproduce the problem either. For end_sync<false>, sc1 (device scope) inside the poll loop — exactly what main already does — is both correct and sufficient; sc0 sc1 (system scope) adds latency risk with no measurable benefit, and moving the fence outside the loop is not warranted. main currently runs stably, and it is not worth introducing system-scope overhead and extra cache instructions for this. I recommend closing this PR or keeping main.