[phambinhfin]

Summary

Add an fp8_large_head entry to the gfx942 Triton MHA forward config and select
it for FP8 at head_dim_v == 256, where fwd["default"] is 3.2-3.8x slower.

Claims

• 3.80x at 30720q x 32768k non-causal, 3.64x at 16384q x 16384k causal, 3.21x at
  4096q x 4096k causal on MI325X, vs today's default at the same shapes.
• Within 0.02-2.4% of the best of 108 swept configs at the two primary shapes on
  two MI325X hosts.
• Scoped to head_dim_v == 256 and FP8 only: d255, d320, d512, bf16/fp16 at
  d256, and fp32/dropout all keep their current config. Above d256 the head dim
  pads to 512 and this entry exceeds LDS, so it cannot widen.
• gfx942 only. No other arch has the entry, so _get_config is unchanged there
  by construction.

Validation

Setup: 1x MI325X (gfx942), vllm/vllm-openai-rocm:nightly-5d35ce3934a7,
torch 2.13.0+rocm10.0.0, aiter at this branch via PYTHONPATH. Before is
fwd["default"], after is fwd["fp8_large_head"], same process, same GPU.

PYTHONPATH=$PWD python3 -m pytest op_tests/triton_tests/attention/test_mha.py \
    -k "fp8_large_head" -v

Kernel

Head dim 256, 24 query heads, 4 KV heads, FP8 Q/K/V, varlen. Median of 10 after
3 warmups.

Shape	default -> fp8_large_head
30720q x 32768k, non-causal	145.401 -> 38.266 ms (3.80x)
16384q x 16384k, causal	22.561 -> 6.203 ms (3.64x)
4096q x 4096k, causal	1.743 -> 0.543 ms (3.21x)

vs CK bf16 at the first shape: 58.699 -> 38.837 ms, relative error 0.0530,
matching the pre-existing FP8 path. Launch parameters change, not the math.

Grid sweep

108 configs over BLOCK_M, BLOCK_N, num_warps, num_stages, PRELOAD_V, two
MI325X hosts, each checked against CK bf16. 72 compiled; the rest hit
OutOfResources at BLOCK_M=256, BLOCK_N=128.

Host	Shape	This entry	Best of 108	Gap
A	30720q x 32768k, non-causal	38.837 ms	38.830 ms	+0.02%
A	30720q x 30720k, causal	21.055 ms	20.558 ms	+2.4%
B	30720q x 32768k, non-causal	39.127 ms	39.103 ms	+0.06%
B	30720q x 30720k, causal	20.998 ms	20.576 ms	+2.0%

Small shapes

Same heads and dtype, varlen with equal lengths per sequence, default vs
fp8_large_head with an A/A default control, 5 alternating rounds, medians.
22 shapes, batch 1-64, seqlen 128-4096, causal and non-causal: 1.87-3.79x,
no shape slower. A/A within 0.5% except two rows (8-12%) on a shared node.

Batch x seqlen	Non-causal	Causal
1 x 128	0.067 -> 0.030 ms (2.26x)	0.070 -> 0.029 ms (2.40x)
1 x 1024	0.308 -> 0.081 ms (3.79x)	0.282 -> 0.081 ms (3.47x)
16 x 128	0.118 -> 0.063 ms (1.87x)	0.132 -> 0.060 ms (2.21x)
64 x 128	0.347 -> 0.175 ms (1.98x)	0.389 -> 0.167 ms (2.33x)

Tests

-k fp8_large_head: 9 passed on MI325X after the rebase onto main.

• test_mha_fwd_fp8_large_head_selection pins the bound. Catches a widened or
  reordered branch.
• test_mha_varlen_fp8_large_head runs d256 FP8 varlen against attention_ref
  and covers causal and non-causal. It catches this entry exceeding LDS.
  [Triton/Gluon] [Feature] MHA: add out_dtype to the Triton forward entry point #7232's test_mha_forward_out_dtype already covers FP8 causal varlen d256.

-k mid_head_selection: 14 skipped, unchanged (gfx950 only).

Not run: the sweep on MI300X; every reachable host was busy, had a profiler
interposition inflating timings ~6x, or a full disk. A narrower A/B there gave
157.645 -> 43.193 ms at 30720q x 32768k non-causal. Re-sweep before this entry
is ever proposed as gfx942's general default.

Details

Compared with default (BLOCK_M=128, BLOCK_N=64, num_warps=4,
num_stages=1, PRELOAD_V=false), fp8_large_head keeps the tile and changes
num_warps 4→8, num_stages 1→2, PRELOAD_V false→true, and adds
matrix_instr_nonkdim=16 and kpack=1. The speedup above is the whole entry
versus default. The 108-config sweep varies BLOCK_M, BLOCK_N, num_warps,
num_stages, and PRELOAD_V together; it does not isolate num_stages.

num_stages=2 is the LDS bound, not a claim that it is the sole source of the
gain: it double-buffers the K tile, and since
BLOCK_DMODEL_POW2 = next_power_of_2(v_head_dim), any head dim above 256
pads to 512 and the two buffers need 65568 B against a 65536 B LDS limit. Hence
== 256, not >= 256 — an earlier revision used >= 256 and turned
d320/d384/d512 FP8 from working-but-slow into OutOfResources.

The branch sits after dropout_or_fp32 and before small_head/mid_head
(head_dim_v <= 128, no overlap), following #4414 and #5858: a targeted entry
rather than changing default.

Testing

• Unit tests added or updated
• Performance benchmarks run
• Tested on MI250X
• Tested on MI300X

Documentation

• No documentation change is needed, with the reason stated below

The config space has no prose docs; the JSON and the _get_config branch are
the only places it is described, both updated here.

Dependencies

• No new third-party dependencies added

Breaking Changes

None. The branch is additive and gated on "fp8_large_head" in fwd_cfg, true
only for gfx942 after this change.

PR Self-Review

Reviewed revision

a41ecd1584060b9a2de7c1ee02615df232112536

Behavior and intent

On gfx942, FP8 Triton MHA forward at head_dim_v == 256 now uses fwd["fp8_large_head"] instead of fwd["default"]. Versus default the entry changes num_warps 4→8, num_stages 1→2, PRELOAD_V false→true, and adds matrix_instr_nonkdim=16 and kpack=1. The measured speedup is for that whole entry. Every other dtype, head dim, dropout/fp32, PE path and architecture keeps its current config.

Validation and unverified scope

On MI325X (gfx942), after rebasing onto main, python3 -m pytest op_tests/triton_tests/attention/test_mha.py -k "fp8_large_head or mid_head_selection" gave 9 passed, 14 skipped. Kernel timings are in Validation above. Not run: the 108-config sweep on MI300X, a single-knob ablation of num_stages, and FP8 e5m2 inputs (only e4m3 was timed).

PR-specific question

The only key that consumes fp8_large_head is gfx942 _get_config with no dropout, no fp32, no PE, an FP8 dtype from get_fp8_dtypes() and head_dim_v == 256. test_mha_fwd_fp8_large_head_selection pins that d255, d320, d512, d128, bf16 at d256 and dropout at d256 keep their previous config. On other architectures the test asserts the JSON has no fp8_large_head key, so their selection is unchanged.

Changed op tests

op_tests/triton_tests/attention/test_mha.py: test_mha_fwd_fp8_large_head_selection pins the dispatch bounds above. test_mha_varlen_fp8_large_head runs d256 FP8 varlen against attention_ref, causal and non-causal. #7232's test_mha_forward_out_dtype already covers FP8 causal varlen d256. op_tests/triton_tests/attention/mha_test_utils.py: shared varlen FP8 setup used by this test and test_mha_varlen_fp8_gluon; no new cases.

[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: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 7516 --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.

[aiter-bot]

READY FOR HUMAN REVIEW

PR Self-Review

Revision a41ecd158406 · mode shadow

Author self-review accepted. Readiness report:

Author answers

• PASS — current-SHA answers are present and structurally specific

Show the author's self-review answers

Behavior and intent

On gfx942, FP8 Triton MHA forward at `head_dim_v == 256` now uses `fwd["fp8_large_head"]` instead of `fwd["default"]`. Versus `default` the entry changes `num_warps` 4→8, `num_stages` 1→2, `PRELOAD_V` false→true, and adds `matrix_instr_nonkdim=16` and `kpack=1`. The measured speedup is for that whole entry. Every other dtype, head dim, dropout/fp32, PE path and architecture keeps its current config.

Validation and unverified scope

On MI325X (gfx942), after rebasing onto main, `python3 -m pytest op_tests/triton_tests/attention/test_mha.py -k "fp8_large_head or mid_head_selection"` gave 9 passed, 14 skipped. Kernel timings are in Validation above. Not run: the 108-config sweep on MI300X, a single-knob ablation of `num_stages`, and FP8 e5m2 inputs (only e4m3 was timed).

PR-specific question

The only key that consumes `fp8_large_head` is gfx942 `_get_config` with no dropout, no fp32, no PE, an FP8 dtype from `get_fp8_dtypes()` and `head_dim_v == 256`. `test_mha_fwd_fp8_large_head_selection` pins that d255, d320, d512, d128, bf16 at d256 and dropout at d256 keep their previous config. On other architectures the test asserts the JSON has no `fp8_large_head` key, so their selection is unchanged.

Changed op tests

`op_tests/triton_tests/attention/test_mha.py`: `test_mha_fwd_fp8_large_head_selection` pins the dispatch bounds above. `test_mha_varlen_fp8_large_head` runs d256 FP8 varlen against `attention_ref`, causal and non-causal. #7232's `test_mha_forward_out_dtype` already covers FP8 causal varlen d256. `op_tests/triton_tests/attention/mha_test_utils.py`: shared varlen FP8 setup used by this test and `test_mha_varlen_fp8_gluon`; no new cases.

PR state

• PASS — pull request is marked ready for review

Copilot findings

• PASS — no active Copilot finding lacks a substantive author reply
• ADVISORY — Copilot did not review the current head SHA

Code and op-test hygiene

• PASS — no committed artifact or conflict marker detected

Required CI

• PASS — Black: passed
• PASS — Ruff: passed
• PASS — Repository dependency: passed
• PASS — Aiter Test Gate: passed

_{Shadow mode reports what would block but keeps the status successful. Set AITER_PR_SELF_REVIEW_MODE=enforce only after reviewing real false positives.}

[akii96]

Thanks for the good work! I took a swing at checking a few things in the hopes that once the needed changes are in before other reviewers/maintainers take a look.

[Copilot]

🟡 Changes recommended

The new test duplicates input setup, quantization, and reference checks from an existing test, leaving shared logic to be extracted into mha_test_utils.py while keeping backend-specific calls separate.

1 open issue

• Extract duplicated varlen FP8 test setup into shared utilities · New

What changed in this PR

This PR adds a targeted gfx942 configuration for Triton MHA forward with FP8 inputs and a value head dimension of 256, without changing the general default.

Changes:

• Adds fp8_large_head with eight warps and two pipeline stages.
• Selects it only for FP8 at head_dim_v == 256, preserving existing dispatch precedence.
• Adds selection-boundary tests and causal/non-causal variable-length correctness tests.

File	Description
op_tests/​triton_tests/​attention/​test_mha.py	Adds config-selection and FP8 d256 correctness tests.
aiter/​ops/​triton/​configs/​gfx942/​triton/​attention/​mha/​DEFAULT.json	Defines the gfx942 fp8_large_head tuning entry.
aiter/​ops/​triton/​_triton_kernels/​attention/​mha.py	Adds guarded selection of the new FP8 configuration.

_{Review effort: Balanced}

[phambinhfin]

@akii96 @Boss2002n Could u take a look again?

[phambinhfin]

@Boss2002n this ticket prevent vllm-project/vllm#62913 move on. Could you take alook if you have time?

[akii96]

Thanks for working through the first round 🙏

You might want to look at two main-side PRs are worth a look, since the branch predates both:

• #7232 added test_mha_forward_out_dtype (covers FP8 causal varlen d256) and moved _flash_attn_forward / get_arch to the top-level imports. So the "only test reaching d256" line in the description is out of date, and the in-function imports can go. Your test is still worth keeping: it checks against attention_ref and also covers non-causal.
• #6502 changed _flash_attn_forward to also return rng_state, so the 5-value unpack in the new test will fail after the rebase (inline).

Happy to take a second look once it's rebased!

[akii96]

Nit (PR description only): compared to default, this changes num_warps (4→8), num_stages (1→2) and PRELOAD_V (false→true), and adds matrix_instr_nonkdim=16 and kpack=1. The description only mentions the first two and credits the gain to num_stages=2. Could you list all of them, and only credit num_stages if you have a sweep that shows it?

[phambinhfin]

Updated the description. Versus default this entry changes num_warps 4→8, num_stages 1→2, PRELOAD_V false→true, and adds matrix_instr_nonkdim=16 and kpack=1.

The reported speedup is the whole entry versus default. The 108-config sweep varies BLOCK_M, BLOCK_N, num_warps, num_stages, and PRELOAD_V together, so it does not isolate num_stages. num_stages=2 stays in the description only as the LDS reason for head_dim_v == 256: above that, BLOCK_DMODEL_POW2 pads to 512 and the double-buffered K tile needs 65568 B against a 65536 B limit.

[akii96]

Minor nit: after rebasing, main already imports _flash_attn_forward and get_arch at the top of this file, so these can go maybe please check

[phambinhfin]

Dropped the in-function imports. After the rebase, _flash_attn_forward and get_arch are already imported at the top of the file.

[akii96]

Needs a fix after rebase: #6502 made _flash_attn_forward return 6 values (it added rng_state), so this unpack will fail. _flash_attn_forward(...)[0] avoids depending on the length.

[phambinhfin]

Rebased onto main and switched to _flash_attn_forward(...)[0], so the test no longer depends on the return length after #6502. -k fp8_large_head: 9 passed on MI325X after the rebase.

[phambinhfin]

Rebased onto main.

• _flash_attn_forward(...)[0] after [Triton/Gluon] [CI] Streamline Flash Attentions api #6502, and dropped the in-function imports that [Triton/Gluon] [Feature] MHA: add out_dtype to the Triton forward entry point #7232 already put at the top of the file.
• Kept test_mha_varlen_fp8_large_head: it checks against attention_ref and covers non-causal. The description no longer says it is the only d256 FP8 path; [Triton/Gluon] [Feature] MHA: add out_dtype to the Triton forward entry point #7232's test_mha_forward_out_dtype already covers FP8 causal varlen d256.
• Description lists every field that differs from default. The speedup is for the whole entry; there is no single-knob sweep that isolates num_stages. num_stages=2 is cited only as the LDS bound.

-k "fp8_large_head or mid_head_selection" after the rebase: 9 passed, 14 skipped, on MI325X. test_mha_forward_out_dtype FP8 causal varlen d256 also passed.