[Triton/Gluon] [CI] Streamline Flash Attentions api by micmelesse · Pull Request #6502 · ROCm/aiter

micmelesse · 2026-10-13T20:58:03Z

Motivation

This PR reorganizes the aiter code that upstream FlashAttention uses. It follows the discussions in Dao-AILab/flash-attention#2981 and #5945. FlashAttention uses aiter's Triton attention through a private package. This PR moves that code and exposes it through the public Triton MHA API. FlashAttention can then use a public entry point. The matching FlashAttention change is Dao-AILab/flash-attention#3064. With it, FlashAttention can use aiter's Triton kernels or aiter's default dispatcher. The default dispatcher includes FMHA v3.

The PR also moves the test and bench commands of the Flash Attention workflow into scripts. The post-merge suite ran each test file in one pytest process. On main, a native abort stops the FA2 file at 4%. The abort hides all later tests.

Technical Details

The Dao-AI kernels move from _triton_kernels/flash_attn_triton_amd/ to _triton_kernels/attention/dao_ai/. The old package forwards what FlashAttention main imports: flash_attn_2, flash_attn_3 and utils.
flash_attn_func and flash_attn_varlen_func in aiter/ops/triton/attention/mha.py accept implementation="default" or implementation="dao_ai". None uses the default from mha_set_impl. Backward uses the implementation and the dropout RNG state of its forward. MXFP8 and block-scale FP8 calls refuse implementation="dao_ai". aiter/ops/mha.py does not change.
mha.py adds _flash_attn_backward next to _flash_attn_forward. Both take implementation=. The two autograd functions call them. FlashAttention calls them with implementation="dao_ai" and does not import the Dao-AI package.
mha_v3.py adds _flash_attn_forward, _flash_attn_backward and _flash_attn_combine. Its autograd functions and its KV-cache path call them. FlashAttention's FA3 interface calls them and does not import the Dao-AI package.
Dao-AI returns rng_state on the device and the dropout record in q.dtype, as FlashAttention's fake ops declare for torch.compile. rng_state is filled on the device, and the backward kernels read it there. The varlen backward does not check cu_seqlens values on the host. A dropout backward does not sync and captures in a CUDA graph. Dense dropout without a returned record no longer fails an assertion.
.github/scripts/flash_attention/ holds install.sh, test.sh and bench.sh. install.sh is install_flash_attention.sh with no change. Each test and bench step of the workflow calls one of these scripts. CI and a local container run the same commands.
test.sh runs one pytest process for each test function of the full suites. A failure or an abort in one function does not stop the other functions. The script fails if a line fails.
The PR suites select the same tests with the same flags as before. They now run all lines after a failure.
test.sh deselects 28 cases of test_flash_attn_output that fail on main. All 28 use dropout and ALiBi. 24 are causal MHA cases with head dimension 59 or 64. 4 are non-causal GQA cases with head dimension 40. Some of these cases fail, and some abort the process. #6859 will remove the deselect lines when it fixes these cases.
The path filters of the workflow include attention/dao_ai/** and .github/scripts/flash_attention/**.

Test Plan

All tests ran on MI355X, gfx950, in the base image of the workflow. The branch was installed as the workflow installs it. The tests used test.sh with the flags of the workflow.

test.sh v2-pr, v2-full, v3-pr and v3-full against FlashAttention main, as aiter CI runs them.
test.sh v2-full and v3-full against CI: Temporarily move Aiter build to linux-aiter-mi300x-1 #3064.
test_mha_dao_ai.py, test_mxfp8_attention_kernel.py and test_block_scale_fp8_attention.py.
The two bench commands of bench.sh, two runs for this branch and two runs for main, on one GPU.

Test Result

All four suites passed against FlashAttention main.
In test_flash_attn_output, 5,228 cases passed, 24 skipped and 28 were deselected. All other functions passed or skipped themselves.
v2-full and v3-full passed against CI: Temporarily move Aiter build to linux-aiter-mi300x-1 #3064 with the same counts. The adapter tests of CI: Temporarily move Aiter build to linux-aiter-mi300x-1 #3064 passed.
test_mha_dao_ai.py, test_mxfp8_attention_kernel.py and test_block_scale_fp8_attention.py passed. This includes the new dropout, selector and routing tests.
The MHA bench has 380 cases. Against main, the geometric mean of the throughput ratio is 1.000, with the best of two runs on each side. Main against itself varies from 0.82 to 1.36 for each case.
The FlashAttention bench gives 369.5 to 369.8 backward TFLOPS on both sides.

PR Self-Review

Reviewed revision

ae15b8be1

Behavior and intent

The Dao-AI Triton attention moves from a private package to attention/dao_ai/. The public Triton MHA wrappers select it for each call with implementation="dao_ai". The old package path forwards its imports, so FlashAttention releases that import it keep working. The Flash Attention workflow runs its tests and benches from scripts. CI and local runs use the same commands. The post-merge suite runs one pytest process for each test function. An abort no longer hides later tests. 28 known failures on main are deselected by node ID until #6859 fixes them.

Validation and unverified scope

All validation ran on MI355X, gfx950, in the base image of the workflow. v2-full, v2-pr, v3-full and v3-pr passed against FlashAttention main. v2-full and v3-full passed against #3064. test_mha_dao_ai.py and the MXFP8 and block-scale tests passed. The MHA bench shows no change against main. gfx942, RDNA3 and the Windows and Triton-only smoke jobs were not tested locally. CI runs them. The benches ran with the commands of bench.sh, not through the script.

PR-specific question

The pull_request event starts the workflow, not pull_request_target. A fork PR therefore has a read-only token and no secrets. The new scripts run inside the test container through docker exec. Their environment has no secrets. They run the same commands as the old inline steps. The Docker login steps and the artifact steps do not change. Fork PR code already ran in these containers, because the jobs build and test the aiter of the PR.

Changed op tests

test_mha_dao_ai.py now has one test for each property of the Dao-AI implementation. Each test sweeps the inputs of its property.
test_flash_attn_func_correctness and test_flash_attn_varlen_func_correctness replace test_mha_dao_ai and keep its cases, split by function. Each adds dropout rows that compare output and gradients with a reference that applies the dropout record of the forward. The dropout rows also check that the gradients are closer to that reference than to one with a shifted mask, so they fail if backward loses the RNG state of the forward. They replace test_mha_dao_ai_dropout_scores.
test_flash_attn_with_kvcache_correctness is new. It compares flash_attn_with_kvcache with attention over the valid cache entries of each batch, and checks that the new keys and values are appended to the cache. No other test covers this function.
test_mha_dao_ai_contract calls each function and checks the shapes, dtypes and devices of its returns. _flash_attn_forward and _flash_attn_backward return what FlashAttention's fake ops declare.
test_mha_dao_ai_routing checks that a backward runs the implementation of its forward after mha_set_impl changes, dense and varlen.
test_mha_dao_ai_refusal checks each unsupported call: MXFP8 and block-scale FP8 with dao_ai, unknown selectors, the Dao-AI varlen block_table and out, and KV-cache softcap and split counts. No forward runs.
test_mha_dao_ai_capture replaces the two graph tests. It captures each function in a CUDA graph under both default devices, including dropout training steps, and compares the replay with eager.
test_mha.py, test_mha_fp8.py, test_mha_v3.py, test_block_scale_fp8_attention.py, bench_mha.py, bench_sage.py, bench_fav3_sage.py and bench_block_scale_fp8_attention.py change only import paths.

github-actions · 2026-10-13T20:58:41Z

🏷️ 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 6502 --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.

Copilot

🟡 Changes recommended

Integration CI filters still target the old kernel directory, and the new routing suite is outside the required Triton test folder.

2 open issues

What changed in this PR

This PR exposes Dao-AI FlashAttention selection through Aiter’s public Triton MHA API and relocates its internal package under attention/dao_ai/.

Changes:

Adds per-call implementation selection and preserves implementation, RNG state, and determinism for backward.
Moves internal modules and updates imports, with a compatibility package for existing consumers.
Adds routing tests and documents the API and package layout.

File	Description
`op_tests/triton_tests/attention/test_mha_v3.py`	Updates the FP8 utility import.
`op_tests/triton_tests/attention/test_mha_fp8.py`	Updates the FP8 utility import.
`op_tests/triton_tests/attention/test_mha_dao_ai.py`	Tests per-call implementation selection.
`op_tests/test_mha_backend_dispatch.py`	Adds CPU routing and backward-state tests.
`op_tests/op_benchmarks/triton/bench_sage.py`	Updates the FlashAttention import.
`op_tests/op_benchmarks/triton/bench_mha.py`	Updates imports and implementation help text.
`op_tests/op_benchmarks/triton/bench_fav3_sage.py`	Updates the FlashAttention import.
`aiter/ops/triton/README.md`	Documents selection and compatibility behavior.
`aiter/ops/triton/attention/unified_attention.py`	Updates the architecture utility import.
`aiter/ops/triton/attention/mha.py`	Adds per-call selection and backward-state handling.
`aiter/ops/triton/attention/mha_v3.py`	Uses the relocated internal package.
`aiter/ops/triton/_triton_kernels/flash_attn_triton_amd/__init__.py`	Preserves compatibility imports.
`aiter/ops/triton/_triton_kernels/attention/fav3_sage_attention.py`	Updates the shared attention helper import.
`aiter/ops/triton/_triton_kernels/attention/dao_ai/utils.py`	Relocates existing FlashAttention utilities.
`aiter/ops/triton/_triton_kernels/attention/dao_ai/interface_v3.py`	Updates relocated FA3 interface imports.
`aiter/ops/triton/_triton_kernels/attention/dao_ai/interface_v2.py`	Updates relocated FA2 interface imports.
`aiter/ops/triton/_triton_kernels/attention/dao_ai/fwd_prefill.py`	Updates prefill kernel imports.
`aiter/ops/triton/_triton_kernels/attention/dao_ai/fwd_decode.py`	Updates decode kernel imports.
`aiter/ops/triton/_triton_kernels/attention/dao_ai/common.py`	Updates common helper imports.
`aiter/ops/triton/_triton_kernels/attention/dao_ai/bwd.py`	Updates backward kernel imports.
`aiter/ops/triton/_triton_kernels/attention/dao_ai/__init__.py`	Exposes interfaces from the relocated package.
`.github/instructions/aiter-ops-triton.instructions.md`	Updates documented kernel and test paths.

_{Review effort: Balanced}

Copilot

🟡 Changes recommended

The public-wrapper, PyTorch isolation, and shared-helper requirements remain unresolved.

3 open issues

_{Review effort: Balanced}

aiter-bot · 2026-11-02T21:26:29Z

WAITING ON AUTHOR

PR Self-Review

Revision 8948c3dbe5d4 · mode shadow

Waiting for the PR author. The readiness report will appear after the self-review answers are complete for this revision.

Questions to answer in the PR description

What behavior does this PR change, and why is the existing behavior insufficient?
What did you validate, and what remains unverified? Include commands, hardware/configuration and results.
What permissions and untrusted-input boundary does this CI change have, and can fork PR code or text reach a privileged execution path?
For every changed op_tests/ file, what behavior does it cover and why is it not redundant? Write None if no op test changed.

Edit the PR Self-Review section in the PR description. Include the current revision SHA; this comment refreshes automatically.

Missing or invalid answers

FAIL — Reviewed revision must contain at least the first 7 characters of the current head SHA

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

Copilot

🟡 Changes recommended

The new Dao-AI RNG handoff lacks numerical backward validation with nonzero dropout: GPU cases use zero dropout, routing tests mock backward calls, and dense and varlen gradient checks against a reference using the returned dropout mask are still needed.

4 open issues

_{Review effort: Balanced}

brunomazzottiamd · 2026-11-06T17:23:22Z

We have a conflict in aiter/ops/triton/attention/mha.py. We can merge once it's resolved!

Copilot