Fix SDMA AllReduce corruption across HIP graph replays

[hubertlu-tw]

Motivation

AllreduceSdma could return stale data when one communicator replayed graphs with different message sizes. The reduce-scatter layout used message-dependent rank offsets, so an address written by a CU reduction in one graph could become an SDMA receive slot in another. On gfx950, that role change could expose stale cache lines and corrupt an entire rank shard.

This matters for inference runtimes that capture several batch-size or token-count buckets and reuse the same communicator across them.

Changes

• Use fixed, pack-aligned receive slots derived from the configured maximum message size.
• Reserve a protected reduced-shard slot so AllGather output cannot overwrite data still being read by another rank.
• Track scratch reuse with a rank-uniform overlap threshold and cross-rank handshake.
• Use 64-bit device-side generations for scatter, gather, reuse, and block completion.
• Match peer-published system-scope stores with system-scope polling loads.
• Use a one-shot SDMA path for messages up to 1 MiB and the protected two-shot path for larger messages.
• Make async graph capture enqueue both ReduceScatter and AllGather in start_async; wait_async only performs copy-out and completion.
• Order a different wait_async stream after the start_async stream with a graph-capturable HIP event.
• Preserve the legacy copy_output_to_user=False argument while always returning results through the caller-provided output.
• Add sync and async changing-size graph-replay tests and a public-header compile test.

Validation

All tests below passed on gfx950.

TP	Sync graph replay	Async graph replay	Message sequence
2	10 replays per captured graph	20 alternating replays	128 MiB, 64 MiB, 128 MiB, 1 MiB, 16 KiB, 16 MiB
4	5 replays per captured graph	14 alternating replays	128 MiB, 32 MiB, 128 MiB, 1 MiB, 16 KiB, 16 MiB, 64 MiB
8	5 replays per captured graph	14 alternating replays	128 MiB, 32 MiB, 128 MiB, 1 MiB, 16 KiB, 16 MiB, 64 MiB

The graph tests capture every size before replay, change input values on each operation, and cover one-shot/two-shot transitions plus compact-output reuse.

Additional checks:

• TP=4 eager out-of-place and in-place AllReduce at 1 MiB with BF16, FP16, and uint32.
• TP=4 async BF16 at 16 MiB with copy_output_to_user=False.
• TP=8 async BF16 at 16 MiB with normal copy-out.
• TP=4 async graph replay with start_async on a side stream and wait_async on the capture stream.
• Mixed eager/graph and sync/async correctness at 16 KiB, 1 MiB, 16 MiB, 64 MiB, and 128 MiB.
• Public umbrella-header compilation for gfx950.
• Editable package rebuild.

Reproduction

Run the following commands from the repository root after installing the package:

# Changing-size graph replay
HIP_VISIBLE_DEVICES=4,5 python -m tests.python.ccl.test_allreduce_sdma_graph --world-size 2 --sizes 128M,64M,128M,1M,16K,16M --replays 10
HIP_VISIBLE_DEVICES=4,5 python -m tests.python.ccl.test_allreduce_sdma_async_graph --world-size 2 --sizes 128M,64M,128M,1M,16K,16M --replays 20
HIP_VISIBLE_DEVICES=4,5,6,7 python -m tests.python.ccl.test_allreduce_sdma_graph --world-size 4 --sizes 128M,32M,128M,1M,16K,16M,64M --replays 5
HIP_VISIBLE_DEVICES=4,5,6,7 python -m tests.python.ccl.test_allreduce_sdma_async_graph --world-size 4 --sizes 128M,32M,128M,1M,16K,16M,64M --replays 14
HIP_VISIBLE_DEVICES=4,5,6,7 python -m tests.python.ccl.test_allreduce_sdma_async_graph --world-size 4 --sizes 128M,32M,128M,1M,16K,16M,64M --replays 14 --cross-stream
HIP_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 python -m tests.python.ccl.test_allreduce_sdma_graph --world-size 8 --sizes 128M,32M,128M,1M,16K,16M,64M --replays 5
HIP_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 python -m tests.python.ccl.test_allreduce_sdma_async_graph --world-size 8 --sizes 128M,32M,128M,1M,16K,16M,64M --replays 14

# Eager compatibility
HIP_VISIBLE_DEVICES=4,5,6,7 python tests/python/ccl/test_allreduce.py --world-size 4 --elems 524288 --iterations 3 --warmup 3 --dtype bf16
HIP_VISIBLE_DEVICES=4,5,6,7 python tests/python/ccl/test_allreduce.py --world-size 4 --elems 524288 --iterations 3 --warmup 3 --dtype fp16
HIP_VISIBLE_DEVICES=4,5,6,7 python tests/python/ccl/test_allreduce.py --world-size 4 --elems 262144 --iterations 3 --warmup 3 --dtype uint32
HIP_VISIBLE_DEVICES=4,5,6,7 python tests/python/ccl/test_allreduce_async.py --world-size 4 --elems 8388608 --iterations 3 --warmup 3 --dtype bf16 --no-copy
HIP_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 python tests/python/ccl/test_allreduce_async.py --world-size 8 --elems 8388608 --iterations 3 --warmup 3 --dtype bf16

# Public-header compile check
/opt/rocm/bin/hipcc -std=c++17 --offload-arch=gfx950 -Iinclude -I. -fsyntax-only tests/cpp/collective/test_sdma_allreduce_public_header.hip

Performance

Measurements use BF16, 20 warmups, 100 measured iterations, five repetitions, and rank-averaged device-event timing. The baseline always uses separate ReduceScatter and AllGather kernels. The PR uses the new one-shot kernel through 1 MiB and the protected two-shot path for larger messages. Negative changes are improvements. PR measurements were refreshed after adding cross-stream event ordering. The 1 MiB graph and 64 MiB eager rows use isolated reruns because the full sweeps contained scheduler outliers.

TP=2

Size	Baseline path	PR path	Graph baseline	Graph PR	Change	Eager baseline	Eager PR	Change
16 KiB	two-shot	one-shot	30.62 us	31.21 us	+1.9%	25.25 us	25.08 us	-0.7%
1 MiB	two-shot	one-shot	50.84 us	51.15 us	+0.6%	45.60 us	45.97 us	+0.8%
16 MiB	two-shot	two-shot	333.81 us	336.72 us	+0.9%	327.94 us	330.46 us	+0.8%
64 MiB	two-shot	two-shot	1.239 ms	1.222 ms	-1.3%	1.232 ms	1.218 ms	-1.1%
128 MiB	two-shot	two-shot	2.457 ms	2.414 ms	-1.7%	2.444 ms	2.408 ms	-1.5%

TP=4

Size	Baseline path	PR path	Graph baseline	Graph PR	Change	Eager baseline	Eager PR	Change
16 KiB	two-shot	one-shot	34.29 us	32.15 us	-6.2%	28.35 us	26.06 us	-8.1%
1 MiB	two-shot	one-shot	46.74 us	45.16 us	-3.4%	41.08 us	39.63 us	-3.5%
16 MiB	two-shot	two-shot	191.89 us	194.77 us	+1.5%	186.73 us	189.50 us	+1.5%
64 MiB	two-shot	two-shot	661.10 us	654.05 us	-1.1%	655.19 us	645.81 us	-1.4%
128 MiB	two-shot	two-shot	1.295 ms	1.266 ms	-2.3%	1.281 ms	1.260 ms	-1.6%

TP=8

Size	Baseline path	PR path	Graph baseline	Graph PR	Change	Eager baseline	Eager PR	Change
16 KiB	two-shot	one-shot	39.83 us	36.05 us	-9.5%	36.73 us	30.28 us	-17.6%
1 MiB	two-shot	one-shot	48.30 us	45.26 us	-6.3%	45.20 us	40.53 us	-10.3%
16 MiB	two-shot	two-shot	123.41 us	126.20 us	+2.3%	119.21 us	119.80 us	+0.5%
64 MiB	two-shot	two-shot	376.16 us	371.31 us	-1.3%	369.68 us	363.50 us	-1.7%
128 MiB	two-shot	two-shot	715.91 us	695.11 us	-2.9%	704.51 us	690.91 us	-1.9%

Cross-stream event overhead

The TP=4 async modes were also compared against the saved pre-event PR run using the same benchmark protocol. Graph replay shows no meaningful regression. Eager async adds approximately 3–5 us at 16 MiB and above, consistent with recording the completion event before the host-side wait.

Size	Async eager before	Async eager after	Change	Async graph before	Async graph after	Change
16 KiB	40.12 us	39.96 us	-0.4%	33.88 us	33.02 us	-2.6%
1 MiB	55.42 us	54.90 us	-0.9%	48.51 us	47.63 us	-1.8%
16 MiB	201.39 us	205.44 us	+2.0%	195.10 us	194.75 us	-0.2%
64 MiB	661.37 us	664.71 us	+0.5%	654.29 us	654.03 us	0.0%
128 MiB	1.272 ms	1.277 ms	+0.4%	1.264 ms	1.266 ms	+0.1%

Trade-offs

• Symmetric scratch grows from the maximum message size to approximately (TP + 1) / TP times that size: 192 MiB at TP=2, 160 MiB at TP=4, and 144 MiB at TP=8 for a 128 MiB maximum.
• copy_output_to_user=False remains source-compatible but no longer suppresses copy-out.
• Sync-graph latency improves at 10 of 15 measured points. The largest regression is TP=8 at 16 MiB: +2.79 us (+2.3%).
• Sync-eager latency improves at 11 of 15 measured points. The largest regression is TP=4 at 16 MiB: +2.77 us (+1.5%).
• Cross-stream safety has no material graph-replay cost in the measured TP=4 sweep. Async eager has a small fixed overhead at larger sizes, peaking at +5.06 us (+0.4%) for 128 MiB.

[wuyl1]

Thanks for the great work! wait_async() may race when using a different stream from start_async(). Could we add stream synchronization and a cross-stream test?

[hubertlu-tw]

Thanks for the great work! wait_async() may race when using a different stream from start_async(). Could we add stream synchronization and a cross-stream test?

@wuyl1 Thank you for reviewing the PR. I have refactored the PR and added stream synchronization and a cross-stream test. Notice that I haven't added the tests I created to mori's nightly CI. All the tests were run locally on my MI355x server.