[vishalshetagar]

The Gluon all-gather kernel computes translated destination bases before the tile loop and reuses them for every tile. In the reviewer's MI355X evaluation, this removes a per-peer, per-tile vmcnt(0) drain, allowing more memory operations to overlap. It also computes output_ptrs once per tile and preserves the logical destination order and local .wt store policy.

Related to #591. The reviewer measured a 4.47% median latency reduction at eight ranks; four-rank performance was effectively unchanged. The published commit's CCL CI job reports success at 1/2/4/8 ranks. These external results apply to commit 0d94e89, before the minor review edits described below.

The patch also adds partial-tile/multiple-iteration regression cases, retains random input data with exact comparisons, covers the Gluon variant in the existing subgroup tests, and adds an explicit gluon benchmark variant. The existing persistent benchmark selects the Triton path, so it cannot measure this change.

The review edits use gl.static_range and gl.cast consistently and set the Gluon benchmark tile explicitly to 8x256. The reviewer confirmed that the main benchmark used this tile on both revisions: the launcher remapped the default Config values of 32x64 to 8x256, and the compiled kernels confirmed that layout.

Reviewer performance evaluation

Results supplied by @nirvedhmeshram, not measured on the author's laptop.

• MI355X (gfx950), ROCm 7.2.2, torch 2.14.1+rocm7.2, Python 3.12; Triton built at 1369c18a in a fresh environment.
• Both revisions used identical tests and benchmark harness, copied from the PR into the baseline. Only iris/ccl/gluon/all_gather.py differed.
• variant=gluon, 8x256 tile. 54 configurations: M=1024/4096/16384, N=4096/8192, fp16/bf16/fp32, comm_sms=64/128/256; 25 warmup and 100 timed iterations each.
• Three alternating PR/base runs; each configuration reports the median of its three runs. Eight-rank runs were serialized. Heap sizes were 64 GiB per rank at four ranks and 128 GiB at eight ranks because allocations accumulate across configurations.

Latency delta is negative when the PR is faster:

Slice	4 ranks	8 ranks
All 54 configurations, median	-0.04%	-4.47%
Minimum / maximum	-0.69% / +0.23%	-15.42% / +5.97%
M=1024, N=4096	-0.27%	-6.01%
M=1024, N=8192	-0.04%	-7.60%
M=4096, N=4096	-0.03%	-4.51%
M=4096, N=8192	-0.09%	-4.42%
M=16384, N=4096	+0.04%	-2.77%
M=16384, N=8192	-0.03%	-3.17%
comm_sms 64 / 128 / 256	-0.01 / -0.04 / -0.08%	-4.13 / -5.15 / -5.55%

The eight-rank mean delta was -4.85%. Low-noise examples include 4096x8192 fp32 at comm_sms=128, improving from 3.0355 ms to 2.7193 ms (-10.4%), and regressions at 16384x8192, comm_sms=64: fp16 from 5.6299 ms to 5.7623 ms (+2.4%), and bf16 from 5.6188 ms to 5.7215 ms (+1.8%).

The reviewer considers these regressions nonblocking. Follow-up measurements suggest that configurations already near link saturation benefit less. A loss of peak link efficiency from more peer stores in flight is an unverified explanation requiring fabric counters. Timing peer groups of 2-4 stores per wait remains a possible follow-up; this PR does not add grouping based only on register measurements.

The follow-up tile sweep actually used 16x64, 16x128, 8x256 and 32x128. Its row originally labeled 32x64 also hit the launcher's remap. Although 8x256 and 32x64 both contain 2048 elements, their contiguous FP16 row spans differ (512 B versus 128 B), so that row is not a controlled test of burst size alone. The reviewer reports that this correction does not change the overall conclusion.

The launcher currently remaps even explicitly supplied block_size_m=32, block_size_n=64, because it matches their values rather than whether they were omitted. A true 32x64 tile therefore cannot be selected through this Config launch path. Distinguishing defaults with None or a sentinel is a separate configuration follow-up; this PR makes the measured 8x256 tile explicit.

Multi-GPU CI

The CCL job for commit 0d94e89 and its step running CCL tests at 1/2/4/8 ranks report success. The Gluon all-gather tests compare against torch.distributed.all_gather_into_tensor. Detailed job logs were access-restricted during this review, so individual pass/skip counts have not been independently checked. The review edits are pushed as 8e87511c; their Iris Tests run is queued at the time of this update, with no new multi-GPU pass claimed yet.

Compiler inspection and grouping

Using Triton 3.8.0, I compared the original kernel, the hoisted version, and experimental INLINE=4/8 variants for gfx942, gfx950 and gfx1250. These are offline compiler results, not measurements on those GPUs or on the Dockerfiles' exact compiler revision.

• The original range(world_size) does not always unroll. For generic M/N on gfx942/gfx950, the 8x256 cases unroll at 2/4/8 ranks, while 32x256 retains the destination loop at 4/8 ranks. Alignment specialization and target architecture change the result.
• The patched static destination loop has no destination-loop backedge in the inspected ISA. In the 24 main hoisted IR variants, all heap-base loads precede the tile loop; the only load inside that loop is the input tile.
• The grouping experiment constructs up to four/eight destination pointers and then issues that group's stores, with static loops. It did not consistently improve register usage, so it is not included in the production patch.

Representative VGPR counts, eight ranks, FP16, four warps:

Target	Tile	M/N specialization	Original	Hoisted	INLINE=4	INLINE=8
gfx942	8x256	generic	38	36	36	36
gfx942	32x256	generic	152	128	128	128
gfx942	8x256	divisible by 16	12	24	63	63
gfx942	32x256	divisible by 16	164	100	162	162
gfx1250	8x256	divisible by 16	20	44	44	43

The 112 offline compilations reported zero VGPR/SGPR spills. That does not extend to larger tiles: on the local gfx1032 run, for example, four-rank FP16 with 32x512 tiles reported 100 VGPR spills before and 127 after. Hoisting increases pointer lifetimes; the MI355X timings above establish the measured tradeoff for the 8x256 configuration, not for every tile size.

I also checked gluon/all_to_all.py: its ctx.store calls translate addresses inside the tile/destination/row loops and have the same heap-base-load pattern. This PR leaves that collective unchanged.

Local validation

Hardware: one RX 6700S under WSL with an experimental RDNA2 bridge. PyTorch 2.13.0+rocm7.1, HIP 7.1.52802, Triton 3.8.0. The PyTorch tests use gfx1030 compatibility mode and disabled profiler registration. The wheel's Triton 3.7.1 lacks hip_utils.num_threads, already used by the original kernel, so the test environment loads Triton 3.8.0 separately.

• 15 repository Gluon all-gather cases passed, including 8192x8192 inputs and the new edge cases, against torch.distributed.all_gather_into_tensor. An external wrapper initialized a one-rank Gloo group; the test bodies and kernel were unchanged.
• 180 supplemental kernel cases passed, 90 original and 90 patched, totaling 1,320 GPU launches through PyTorch's HIP runtime. Independent allocations on the same GPU, launched sequentially, represented logical ranks and strided subgroups. Exact NumPy-reference/full-heap comparisons covered outputs, padding, guards and inactive heaps. These checks do not validate inter-GPU transport or concurrent ranks.
• ruff check . --fix, ruff format . (309 files unchanged), bytecode compilation and git diff --check passed. Both new Gluon subgroup variants collect locally; see the multi-GPU CI result above for external coverage.

The normal two-rank runner fails before collection with invalid device ordinal, since device 1 does not exist. RCCL initialization also failed locally before reaching the kernel; the successful repository tests above used Gloo.

The review edits in 8e87511c were revalidated locally: all 15 one-rank repository GPU cases, 180 supplemental logical-heap cases (1,320 launches), 25 host tests, and nine host-only benchmark configuration checks passed. Lint, formatting, benchmark imports and whitespace checks passed. Comparing 0d94e89 with the API edits across 18 offline specializations (gfx942/gfx950/gfx1250, 2/4/8 ranks, 8x256/32x256 FP16 tiles) produced matching instruction streams excluding debug directives/comments and unchanged register/spill counts. These checks do not constitute new multi-GPU timing results.

Reproducing distributed validation

The commands below retain the requested correctness-before-timing workflow. Reviewer performance results and remote CI are reported above; the local single-GPU checks do not replace distributed validation.

Suggested validation commands

Use the same updated tests and benchmark harness for the baseline (cfcb96c9bd3d3a4feeb03623f89fbbf1bb99f843) and this branch, in the project's intended ROCm/Triton environment. In particular, copy the updated benchmark to the baseline checkout so variant=gluon measures the same kernel family on both revisions.

for n in 2 4 8; do
  python tests/run_tests_distributed.py tests/ccl/test_all_gather_gluon.py --num_ranks "$n" -v || exit 1
done
for n in 4 8; do
  python tests/run_tests_distributed.py tests/ccl/test_process_groups.py --num_ranks "$n" -k 'test_all_gather_with_groups and gluon' -v || exit 1
done

After correctness passes in both checkouts, run matched configurations with distinct output paths:

python -m benchmark.ccl.bench_all_gather \
  --axis_variant=gluon --axis_backend=iris --axis_num_ranks=2,4,8 \
  --axis_M=1024,4096,16384 --axis_N=4096,8192 \
  --axis_dtype=fp16,bf16,fp32 --axis_comm_sms=64,128,256 \
  --benchmark_format=json --benchmark_out=/absolute/path/to/before-or-after.json

Use the same GPU topology, software and power settings; repeat measurements to assess variation. The benchmark defaults to a 16 GiB heap per rank, so adjust --heap_size for the selected tensors/hardware. Report latency/bandwidth alongside VGPR/SGPR counts and spills.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

[Copilot]

Copilot encountered an error and was unable to review this pull request. You can try again by re-requesting a review.

Note

This error may be related to your runner configuration. You can now configure runners for Copilot code review separately from Copilot cloud agent by creating a copilot-code-review.yml file with your setup steps. Read the docs for details.

[nirvedhmeshram]

LGTM, just some minor comments to address, I was able to do the perf evaluation requested in the PR description (see below) and we can replace the language in the PR description to make the following stronger claim.

"Removes a per-peer, per-tile vmcnt(0) drain" is a stronger and more durable justification than "fewer heap-base loads".

Also, feel free to include the perf numbers below in the PR description,

Environment

• Hardware and software: MI355X (gfx950); ROCm 7.2.2, torch 2.14.1+rocm7.2, Python 3.12.
• Triton was built from source at 1369c18a (main as of 2026-10-16, the PR base date) in the fresh venv /tmp/iris_venv.
• Both checkouts ran identical tests and benchmark harness. The PR's tests/ccl/* and benchmark/ccl/bench_all_gather.py were copied into the baseline worktree, so the two differ only in iris/ccl/gluon/all_gather.py.

Performance (variant=gluon, default Config tile 32x64)

• Sweep: M ∈ {1024, 4096, 16384}, N ∈ {4096, 8192}, fp16/bf16/fp32, comm_sms ∈ {64, 128, 256}. That is 54 configs, each 25 warmup + 100 timed iterations.
• The runs alternated PR, base, PR, base, PR, base, and each config reports the median over the 3 runs. The 8-rank runs went strictly one after another.
• The heap was 64 GiB per rank at 4 ranks and 128 GiB at 8, because the harness never frees between configs.

Slice (median delta, negative = PR faster)	4 ranks	8 ranks
all 54 configs	-0.04%	-4.47% (mean -4.85%)
min / max	-0.69% / +0.23%	-15.42% / +5.97%
M=1024, N=4096	-0.27%	-6.01%
M=1024, N=8192	-0.04%	-7.60%
M=4096, N=4096	-0.03%	-4.51%
M=4096, N=8192	-0.09%	-4.42%
M=16384, N=4096	+0.04%	-2.77%
M=16384, N=8192	-0.03%	-3.17%
comm_sms 64 / 128 / 256	-0.01 / -0.04 / -0.08%	-4.13 / -5.15 / -5.55%

Some 8-rank configs where the noise (max-min spread over 3 runs) is under 1%:

Config	base ms	PR ms	delta	noise
1024x8192 fp32, sms 128	0.8785	0.7995	-9.0%	0.4%
4096x8192 fp32, sms 128	3.0355	2.7193	-10.4%	0.3%
4096x4096 fp32, sms 256	1.6451	1.5015	-8.7%	0.4%
16384x8192 fp32, sms 128	11.3869	10.6159	-6.8%	0.2%
4096x8192 bf16, sms 256	1.4785	1.4126	-4.5%	0.4%
16384x8192 fp16, sms 64	5.6299	5.7623	+2.4%	0.1%
16384x8192 bf16, sms 64	5.6188	5.7215	+1.8%	0.7%

About the few regressions, I dont think those should be blockers, here is what the investigation (by agent) revealed about that.

• In the main sweep the only clean regression is the largest 16-bit shape at comm_sms=64, at about +2%.
• I ran two follow-up experiments at 8 ranks, M=16384, N=8192, 3 alternations each:

Experiment	Result
comm_sms 16 / 32 / 64 (fp16, fp32; M 16384 and 32768)	The PR is 5-10% faster at 16 and 32. The +1.5% / +2.4% fp16 regression appears only at 64. So "too few CTAs" is ruled out.
tiles 16x64, 16x128, 32x64, 32x128 (512 B to 4 KB per peer per wave), comm_sms 32 and 64	Burst size doesn't predict the regression. fp32 16x64 (1 KB) regresses by +1.6-1.8% with noise ≤0.2%, but fp16 16x128 (1 KB) and 16x64 (512 B) don't.

• The variable that fits is whether the baseline was already link-saturated.
  • The 4 configs where the baseline already reached ~330 GB/s or more come out -0.8% to +2.3%.
  • The other 12 range from -4% to -27%, except one noisy +2.5%.
  • The PR pulls nearly every config up to the same ~315-350 GB/s ceiling. Peak per dtype barely moves: fp32 356 -> 350 GB/s, fp16 333 -> 336 GB/s.
• Unverified interpretation (would need fabric counters). At saturation, the baseline's one-peer-at-a-time pattern is the stagger the kernel comment intends. The PR's all-peers-at-once pattern loses roughly 1-2% of peak link efficiency.
• Possible follow-up. Group peers: store to 2-4 peers per s_waitcnt. The author's INLINE=4 experiment is close to this, but they only evaluated it for registers, not for timing.

[vishalshetagar]

Thanks for running the evaluation and looking into the regressions. I’ll address the two comments and update the PR description with your measurements and the vmcnt(0) explanation.

One detail on the benchmark: the launcher already maps the default 32×64 tile to 8×256 in both revisions. Did your runs use that launch path? I’ll make the tile explicit in the benchmark and want to ensure the results are labelled correctly.

[nirvedhmeshram]

Yes, we used that launch path, and you're right about the remap. Our benchmark runs went through the default Config(), so launch() turned the default 32×64 into 8×256 on both revisions. The compiled kernels confirm it: the 8-elements-per-thread kernels contain only 256/255 constants and no 64/63 ones, on both the PR and the baseline. So every number in the review ran 8×256. We labeled it "32×64" because we read the Config default instead of the launcher.

Please label the benchmark results as 8×256. Making the tile explicit is the right fix, but note that an explicit block_size_m=32, block_size_n=64 also triggers the remap, because launch() matches on the values and can't tell them from the defaults. So 32×64 can't be benchmarked through Config at all today. Maybe use a sentinel default or None.

One correction to our follow-up experiment: the tile sweep used 16×64, 16×128, "32×64" and 32×128. The "32×64" row was really 8×256, because our explicit 32/64 also hit the remap. The other three ran as labelled. This doesn't change the conclusion, since 8×256 and 32×64 both have 2048 elements and so the same bytes per peer per wave. But 8×256 rows are 512 B contiguous (fp16) where 32×64 rows are 128 B, so that row isn't a pure test of burst size.

[vishalshetagar]

Thanks for confirming. I’ll label the results as 8×256 and correct the follow-up sweep description, including the contiguous-layout caveat. The explicit 8×256 benchmark configuration is already pushed. I’ll leave the Config default/remapping behavior for a separate follow-up.

[nirvedhmeshram]

LGTM, after updating the PR description.

[vishalshetagar]

Updated the PR description with the confirmed 8×256 tile and your performance results. Thanks for the review and validation!