[HIP] fix(custom-ar): gfx1250 fused all-reduce + RMSNorm by XiaobingSuper · Pull Request #6699 · ROCm/aiter

XiaobingSuper · 2026-10-20T06:09:54Z

Motivation

On gfx1250 (MI455X), MiniMax-M3 TP4 + Eagle3 (K=3) produced wrong tokens. TP ranks returned different results from the same collective, so their speculative state diverged. No-spec decoding was unaffected.

gfx1250 had no fused all-reduce kernels, and the old-arch ones ran with a gfx1250 handle. fused_allreduce_rmsnorm* (including the per-group and MXFP4 quant forms) and fused_qknorm_allreduce* are built only into module_custom_all_reduce, and they cast _fa to the old-arch CustomAllreduce. On gfx1250, _ptr points to the gfx1250 class. These kernels then synchronize through the gfx1250 Signal with their own counters, which releases later gfx1250 barriers early. The next all-gather or all-reduce reads peers' buffers before they are staged, and returns stale rows that differ per rank. The Eagle3 draft model's norms use fused all-reduce + RMSNorm.

The other bug found in the same investigation, the LL all-reduce block-epoch race that shows up when small and large all-reduces alternate, is fixed on main by #6756. This PR is rebased on it and only fixes the fused path.

Changes

gfx1250 fused all-reduce + RMSNorm kernels: a direct-read one-stage path plus an exact two-stage path for TP4+ rows of at least 6144 elements. Two-stage reduce-scatters hidden shards into a peer-visible staging region, then gathers them while applying residual + RMSNorm. It keeps fixed rank-order accumulation and the same bf16/fp16 rounding as the unfused path. Both paths support gemma_norm, per-token FP8 output and the optional bf16 mirror.
Above the 4 MiB LL range, TP4+ wide rows use two-stage when the reduced shard fits its 16 MiB staging region; TP2 and narrower rows keep one-stage. Three single-block barriers publish the registered input and the shards without paying a barrier in every data block. Within the LL range, LL all-reduce followed by RMSNorm stays faster. Padded or strided rows take the unfused path.
The staging region sits after the LL scratch (meta_size grows by 16 MiB per rank), so the LL kernel, which accepts a line once its flag words equal the current epoch, never polls data it did not write.
Per-token FP8 output: on gfx1250 the fused per-token FP8 limit defaults to 64 MiB instead of the old-arch 4 MiB, which had left the gfx1250 FP8 kernel with no caller. The epilogue quantizes the bf16/fp16-rounded output with the steps of dynamic_per_token_scaled_quant, so fused and unfused FP8 bytes and scales are identical. The old-arch default and AITER_AR_QUANT_MAX_BYTES are unchanged.
LL all-reduce copy-in skip: gfx1250 always captures on the copy-in path, so every all-reduce first copied its input into the registered pool. The LL kernel (up to 4 MiB) reads its input in place, so the copy is skipped for LL sizes. This speeds up every small gfx1250 all-reduce, including the LL + RMSNorm dispatch that decode-sized fused calls take.
The per-group and MXFP4 quant forms and the qk-norm kernels still have no gfx1250 version. On gfx1250 the communicator takes its unfused path for per-group and MXFP4, and the direct qk-norm entry points raise instead of corrupting later collectives.
New op_tests/multigpu_tests/test_custom_ar_fused_interleave.py mixes fused, small, large and gathered collectives while delaying each rank by a different amount every step. It runs on gfx1250 only; CI has no gfx1250 runner.

Non-gfx1250 paths are unchanged: can_fuse_ar_rms, should_fuse_ar_rms and has_old_arch_fused_kernels are all true there, and the op table maps to the existing ops.

Test results

One MI455X node, rocm/fw-bringup:gfx1250-atom-k3-20261014 (ROCm 10.1.0, HIP 7.16.26315, PyTorch 2.11.0+rocm10.1.0a20260811, Triton 3.8.0). TP8 is not validated: no 8-GPU gfx1250 node was available.

Kernel check with op_tests/multigpu_tests/test_fused_ar_rms.py (bf16 and fp16, eager and CUDA graph, fused and per-token FP8; the file is in CI's multi-GPU skip list, so it was run by hand):

python3 op_tests/multigpu_tests/test_fused_ar_rms.py -t 4 --test fused quant \
  -s 16,6144 4,2880 384,6144 1024,6144 2730,6144 2800,6144 1024,7168 512,8192 1024,4096 1024,2880
python3 op_tests/multigpu_tests/test_fused_ar_rms.py -t 2 --test fused quant \
  -s 16,6144 384,6144 1024,6144 1024,4096 1024,2880 512,8192

All 128 runs pass, covering the LL range, one-stage (TP2, hidden 2880/4096), two-stage and FP8 output above 4 MiB. After the staging-region commit, TP4 was re-run at 384, 2800 and 5400 rows × 6144 and 1024×7168 (32 runs), and after the FP8 exact-match commit the per-token FP8 shapes were re-run (16 runs); all pass. Comparing fused FP8 with fused bf16 + dynamic_per_token_scaled_quant (TP4, 1024×6144, 2800×6144 and 1024×4096, with and without gemma_norm) gives identical bytes and scales.

Latency, TP4, bf16, hidden 6144, CUDA graph, µs per call, with the script below (python3 bench_ar_rms.py --rows 1,16,64,256,384,1024,2048,4096):

rows	MiB	LL or naive all-reduce	add + RMSNorm	unfused total	fused dispatch	vs unfused
1	0.01	8.9	2.8	11.7	11.5 (LL + RMSNorm)	1.02x
16	0.19	10.0	3.1	13.1	12.5 (LL + RMSNorm)	1.05x
64	0.75	12.8	3.2	16.0	15.7 (LL + RMSNorm)	1.02x
256	3	26.1	3.9	30.0	29.7 (LL + RMSNorm)	1.01x
384	4.5	64.3	4.2	68.5	44.1	1.55x
1024	12	120.2	6.3	126.5	64.0	1.98x
2048	24	207.0	9.8	216.8	95.1	2.28x
4096	48	382.7	17.5	400.1	148.0	2.70x

In the LL range the copy-in skip is the gain: with the same script, the LL all-reduce goes from 10.4 to 8.8 µs at 1 row and from 28.5 to 26.1 µs at 256 rows. A barrier-synchronized fused kernel is slower there (19.6 µs at 1 row), and a fused kernel on the LL protocol did not beat LL + RMSNorm below 64 rows, since RMSNorm needs a whole row in one block and that leaves few CUs to poll. Absolute numbers are from one node and vary between nodes.

bench_ar_rms.py

#!/usr/bin/env python3
"""TP graph-mode latency (us/call) of: LL all-reduce alone, add+RMSNorm alone, fused AR+RMSNorm dispatch.
Usage: bench_ar_rms.py [--tp 4] [--hidden 6144] [--rows 1,4,16,64,128,256,341] [--gemma]"""
import argparse

import torch
import torch.distributed as dist
from multiprocessing import Pool, set_start_method


def worker(rank, args, init_method):
    from aiter import rmsnorm2d_fwd_with_add
    from aiter.dist.communication_op import (
        tensor_model_parallel_all_reduce,
        tensor_model_parallel_fused_allreduce_rmsnorm,
    )
    from aiter.dist.parallel_state import (
        destroy_distributed_environment, destroy_model_parallel, ensure_model_parallel_initialized,
        get_tp_group, graph_capture, init_distributed_environment, set_custom_all_reduce)

    torch.cuda.set_device(rank)
    set_custom_all_reduce(True)
    init_distributed_environment(world_size=args.tp, rank=rank, distributed_init_method=init_method)
    ensure_model_parallel_initialized(args.tp, 1)
    H, eps, inner = args.hidden, 1e-6, 20
    w = torch.randn(H, device="cuda", dtype=torch.bfloat16)
    rows_list = [int(r) for r in args.rows.split(",")]
    results = {}
    for T in rows_list:
        x = torch.randn(T, H, device="cuda", dtype=torch.bfloat16)
        res = torch.randn(T, H, device="cuda", dtype=torch.bfloat16)
        out, res_out = torch.empty_like(x), torch.empty_like(x)
        ops = {
            "ar": lambda: tensor_model_parallel_all_reduce(x),
            "rms": lambda: rmsnorm2d_fwd_with_add(out, x, res, res_out, w, eps, args.gemma, 0),
            "fused": lambda: tensor_model_parallel_fused_allreduce_rmsnorm(
                x, res, w, eps, gemma_norm=args.gemma),
        }
        for name, fn in ops.items():
            for _ in range(3):
                fn()
            torch.cuda.synchronize()
            g = torch.cuda.CUDAGraph()
            with graph_capture() as gc, torch.cuda.graph(g, stream=gc.stream):
                for _ in range(inner):
                    fn()
            times = []
            for _ in range(30):
                get_tp_group().barrier()
                s, e = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)
                s.record()
                g.replay()
                e.record()
                e.synchronize()
                times.append(s.elapsed_time(e) * 1000 / inner)
            results[(T, name)] = sorted(times)[len(times) // 2]
            del g
    torch.cuda.synchronize()
    get_tp_group().barrier()
    destroy_model_parallel()
    destroy_distributed_environment()
    if dist.is_initialized():
        dist.destroy_process_group()
    return results


if __name__ == "__main__":
    set_start_method("spawn", force=True)
    p = argparse.ArgumentParser()
    p.add_argument("--tp", type=int, default=4)
    p.add_argument("--hidden", type=int, default=6144)
    p.add_argument("--rows", default="1,4,16,64,128,256,341")
    p.add_argument("--gemma", action="store_true")
    args = p.parse_args()
    from aiter.dist.utils import get_distributed_init_method, get_ip, get_open_port
    init = get_distributed_init_method(get_ip(), get_open_port())
    with Pool(args.tp) as pool:
        res = pool.starmap(worker, [(r, args, init) for r in range(args.tp)])
    print(f"hidden={args.hidden} tp={args.tp} gemma={args.gemma}  (us/call, max over ranks)")
    print(f"{'rows':>5} {'MiB':>6} {'LL AR':>7} {'add+RMS':>8} {'sum':>7} {'fused dispatch':>15}")
    for T in [int(r) for r in args.rows.split(",")]:
        v = {k: max(r[(T, k)] for r in res) for k in ("ar", "rms", "fused")}
        mib = T * args.hidden * 2 / 2**20
        print(f"{T:5d} {mib:6.2f} {v['ar']:7.1f} {v['rms']:8.1f} {v['ar'] + v['rms']:7.1f} {v['fused']:15.1f}")

test_custom_ar_fused_interleave.py, 300 steps, alternating 384/1024-row fused calls with LL all-reduce and all-gather:

Build	Result
without the fused-kernel fix	fused AR + RMSNorm corrupted (max err 16.5)
this PR (on main, includes #6756)	PASSED

Without #6756, the same test also fails on the plain 16-row all-reduce, independent of this PR.

End to end, MiniMax-M3 MXFP4 TP4 + Eagle3 GQA K3, full GSM8K (1319, 5-shot, c4), default custom all-reduce settings, with ROCm/ATOM#2710 for FlyDSL paged decode (ATOM_PA_FLYDSL=1):

	KV	flexible	strict	acceptance
before	BF16	0.9151	0.9121	75.6% (inflated by corrupted verify rows)
this PR	BF16	0.9386	0.9393	72.7%
this PR	FP8	0.9318	0.9318	72.8%
no-spec reference	BF16	0.9386	0.9393	–

Measured at eeb90b37a. The later commits move the two-stage staging buffer and route FP8 output above 4 MiB to the fused kernel, which now produces the same bytes as the unfused path it replaces. An earlier version of that kernel quantized the fp32 value instead and moved BF16 GSM8K to 0.9303/0.9310; the exact-match version is what this PR ships. The "before" row used the earlier unified-attention decode stack, which ROCm/ATOM#2710 replaces.

github-actions · 2026-10-20T06:10:23Z

🏷️ 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:gfx1250-ffm-triton`	Run the five-shard gfx1250 FFM Triton test suite
`ci:triton-300x`	Run an additional Triton test job on MI300X in PRs (added automatically when gfx942 configs change); main branch always runs both MI35X and MI300X
`ci:triton-355`	Run the full Triton test suite on MI35X, not only the tests the change affects
`ci:windows`	Windows JIT op smoke test on the gfx1201 ROCm runner
`ci:multi-gpu`	Aiter multi-GPU tests on the 8-GPU runner
`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: DeepSeek-V4-Pro, MiniMax-M3-MXFP4
`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 6699 --add-label <label>

One backend per PR:
A PR changes one kernel backend: [Triton/Gluon] (Triton and Gluon count as one), [HIP], [ASM], [CK], [OPUS] or [FlyDSL]. If the title ends up with two backend tags, split the PR -- as stacked pull requests when one part cannot merge without the other.

PR title tags & labels:
Component tags ([Triton/Gluon], [HIP], [CK], [ASM], ...) are added to the PR title and as PR labels automatically from the changed files and re-synced on every push — change-type tags like [fix]/[Perf], op tags like [MLA], and human labels (ci:*) are left untouched. Add the no-auto-title label to stop the title rewrites; labels stay in sync either way.

fused_ar_rms and its per-token FP8 form only had old-arch kernels, which cast _ptr to the old-arch CustomAllreduce. On gfx1250 _ptr is the gfx1250 class, so those kernels synchronized through the gfx1250 Signal with their own counters and released later gfx1250 barriers early. Add a gfx1250 kernel (one block per token, ranks summed in rank order) with the same numerics as the old-arch kernel, including gemma_norm and the optional bf16 mirror, and dispatch fused_ar_rms to it on gfx1250. The kernel syncs through the barrier, so the communicator uses it above the 4 MiB LL all-reduce range, where it is 1.6-2.4x faster than all-reduce plus RMSNorm on MI455X. Within the LL range, the LL all-reduce plus RMSNorm stays faster. Padded or strided rows have no gfx1250 kernel and take the unfused path. Old-arch dispatch is unchanged.

The per-group and MXFP4 quant forms of fused all-reduce + RMSNorm and the qk-norm all-reduce kernels exist only in the old-arch module and would hit the same Signal mix-up on gfx1250. Skip them there: the communicator takes its unfused path for per-group and MXFP4, and the direct qk-norm entry points raise instead of corrupting later collectives.

… rank skew Speculative decoding alternates small fused all-reduce + RMSNorm calls (the draft's norms) with larger ones (the target), plus plain all-reduces and the barrier-synchronized logits all-gather, all on one custom all-reduce instance while the ranks drift apart. This TP4 test runs that mix with a different GPU-side delay per rank every step and checks every result against an fp32 reference. It fails when gfx1250 runs the old-arch fused kernels (they corrupt the shared Signal and later collectives read stale rows) and on the LL block-epoch race fixed in #6756.

Co-authored-by: Cursor <cursoragent@cursor.com>

gfx1250 always captures on the copy-in path, so every all-reduce first copies its input into the registered pool. The LL kernel reads the input in place and never touches that pool, so for LL sizes the copy is pure overhead: on MI455X TP4 it costs 1.5-2.5 us per call in CUDA graphs (hidden 6144: 1 row 10.4 -> 8.8 us, 256 rows 28.5 -> 26.1 us).

It guards gfx1250 kernels and dispatch; other archs keep their own fused path, and the CI multi-GPU runners are MI350X.

Copilot