[triton-mha] add gfx1151 tuning config - #3423
Conversation
`_attn_fwd` (the `default` MHA impl) reads its tuning from
`aiter/ops/triton/configs/<gfx>-MHA-DEFAULT.json` at call time. Without
a matching file the wrapper raises `FileNotFoundError`, which made
`mha_set_impl("default")` unusable on Strix Halo.
Add `gfx1151-MHA-DEFAULT.json`. Block sizes match the existing gfx950
tuning (BLOCK_M=128, BLOCK_N=64, waves_per_eu=2, num_stages=1); the
`default` forward profile picks `num_warps=8` to mirror the choice
used by the `flash_attn_triton_amd` (`dao_ai`) prefill kernel on the
same arch, which empirically halves loop-dispatch overhead on the
Qwen3-Omni ViT prefill shape (head_dim=72) versus `num_warps=4`.
Backward, PE, and dropout entries follow the gfx950 defaults pending
arch-specific validation.
On gfx1151 with the `dao_ai`-style head-stride alignment hint applied
to `_attn_fwd`, the Qwen3-Omni ViT prefill shape
(B=1, S=3200, H=16, head_dim=72, fp16) goes from a `[SKIP]` (no config)
to 2.39 ms via `bench_mha -impl default -fn fwd_varlen`, matching the
`dao_ai` impl within measurement noise (2.41 ms).
🏷️ CI GuideRuns automatically on every PR:
Extended tests (opt-in via labels):
|
There was a problem hiding this comment.
Pull request overview
Adds a new Triton MHA default tuning JSON for gfx1151 so the default MHA forward implementation can load an architecture-specific config instead of failing with FileNotFoundError on Strix Halo / RDNA3.5.
Changes:
- Add
aiter/ops/triton/configs/gfx1151-MHA-DEFAULT.jsoncontaining forward + backward tuning parameters for gfx1151. - Set
fwd.default.num_warps=8for gfx1151 (relative to other arch configs).
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
|
@vgokhale, the tuning in here focusses on the ViT prefill shape: B=1, S=3200, H=16, head_dim=72, fp16, varlen, non-causal. Are there scripts I can run to benchmark/automatically tune the remaining entries of this json file in a follow-up PR? |
Hi @mgehre-amd. You can try I'm not aware about any tuning script for Triton MHA. I usually hack some dirty shell script to edit the config file in-place and then run |
gfx1101 (RDNA3, e.g. RX 7800 XT) ships no MHA config, so `_get_config` in `_triton_kernels/attention/mha.py` finds no `configs/gfx1101/triton/attention/mha/DEFAULT.json` and every call to `aiter.ops.triton.attention.mha.flash_attn_func` fails before a kernel runs. Of the architectures in RDNA_ARCHS, only gfx1151 ships one today. Nine of the eleven entries are taken verbatim from the gfx1151 donor (RDNA3.5, added in ROCm#3423, tuned in ROCm#3560), which is the nearest tuned architecture. Two forward entries are tuned on gfx1101 instead of inherited, and they differ from each other only in `num_warps` and `num_stages`: fwd/default BLOCK_M 128, num_warps 8, num_stages 3 (donor: 64 / 4 / 2) fwd/small_head BLOCK_M 128, num_warps 4, num_stages 1 The split uses the `small_head` bucket added in ROCm#4414, which is opt-in per architecture by the mere presence of the key, so this stays a data-only change. It is needed because one `fwd/default` cannot serve both halves on this card: measured against the donor, `M128 w4 s1` is 0.935x on head_dim 64 but 1.269x on head_dim 128, while `M128 w8 s3` is 0.914x on head_dim 128 but 1.062x on head_dim 64. Measured on Windows native ROCm, triton 3.8.0, fp16, 10 independent repeats of 20 iterations after 5 warmups, `torch.cuda.synchronize()` per iteration; a result counts only when the [min, max] intervals across repeats are disjoint. Ratios are against the gfx1151 donor entry, i.e. against what a donor-inherited config would do. head_dim 128, `default` Flux joint 0.914x / 0.891x / 0.892x on three independent torch+ROCm stacks (2.11/7.15, 2.11/10.1, 2.15/10.1), all pinned to the same triton; llama3-8B 0.882x, mixtral-7B 0.880x, kimik25-tp4 0.849x at seqlen 16384 head_dim <= 64, `small_head` SDXL self-attn 0.935x; deepseek-V3 0.658x and glm47fp8-tp4 0.857x at seqlen 16384 The LLM shapes come from `op_tests/op_benchmarks/triton/utils/model_configs.json` (prefill, causal, GQA, batch 1); the sequence lengths are not in that file and are chosen here. Not covered: batch > 1, varlen/thd, sliding window, decode. Note that the `small_head` comment in `_get_config` does not describe gfx1101. It states that 16 < d <= 64 suffers a num_stages=1 pipelining pathology which num_stages=3 cures. On this card the ordering is the opposite -- on the tuned M128/N32/w4 tile, num_stages 1/2/3 measured 2.300 / 2.399 / 2.492 ms. The bucket is still the right mechanism here, for a different reason: the two head_dim ranges want a different `num_warps`, not a different `num_stages`. Signed-off-by: Martin Domanský <ragua@email.cz>
Summary
Adds
aiter/ops/triton/configs/gfx1151-MHA-DEFAULT.jsonso that AITER'sdefaultMHA forward impl (aiter.ops.triton.attention.mha._attn_fwd) is usable on gfx1151 (Strix Halo / RDNA3.5).Without a matching
<gfx>-MHA-DEFAULT.jsonthe wrapper raisesFileNotFoundErrorat call time, which mademha_set_impl("default")unusable on Strix Halo. The other shipped tunings cover gfx942 (CDNA3), gfx950 (CDNA4), and gfx1250 (RDNA4) — gfx1151 was the only consumer-RDNA gap.Block sizes mirror the gfx950 tuning (
BLOCK_M=128,BLOCK_N=64,waves_per_eu=2,num_stages=1). Thedefaultforward profile picksnum_warps=8(vs. gfx950's4) to match theflash_attn_triton_amd(dao_ai) prefill kernel's choice for the same arch — empirically it roughly halves loop-dispatch overhead on the Qwen3-Omni ViT prefill shape (B=1, S=3200, H=16, head_dim=72). Backward, PE, and dropout entries follow the gfx950 defaults pending arch-specific validation.Benchmark — gfx1151,
bench_mha -impl defaultQwen3-Omni ViT prefill shape: B=1, S=3200, H=16, head_dim=72, fp16, varlen, non-causal. Toolchain: torch 2.11.0+rocm7.14.0a20260529, triton 3.7.0 (built from
triton-lang/tritonmain,d92727c2).bench_mhatimeorigin/main(no gfx1151 config)FileNotFoundError: gfx1151-MHA-DEFAULT.jsonThis PR by itself moves the
defaultimpl from "unusable" to "runs but ~2× slower than necessary onhead_dim=72". To close the remaining gap and reach parity with thedao_aiimpl (~2.41 ms on the same shape), pair it with the head-stridetl.multiple_of(..., 8)hint shipped as a separate, independent PR.Reproducer (in-tree)
python op_tests/op_benchmarks/triton/bench_mha.py \ -fn fwd_varlen -equal_seqlens \ -b 1 -hq 16 -hk 16 -sq 3200 -sk 3200 -d 72 \ --dtype fp16 -causal False \ -impl default -metric timeCorrectness
This PR adds a JSON configuration file only; no kernel code changes.
op_tests/triton_tests/attention/test_mha.py::test_mhanow collects and runs (previously: hardFileNotFoundErrorat collection / first call). 171 PASSED, 9 pre-existingtorch.OutOfMemoryErrorfailures (attention_ref(upcast=True)materialising[B, H, S, S]fp32 scores at B∈{30,50}, H=48, S=2048 — exceeds 64 GiB Strix Halo VRAM; identical behavior onorigin/mainwith any impl).Test plan
bench_mha -impl defaulton gfx1151op_tests/triton_tests/attention/test_mha.py::test_mha(171 PASSED, 9 pre-existing OOMs)