[triton-mha] add gfx1101 tuning config

[Ragua1]

Motivation

aiter.ops.triton.attention.mha.flash_attn_func cannot run at all on gfx1101 (RDNA3). _get_config in _triton_kernels/attention/mha.py opens configs/{arch}-MHA-DEFAULT.json unconditionally, and with no per-arch file and no fallback the call raises FileNotFoundError before any kernel is compiled.

RDNA_ARCHS in _triton_kernels/flash_attn_triton_amd/utils.py lists eight architectures and only gfx1151 ships an MHA config, so the default MHA path is unreachable on the other seven. This PR adds the one architecture available to me for testing on hardware, following #3423 and #3560, which between them added and then tuned gfx1151-MHA-DEFAULT.json.

Technical Details

One new file, aiter/ops/triton/configs/gfx1101-MHA-DEFAULT.json: 92 lines added, nothing removed, no code touched.

Nine of the ten entries are taken verbatim from gfx1151-MHA-DEFAULT.json, the nearest tuned architecture — RDNA3.5 against RDNA3, same Wave32 execution and the same WMMA generation. fwd/default is tuned on this card instead of inherited, so the difference against the gfx1151 file is two numbers: BLOCK_M 64 → 128, num_stages 2 → 1.

That entry and not the others because fwd/default is what _get_config selects whenever dropout is disabled and the dtype is not fp32 — the branch tests dtype == torch.float32, so fp16, bf16 and fp8 all resolve to it. fwd/dropout_or_fp32, both pe entries and both backward sections are inherited unchanged; only the forward path was exercised, as this host is used for inference, so no claim is made about backward here.

Outside the scope of this change: a second call in the same process reports KeyError: 'default' rather than the original error, because the memo on _get_config is populated before the file is opened. That is addressed in #4447.

Test Plan

The failure is total — the op raises before any kernel runs — so verification has two parts: the call must complete, and the result must be numerically correct. Compare the output against an fp32 SDPA-MATH reference, with torch's own fp16 SDPA measured against that same reference in the same process so the tolerance comes from a control rather than a quoted figure; confirm the failure reappears when the file is removed; and exercise all four forward entries, not only the one a fp16 non-dropout shape selects, since the inherited pe entries carry the largest tile in the file.

fwd/default was tuned in the two stages #3560 used, with 5 warmup and 20 timed iterations per point, torch.cuda.synchronize() per iteration, fp16, seed 424242, median reported. Stage 1: 48 combinations of BLOCK_M ∈ {32, 64, 128, 256} × BLOCK_N ∈ {32, 64} × num_warps ∈ {2, 4, 8} × num_stages ∈ {1, 2} on SDXL self-attention (2, 4096, 4096, 10, 64) — all 48 valid, none exceeding the 64 KiB LDS budget. Stage 2: the best three plus the inherited values across four shapes — SDXL self-attention at 4096 and 1024, SDXL cross-attention, and a Flux joint-attention shape at head_dim 128 — additionally varying PRELOAD_V and waves_per_eu. The leading candidates were then repeated seven times each, so a ratio is only claimed where the ranges are disjoint, and re-verified in bf16.

⚠️ The scope of what this measures should be stated before the numbers. Four diffusion attention shapes is a much narrower basis than the n=112 LLM shape set behind #3560, and I have no way to check these values against the shapes that PR targeted. If that basis is too narrow to accept, the alternative is to inherit fwd/default from gfx1151 unchanged along with the other nine entries, which is a two-number revert.

op_tests/triton_tests/attention/test_mha.py could not be run on this host: AITER_TRITON_ONLY is unconditional on win32, and aiter.test_mha_common imports aiter.dtypes, which the Triton-only branch of aiter/__init__.py does not bind. This is independent of the present change, and I preferred not to introduce a workaround into the verification being reported.

Reproducer, run with the AITER checkout on PYTHONPATH (no install):

import torch
from torch.nn.attention import SDPBackend, sdpa_kernel

from aiter.ops.triton.attention import mha
from aiter.ops.triton.utils._triton import arch_info

b, sq, sk, h, d = 2, 4096, 4096, 10, 64
gen = torch.Generator(device="cuda").manual_seed(424242)
opts = {"dtype": torch.float16, "device": "cuda"}
q = torch.randn(b, sq, h, d, generator=gen, **opts)
k = torch.randn(b, sk, h, d, generator=gen, **opts)
v = torch.randn(b, sk, h, d, generator=gen, **opts)

print("arch:", arch_info.get_arch())

qt, kt, vt = (x.transpose(1, 2) for x in (q, k, v))
with sdpa_kernel(SDPBackend.MATH):
    ref = torch.nn.functional.scaled_dot_product_attention(
        qt.float(), kt.float(), vt.float()
    ).transpose(1, 2)
    sdpa16 = torch.nn.functional.scaled_dot_product_attention(qt, kt, vt).transpose(1, 2)

out = mha.flash_attn_func(q, k, v, causal=False)  # raises without this PR

print("out:", tuple(out.shape), out.dtype)
print("max abs diff vs fp32 reference : aiter %.3e | torch SDPA fp16 %.3e" % (
    (out.float() - ref).abs().max(),
    (sdpa16.float() - ref).abs().max(),
))

Tested on: RX 7800 XT (gfx1101, RDNA3), Windows, torch 2.11.0+rocm7.15.0a20260728 (HIP 7.15.26290), triton 3.7.1, AITER_TRITON_ONLY implicit. Not tested on CDNA or on any other RDNA part; no other hardware is available to me.

Test Result

With the file in place:

[aiter] Triton ops only: CK and HIP ops (and their JIT build) are skipped.
arch: gfx1101
out: (2, 4096, 10, 64) torch.float16
max abs diff vs fp32 reference : aiter 6.470e-05 | torch SDPA fp16 6.013e-05

A control that accounts for the residual: with TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1, torch's own AOTriton flash kernel becomes selectable on this card, and against the same fp32 reference it returns exactly the figures this config produces — max abs 6.470e-05, normalized 2.294e-04. The gap to the fp16 SDPA-MATH figure is therefore attributable to flash-style tiling and online softmax rather than to these tile sizes.

Performance, median of 20 iterations, seven independent repeats (range across repeats):

Shape	Inherited	This PR	Ratio of medians	Intervals
SDXL self 4096	2.820 – 3.246 ms	2.195 – 2.246 ms	0.730×	disjoint
SDXL self 1024	0.614 – 0.684 ms	0.537 – 0.552 ms	0.843×	disjoint
SDXL cross	0.314 – 0.417 ms	0.360 – 0.371 ms	0.895×	overlapping
Flux joint, head_dim 128	6.822 – 7.169 ms	7.134 – 7.620 ms	1.064×	overlapping

The claim is therefore confined to diffusion self-attention. On cross-attention and at head_dim 128 no reliable difference was established — the ranges overlap, the inherited entry has a coefficient of variation of 11.5 % on cross-attention, and across three independent runs the Flux ratio measured 0.988×, 1.064× and 1.15×. bf16 gives the same ordering, with the chosen entry at 0.792× of the inherited one on the dominant shape.

Two supporting results behind the two changed numbers:

• num_stages: 1. BLOCK_M 128 with two stages was the single fastest cell on SDXL but 1.35× slower than inherited at head_dim 128; one stage costs 0.028 ms on SDXL and carries no such regression. On this architecture an extra stage can only buffer in registers — is_pingpong_schedule_enabled in the AMD backend admits gfx942 and gfx950 only, and enabling async copy makes num_stages ≥ 2 fail to compile outright (ttg.async_copy_global_to_local ... explicitly marked illegal) — and n_spills on the compiled kernel rises monotonically with the value.
• BLOCK_M: 128. Sweeping BLOCK_M × BLOCK_N with n_regs/n_spills read off each kernel, BLOCK_M 128 / BLOCK_N 32 is the only combination reaching n_regs 256 with zero spills at head_dim 64. BLOCK_M 64 leaves 68 registers unused while reading K and V twice as many times; BLOCK_M 256 halves that volume again but spills 168 registers and is slower, 3.003 ms against 2.219 ms.

Every other fwd/default in aiter/ops/triton/configs pairs BLOCK_M: 128 with num_stages: 1 — gfx942, gfx950 and gfx1250 — and gfx1151 is the only entry at 64 and 2. The sweep here converged on the same pair independently. That is offered as corroboration for the gfx1101 values, not as a claim about gfx1151, which is an APU with a different cache hierarchy.

📌 All thirteen configurations in stage 2 produced an identical maximum absolute difference per shape — one distinct value per shape, including for the inherited entry. Tile parameters therefore have no numerical effect on this kernel; the choice is purely a performance decision. All four forward entries were also exercised in separate processes, since _get_config memoises, and each ran with the expected entry selected.

Why fwd/dropout_or_fp32 was left inherited, as a measurement rather than as caution. That entry was swept as well and a configuration exists that is faster in fp32 — BLOCK_M 128, BLOCK_N 32, 4 warps, 2 stages at 0.81×. It was not applied because _get_config routes two different workloads to that one entry, fp32 without dropout and fp16 with dropout, and that configuration is 25.3× slower on the dropout path and 4.24× slower in fp32 at head_dim 128, with the runner-up failing to launch (OutOfResources) — reproduced in two independent runs. The inherited BLOCK_M 32 with 2 warps is the only one acceptable across all four combinations. That hazard is also why fwd/default was verified across four shapes and two dtypes rather than on the shape it was selected on.

Related: a fallback for the same defect, prepared but not included here

_get_config has no fallback at all, so 11 of the 14 architectures listed in CDNA_ARCHS and RDNA_ARCHS raise FileNotFoundError on this path — only gfx942, gfx950 and gfx1151 ship a config. This PR is the narrow remedy for one of the eleven. A same-family fallback is implemented and tested: +56/−1 in _get_config plus an 89-line test, CDNA falling back to gfx942 and RDNA to gfx1151 with a single warning naming the donor architecture, covering 10 of the 11. gfx1030 is excluded deliberately — it is listed in RDNA_ARCHS but has no WMMA units. Because path resolution needs no GPU the test covers all fourteen architectures on any host (28 passed, 1 skipped), and on this card the fallback runs end to end with no gfx1101 config present, producing the same 6.470e-05.

The two changes are complementary rather than alternatives: a fallback would hand gfx1101 the gfx1151 values, which are the ones in the inherited column above, so it makes the op reachable on ten architectures while this file makes it fast on the one that could be measured. I would prefer a maintainer to choose between (a) the fallback added here as a second commit, (b) the fallback as a separate PR with this one remaining data-only, or (c) no fallback, retaining one config file per architecture as in #3423 and #3560 — that is a policy question about how the repository handles missing configs rather than a technical one. In any order, the fallback modifies the line immediately above the one #4447 (also mine) moves, and #4414 edits the same function further down; both resolve as one-line rebases.

Diff shape:

$ git diff --numstat main
92      0       aiter/ops/triton/configs/gfx1101-MHA-DEFAULT.json

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

[carlushuang]

LGTM