[dsidler]

Summary

ctx.all_gather called ctx.store once per destination, and every call did two tl.load(heap_bases + rank) (source and destination base). A load issued after a store cannot be hoisted above it, and waiting for that load also waits for every earlier store, so the compiler put an s_waitcnt vmcnt(0) in front of each store. On gfx942 vmcnt counts stores too, so the remote stores were effectively serialized.

This loads the local and all peer heap bases once before the first store, keeps each peer as a scalar pointer delta (peer_base - local_base), and then issues plain stores. The local copy goes first and peers are written in a per-rank rotated order (rank + 1, rank + 2, ...), the same idea as the Gluon all-gather in #594. No API change.

Assembly (256x256 fp16 tile, _fused_matmul_all_gather_kernel, gfx942)

	before	after
s_waitcnt vmcnt(0) in the epilogue	128	1
stores preceded by a drain	128 of 128	0 of 128
global loads in the epilogue	14 (re-issued per destination)	0 (hoisted above the K loop)
epilogue instructions	1843	1251

Benchmark

benchmark/ops/bench_matmul_all_gather.py, the CU-based matmul_all_gather, 8 ranks, fp16, N=3584, K=8192, two runs each on current main:

M_local	main	this PR
1024	1.39 ms	1.01 ms
4096	4.76 ms	2.30 ms
16384	13.8 ms	7.64 ms

Triton CCL all-gather (iris/ccl/triton/all_gather.py)

The same pattern is applied to persistent_all_gather, which called iris.store per remote destination inside the tile loop. The output base is translated once per peer before the tile loop; the local copy is written first and peers in rotated order, like the Gluon version in #594. persistent_all_gather_partitioned is left as it was: hoisting its single destination base out of the tile loop made it 5-12% slower at 64 or fewer programs.

Assembly of persistent_all_gather (gfx942): global loads 13-14 -> 1, vmcnt(0) drains 9 -> 2-3, stores immediately after a drain 8/8 -> 5/8 (I did not look into the remaining 5).

benchmark/ccl/bench_all_gather.py, persistent, 8 ranks, fp16, N=4096, two runs each. With 64 or more programs the kernel is link-bound and the change is neutral. With a small CU budget, issue serialization limits it and the change helps:

M	comm_sms	main	this PR
4096	8	2.13 ms	1.38 ms
4096	16	1.16 ms	0.89 ms
4096	32	0.89 ms	0.83 ms
4096	64	0.83 ms	0.82 ms
16384	8	8.39 ms	5.19 ms
16384	16	4.59 ms	3.35 ms
16384	32	3.35 ms	3.15 ms
16384	64	3.09 ms	3.04 ms

🤖 Generated with Claude Code

[mawad-amd]

can we do the traffic shaping here too? this still loops 0..world_size-1 with local in the middle, so every rank writes to the same peer at the same time. context.py does local-first + rotated order like #594

[mawad-amd]

local store lost its hint — the old self.store passed hint=(1, tile.block_n) for every dest including self. the remote stores below still get max_contiguous/multiple_of but this one doesnt