perf: raise fused ar 1stage limit to cover conc64

[ftyghome]

Motivation

Currently, the gating for the all-reduce 1-stage (fused path) is total_bytes <= 128K. Under inference paths like Kimi K2.5, only bs4 goes through this path. However, actual testing shows that the 1-stage path yields performance gains across bs8, 16, 32, and 64 (with a regression at bs128), and the 1-stage kernel itself can cover bs64 and above without modification. This PR therefore raises the limit to 1M to cover the bs64 path, while ensuring the kernel can handle this size.

Test Plan

Profile Kimi K2.5 model inference and compare the performance of 1-stage and 2-stage paths.
(tested bs64 and bs128 to locate the boundary.)

Test Result

fused	non-fused (sum of 2 stages)	non-fused (reduce_scatter_cross_device)	non-fused (local_device_load)
23618	24850	18527	6323
41213	22272	15701	6571

[Copilot]

Copilot wasn't able to review any files in this pull request.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

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

[TennyWang1223]

Hi, I ran the same shapes locally and got the opposite result — pushing everything through the 1-stage (fused) kernel is consistently slower than the original 2-stage path on my setup. Could you double-check the test config /
environment?

Environment

• GPU: MI300
• ROCm: 7.2.4
• dtype: bf16
• TP: 8

Fused (1-stage, this PR)

shape	latency_us
(8, 7168)	47.738
(16, 7168)	42.593
(32, 7168)	50.786
(64, 7168)	79.353

Original (2-stage gating)

shape	latency_us
(8, 7168)	37.144
(16, 7168)	38.788
(32, 7168)	40.134
(64, 7168)	55.674

The regression shows up across all four shapes, not only at bs128. Could you share the exact test command, GPU SKU, and ROCm/RCCL versions you used? Want to make sure we're comparing apples to apples before raising the threshold to
1M.

[ftyghome]

Hi TennyWang1223,

Thanks for checking this. I was able to reproduce a similar regression in the standalone microbench. However, when I profiled the actual decode path, the result looked a bit different from the isolated op test.

My environment:

• GPU: MI350X VF
• ROCm Version: 7.2.3

For the microbench, I used:

AITER_AR_1STAGE=1 python op_tests/multigpu_tests/test_fused_ar_rms.py \
    -t {4,8} -d bf16 -s 4,7168 8,7168 16,7168 32,7168 64,7168 --test fused -g 1

The standalone result is consistent with what you observed. At TP4, the fused path starts to lose after conc16. At TP8, it starts to lose after conc8.

TP4, us, min / max:

shape	fused	non-fused
(4, 7168)	6.10 / 6.46	9.02 / 9.37
(8, 7168)	6.33 / 6.69	9.09 / 9.40
(16, 7168)	8.15 / 8.50	9.77 / 10.14
(32, 7168)	12.02 / 12.08	11.68 / 12.13
(64, 7168)	20.25 / 20.57	16.04 / 16.46

TP8, us, min / max:

shape	fused	non-fused
(4, 7168)	7.78 / 8.05	9.80 / 10.01
(8, 7168)	8.57 / 19.93	10.38 / 19.39
(16, 7168)	10.74 / 11.32	10.17 / 10.54
(32, 7168)	16.85 / 17.41	11.17 / 11.75
(64, 7168)	29.92 / 30.25	13.49 / 15.61

But in the actual Kimi K2.5 MXFP4 decode path on the ATOM server, the fused 1-stage path is still faster, or roughly at parity, up to conc64. I measured this with torch profiler and used the real per-kernel time for each path. The numbers below are p50 in ns.

conc	fused	non-fused, sum of 2 stages	reduce_scatter_cross_device	local_device_load
4	8559	15518	9719	5799
8	9239	15758	9959	5799
16	11119	16958	11159	5799
32	14479	17598	11799	5799
64	22999	22839	16920	5919

p90 / p95, ns

p90:

conc	fused	non-fused, sum of 2 stages	reduce_scatter_cross_device	local_device_load
4	10559	18238	12039	6199
8	11119	18958	12679	6279
16	14839	22718	16479	6239
32	17999	21959	15760	6199
64	30000	30998	24439	6559

p95:

conc	fused	non-fused, sum of 2 stages	reduce_scatter_cross_device	local_device_load
4	11199	19358	12959	6399
8	11959	19999	13440	6559
16	16319	24079	17520	6559
32	19279	23558	16999	6559
64	31839	33239	26360	6879

Thanks again for checking this. Please feel free to let me know if there is any other data or experiment that would be helpful from my side.

[TennyWang1223]

Hi @ftyghome ,

Thanks for the detailed numbers. The gap between the microbench and the real decode path is interesting — before changing the threshold, I'd like to look at the kernel timeline directly.

Could you capture two torch.profiler traces from the Kimi K2.5 MXFP4 decode path on the ATOM server and attach them to this PR?

• trace_fused.json — AITER_AR_1STAGE=1 (1-stage path)
• trace_nonfused.json — default gating (2-stage path)

Ideally:

• Same conc (e.g. 32 and 64, where the two paths diverge most)
• Same TP / dtype / model config you used for the table above
• A few warmup iters, then ~20 profiled steps
• record_shapes=True, with_stack=True if possible

That way I can line up the per-kernel times and the surrounding ops (RMSNorm, comm overlap, etc.) and see why the standalone op test and the end-to-end path disagree.

Thanks!

[ftyghome]

Hi @TennyWang1223 ,

Thanks for following up. I have uploaded the torch profiler traces (since the files are large, I uploaded them in my Cloudflare R2 bucket):

trace_fused_conc32.json.gz
trace_fused_conc64.json.gz
trace_nonfused_conc32.json.gz
trace_nonfused_conc64.json.gz

I used the following SQL snippet in Perfetto to extract the performance data:

WITH params AS (
  SELECT 32 AS sz
),
decode_name AS (
  SELECT 'decode[bs=' || sz || ' tok=' || sz || ' d=' || sz || ']' AS dn
  FROM params
),
d AS (
  SELECT
    CASE
      WHEN child.name LIKE '%allreduce_fusion_kernel%'           THEN 'allreduce'
      WHEN child.name LIKE '%reduce_scatter_cross_device_store%' THEN 'reduce_scatter'
      WHEN child.name LIKE '%local_device_load_rmsnorm%'         THEN 'load_rmsnorm'
    END AS kind,
    child.dur
  FROM slice AS parent,
       descendant_slice(parent.id) AS child
  WHERE parent.name = (SELECT dn FROM decode_name)
    AND (child.name LIKE '%allreduce_fusion_kernel%'
      OR child.name LIKE '%reduce_scatter_cross_device_store%'
      OR child.name LIKE '%local_device_load_rmsnorm%')
),
ranked AS (
  SELECT
    kind, dur,
    ROW_NUMBER() OVER (PARTITION BY kind ORDER BY dur) AS rn,
    COUNT(*)     OVER (PARTITION BY kind)              AS n
  FROM d
)
SELECT
  kind,
  n                                                   AS cnt,
  MAX(CASE WHEN rn = (n + 1) / 2         THEN dur END) AS median_ns,
  MAX(CASE WHEN rn = (90 * n + 99) / 100 THEN dur END) AS p90_ns,
  MAX(CASE WHEN rn = (95 * n + 99) / 100 THEN dur END) AS p95_ns
FROM ranked
GROUP BY kind, n
ORDER BY kind;

Thanks! Let me know if I can help with anything else from my side.

[ftyghome]

Hi @TennyWang1223,

We added a new set of experiments on MI355X with ROCm 7.2.3 to further validate the gating threshold.

The updated result is consistent with our previous TP4 observation: for TP4, the 1-stage path still gives clear benefits for shapes up to (32, 7168).

We also added TP8 measurements. For TP8, the 1-stage path is beneficial up to (16, 7168), while the 2-stage path becomes better from (32, 7168) onward. Based on this, we updated the gating logic to reflect the measured TP4 / TP8 behavior instead of using the same threshold for all TP sizes.

The new e2e profiling result is shown below. We did not include conc >= 128 because the 1-stage kernel is no longer applicable for those shapes.

Baseline: 1-stage vs 2-stage AllReduce path

Median latency is reported in ns, lower is better.
2-stage total = local_device_load_rmsnorm + reduce_scatter_cross_device_store (sum of medians).
1s/2s is the ratio of medians; < 1 means the 1-stage path is faster.

TP4

conc	1-stage (ns)	2-stage (rms+rs, ns)	1s/2s	faster
4	7,120	11,078 (4,479+6,599)	0.643	1-stage (1.56×)
8	8,079	10,998 (4,599+6,399)	0.735	1-stage (1.36×)
16	10,119	11,958 (4,599+7,359)	0.846	1-stage (1.18×)
32	13,919	14,758 (4,599+10,159)	0.943	1-stage (1.06×)
64	23,359	19,758 (4,839+14,919)	1.182	2-stage (1.18×)

TP8

conc	1-stage (ns)	2-stage (rms+rs, ns)	1s/2s	faster
4	8,999	11,678 (4,519+7,159)	0.771	1-stage (1.30×)
8	9,879	12,238 (4,479+7,759)	0.807	1-stage (1.24×)
16	12,439	12,918 (4,519+8,399)	0.963	1-stage (1.04×)
32	19,399	13,478 (4,519+8,959)	1.439	2-stage (1.44×)
64	35,759	16,878 (4,759+12,119)	2.119	2-stage (2.12×)

Overall:

• TP4: 1-stage is used for conc <= 32; 2-stage is used from conc >= 64.
• TP8: 1-stage is used for conc <= 16; 2-stage is used from conc >= 32.

Thanks again for your review. Please let me know if there is any other data you would like us to add.

[TennyWang1223]

Hi @ftyghome, I tested this pr and compire with main branch on MI355x, but test result alse seems a little different. In tp4, after m larger then 16 (include m=16), main branch get a shorter latency, and latency seems equal in tp8, need double check

[ftyghome]

Hi @TennyWang1223,

Thank you for testing this on the latest main branch!

In tp4, after m larger then 16 (include m=16), main branch get a shorter latency, and latency seems equal in tp8, need double check

I may be missing some context here, but I guess your point is that the 1-stage / 2-stage kernels may have been further optimized on the main branch, so the 1-stage gating threshold could have changed.

To double check this, I ran another round of profiling on the latest aiter / atom main branches on MI355X. I did notice that the kernel signature has changed, but the result is still similar to what we observed before: for TP4, the 1-stage path is still faster up to m = 32.

The numbers below are extracted from real e2e traces (using the SQL snippet shared in this PR).

TP4

m	n	1-stage (ns)	2-stage (rms+rs, ns)	1s/2s	faster
4	7168	7,839	10,718 (4,359+6,359)	0.731	1-stage
8	7168	7,599	11,319 (4,479+6,840)	0.671	1-stage
16	7168	9,719	11,398 (4,439+6,959)	0.853	1-stage
32	7168	13,719	14,638 (4,559+10,079)	0.937	1-stage
64	7168	22,119	19,758 (4,799+14,959)	1.119	2-stage
128	7168	—	28,238 (4,999+23,239)	—	2-stage

Would it be possible for you to share the benchmarking script or command you used? I can run the same benchmark on my side to make the comparison more direct.

Thanks again!

[TennyWang1223]

Hi, @ftyghome
Here is my test script, you can just in main branch and this pr and then compare results. Please let me know if you have any update
test_fused_ar_rms_perf.py

[ftyghome]

Hi @TennyWang1223 ,

Thanks for sharing the script!

I ran your script on my MI355X machine, but the result seems different from yours and is aligned with my previous e2e profiling result.

Here are the numbers I got.

All 1-stage

TP8

m	n	latency_us	err
4	7168	6.98	0
8	7168	16.04	0
16	7168	15.86	0
32	7168	14.62	0
64	7168	25.27	0

TP4

m	n	latency_us	err
4	7168	5.83	0
8	7168	5.88	0
16	7168	7.98	0
32	7168	11.55	0
64	7168	20.09	0

All 2-stage

TP8

m	n	latency_us	err
4	7168	15.60	0
8	7168	10.23	0
16	7168	17.09	0
32	7168	11.37	0
64	7168	13.46	0

TP4

m	n	latency_us	err
4	7168	9.19	0
8	7168	9.28	0
16	7168	10.03	0
32	7168	11.79	0
64	7168	16.38	0

I also noticed that some err values in your result seem to be non-zero, while all err values are 0 on my side. Could this be due to some environment difference?

For reference, my environment is:

aiter commit: 40b3427
atom commit: fdb7b02
ROCm: rocm/noble 7.2.3.70203-90~24.04 amd64
OS: Ubuntu 24.04

Both aiter and atom commits are from the latest main branches as of yesterday.

Please let me know if there is any other config I should align with your setup.

[TennyWang1223]

Hi, @ftyghome
I accept this result, but there is a conflict in the current PR. Please resolve this conflict and then I will re-review it.

[ftyghome]

Hi @TennyWang1223 ,

Thanks. I have the conflict resolved.