Tune MHA config & small-head pipeline pathology

[nidal567]

Motivation

On gfx950, the TOT MHA forward path (Triton 3.8 + async_copy) regressed against the golden baseline (Triton 3.4) on a set of shapes: all 18 MLA-prefill PE shapes (DeepSeek-R1 / Kimi-K2, dqk=192/dv=128) up to +8.7%, and a head_dim=64 causal shape by +204%. The latter pathological, with individual _attn_fwd dispatches intermittently stalling up to 1.6 seconds.

This PR retunes gfx950-MHA-DEFAULT.json (plus a minimal head-dim-keyed config selection) to recover them without disturbing the shapes that already improved.

TLDR

Tuning gfx950-MHA-DEFAULT.json (+ a small head-dim-keyed config-selection correction) on the TOT env:

• Regressions: 28 → 8. 20 of 28 initial regressions fixed to ok/improvement; 0 new regressions introduced.
• Improvements: 193 → 222 (out of 240 shapes); ok: 19 → 10.
• All 18 PE-bucket regressions fixed (DeepSeek-R1 / Kimi-K2 MLA-prefill, dqk=192/dv=128) — now 4–13 % faster than golden.
• The catastrophic +204 % outlier fixed (d64 dense causal): 1030 µs → 333 µs, now matching golden; per-dispatch profiling shows TOT num_stages=1 intermittently stalled individual dispatches up to 1.6 s, eliminated by the fix.
• Correctness: 669/669 forward pass; full test_mha.py 905 passed / 132 skipped / 0 failed (backward included, unaffected). The two config changes are numerically transparent (pipelining only) on the shapes they touch; a num_stages=3 Triton-3.8 miscompile for head_dim≤16 was found and explicitly excluded.
• 8 residual regressions remain (6 tiny synthetic test shapes + 2 marginal real d128 at +1.3–1.5 %). Attributed with evidence to Triton 3.8 codegen on small MHA-forward workloads — not async-copy (disabling async-copy makes them worse), and not config-fixable without introducing new regressions or the miscompile.

Technical Details

Two independent, config-first changes:

1. pe.num_stages 1 -> 3 - restores golden's value that the tot config had dropped. This recovers all 18 existing PE regressions; full 32-shape PE bucket is now 4-13% faster than golden.
2. New small_head config bucket + head-dim gate in _get_config. head_dim isn't a config selection input, so d64/d32 shared the default bucket with d128; setting num_stages=3 fixes the small-head pipelining pathology but regresses small-M d128. Fix: _get_config takes head_dim_v (wrapper passes v_head_dim) and routes 16 < d_head_v <= 64 to a new small_head bucket (num_stagesk=3), leaving d128 on num_stages=1. Branch order keeps PR/dropout/fp32 unaffected. Lower bound >16 is a correctness guard: num_stages=3 miscompiles head_dim <= 16 on Triton 3.8, so those stay on default.

Test Plan

Correctness test for validation:
pytest op_tests/triton_tests/attention/test_mha.py -v
Expected: 905 passed, 132 skipped, 0 failed

Running smoke tests with Triton's do_bench

python op_tests/op_benchmarks/triton/bench_mha.py -fn fwd -b 1 -hq 128 -hk 128 -sq 8192 -sk 8192 -d 192 -dv 128 -causal True --dtype bf16 -metric time
python op_tests/op_benchmarks/triton/bench_mha.py -fn fwd -b 8 -hq 48  -hk 8   -sq 2048 -sk 2048 -d 64          -causal True --dtype bf16 -metric time

Numbers are collected usingrocprofv3 --kernel-trace, undergoing a trimmed mean methodology to grab reliable single average numbers for each test case for all existing shapes shared in the default gfx950-MHA-DEFAULT.json config, the model_shapes.json, and the default in Triton's do_bench (240 shapes in total).

Source of shapes:
240 deduplicated shapes. M = N_CTX_Q (sq), N = N_CTX_K (sk), K = d_head_qk. Source tags (unioned on dedup): bench (bench_mha.py make_workloads × the 10 utils/model_configs.json models) on 160, model_shapes (model_shapes.json mha entries incl. DeepSeek-R1/Kimi-K2 PE and GPT-OSS sink+sliding-window) on 136, test (distinct test_mha.py perf shapes: head_dim=64, GQA 48:8, 2048², varlen head_dim 8/32) on 16. TP ∈ {1, 8} where per-rank HQ/HK stay GQA-valid (no KV-head replication; Qwen3 hk=4 → TP1 only)

Test Result

• Correctness: full test_mha.py
905 passed, 132 skipped, 6 warnings in 127.86s (0:02:07)

• Performance (with rocprofv3 --kernel-trace, chronological middle-80% trimmed mean, 3–5 reps, golden vs tot over 240 dedup shapes):

  • regressions 28 → 8, improvements 193 → 222, 0 new regressions.
  • All 18 PE regressions → improvement
  • The +204% outlier → −1.5% (1030 µs → 333 µs, per-dispatch 1.6 s stalls eliminated);
  • PE example DeepSeek-R1 8192²: baseline tot 4901 µs → tuned 4083 µs (−16.7%, beats golden's 4546 µs).

• The remaining 8 regressions:

  • The 8 residuals (6 tiny synthetic test shapes + 2 marginal real-model d128 at +1.3–1.5%) are attributed to Triton-3.8 codegen (async-copy ruled out — disabling it makes them worse), config-unfixable without new regressions.
  • Over the next few weeks it is planned to profile and collect traces of these codegen regressions with ATT, and collect the IRs and ASM for the respective kernels.

• Validation checks with black --check and ruff check clean

Profile 1 — the +204 % outlier dense_b8_hq48_hk8_sq2048_sk2048_dqk64_dv64_c1 (200 dispatches, rep0):

variant	median	mean	max single dispatch	trimmed mean
GOLDEN (3.4, stages=1)	332.8 µs	343.7 µs	457.3 µs	338.2 µs
TOT baseline (3.8, stages=1)	390.4 µs	15223 µs	1,640,403 µs (1.64 s)	pathological
TOT tuned (3.8, stages=3)	326.7 µs	338.4 µs	443.3 µs	333.0 µs

At num_stages=1 on Triton 3.8, individual _attn_fwd dispatches intermittently stall up to 1.64 seconds (median stays ~390 µs). num_stages=3 eliminates the stalls → clean 333 µs, matching golden.

Profile 2 — PE improver dense_b1_hq128_hk128_sq8192_sk8192_dqk192_dv128_c1 (DeepSeek-R1 prefill):

variant	trimmed mean
GOLDEN (3.4, pe stages=3)	4545.7 µs
TOT baseline (3.8, pe stages=1)	4901.2 µs (+7.8 % vs golden)
TOT tuned (3.8, pe stages=3)	4083.2 µs (−16.7 % vs baseline tot, −10.2 % vs golden)

Mechanism (Triton LDS metadata): num_stages=3 kernels allocate more shared/LDS memory for pipelining buffers (e.g. 16384 vs 8192 bytes) — deeper software pipelining hides memory latency, which is why it helps the PE and mid-small-head shapes but pressures occupancy on small-M d128 (hence the head-dim gate).

Final status:

status	baseline	final
improvement	193	222
ok	19	10
regression	28	8

For the following tables, for simplicity, keep in mind the following: M = N_CTX_Q (sq), N = N_CTX_K (sk), and K = d_head_qk

The 20 fixed (representative):

M	N	K	d_v	HQ/HK	TP	layout	bucket	golden	tot_base (Δ%)	tot_final (Δ%)	status
2048	2048	64	64	48/8	1	dense	small_head	338.2	1030.4 (+204.7%)	333.0 (−1.5%)	improvement
8192	8192	192	128	128/128	1	dense	pe	4547.6	4944.0 (+8.7%)	4119.2 (−9.4%)	improvement
1024	1024	192	128	64/64	1	varlen	pe	472.0	512.5 (+8.6%)	447.8 (−5.1%)	improvement
8192	8192	192	128	64/64	8	varlen	pe	333.0	357.5 (+7.4%)	289.4 (−13.1%)	improvement
8192	8192	192	128	128/128	1	dense(c0)	pe	7536.6	7838.0 (+4.0%)	7086.1 (−6.0%)	improvement
8192	8192	128	128	128/8	1	dense	default	2678.4	2799.2 (+4.5%)	2686.6 (+0.3%)	ok
… (all 18 PE shapes -> improvement, + the two (default and small_head) above)

The 8 remaining regressions:

M	N	K	d_v	HQ/HK	TP	layout	golden_us	tot_final_us	Δ%	source
2048	2048	64	64	6/1	8	varlen	66.6	76.5	+14.9%	test:head64
2048	2048	64	64	6/1	8	dense	65.5	72.9	+11.3%	test:head64
2048	2048	64	64	8/8	1	varlen	80.5	86.3	+7.3%	test:head64
2048	2048	32	32	16/16	1	varlen	61.3	65.5	+6.9%	test:vd32
2048	2048	64	64	8/8	1	dense	79.5	84.5	+6.4%	test:head64
8192	8192	128	128	5/1	8	dense(c0)	303.2	307.9	+1.5%	model_shapes:Llama4-Maverick
8192	8192	128	128	8/1	8	dense(c0)	330.0	334.2	+1.3%	model_shapes:Llama3-70B
2048	2048	64	64	48/8	1	varlen	338.3	342.4	+1.2%	test:head64

6 of 8 are synthetic test-coverage shapes (small few-head 2048² workloads, 60–86 µs); 2 are marginal real-model d128 (+1.3–1.5 %, ≈4 µs). num_stages=3 (small_head) improves the synthetic ones vs baseline (e.g. hq6/hk1 varlen +19.8 % → +14.9 %) but does not cross the ok threshold.

Submission Checklist

• Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

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

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

[brunomazzottiamd]

LGTM!

[brunomazzottiamd]

@nidal567, can you please rebase on top of main so CI can run on gfx942? I think Aiter Test / Standard Tests (1 GPU) (linux-aiter-mi35x-1, MI35X, 8, 7) should be fixed as well.