[Triton/Gluon] Fix Flash Attention main ci by micmelesse · Pull Request #6859 · ROCm/aiter

micmelesse · 2026-10-23T03:09:45Z

Summary

Fix gfx950 causal Flash Attention backward failures in aiter's Dao-AI Triton path. The backward fix passed the complete ordered FA2 test file on a local gfx950 node.

Motivation

Flash Attention CI on aiter main found incorrect gradients and native aborts in causal backward cases. The previous dQ mask handling and fused backward program were sensitive to gfx950 code generation. Dao-AI dropout backward also needs the RNG state returned by forward.

Changes

Use the shared attention mask to exclude invalid dQ logits before exponentiation.
Keep dK/dV and dQ in explicit causal phase helpers and run both in each program, without changing the gfx950 tile sizes.
Include the window bounds in the causal backward autotune key so windowed and unwindowed calls can select separate configs.
Pass the Dao-AI forward RNG state through dense and varlen autograd into backward.
Add a gfx950 regression test for dense and varlen causal MHA, MQA, and GQA backward with dropout and ALiBi.

Performance (if applicable)

The comparison below uses the PR base and current head 4c315368d on the same local gfx950 GPU with bench_mha.py -impl dao_ai -metric throughput and equal sequence lengths.

Hardware	Workload/configuration	Before	After	Change
gfx950	bf16 Dao-AI causal varlen backward; B=1, Hq=Hk=128, Sq=Sk=8192, d=56, window=(1024, 0)	163.3 TFLOPS	135.7 TFLOPS	−16.9%

The d56 windowed backward case remains slower than the PR base. A 380-case sweep of a preparatory scratch revision found six cases more than 10% below the base. The full sweep has not been rerun on this head.

Testing

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

Documentation

Docstrings updated
User guide updated
Performance guide updated
No documentation change is needed, with the reason stated below

The change is internal to the existing Flash Attention implementation and does not change its public API.

Dependencies

No new third-party dependencies added
New dependencies are listed and justified below

Breaking Changes

None.

PR Self-Review

Reviewed revision

4c315368d3fabd9b2f7bf5865867704ae7b4de90

Behavior and intent

The PR corrects gfx950 causal backward gradients and prevents the tested native aborts. It masks invalid dQ logits before exponentiation, simplifies the causal phase bodies, gives dropout backward the forward RNG state, and separates windowed and unwindowed autotune entries. The previous behavior failed the FA2 correctness cases described below.

Validation and unverified scope

On a local gfx950 node, 832eafce1 passed the complete ordered FA2 test file with FLASH_ATTENTION_TRITON_AMD_AUTOTUNE=0: 60,576 passed and 27,004 skipped. The 12 new regression cases and four previously failing FA cases also passed.
Current head 4c315368d passed 41 targeted Dao-AI tests with autotuning enabled. The complete FA2 file has not been rerun after the autotune-key change.
The d56 performance result above is from the current head. Five other cases below the 10% target were seen only in a preparatory scratch sweep.
MI250X and MI300X were not tested locally. The compiler lowering behind the earlier bad store address remains unidentified.

PR-specific question

The new op test checks for incorrect or nonfinite dQ, dK, and dV, or an abort, in gfx950 causal backward with dropout and ALiBi. Its dense 113×203 d59 case covers the main CI gradient failure. Existing Dao-AI backward tests do not combine these conditions, and the existing dropout-scores test checks forward output rather than gradients.

Changed op tests

op_tests/triton_tests/attention/test_mha_dao_ai.py adds a public-wrapper autograd comparison against a PyTorch reference for six shapes in both dense and varlen layouts, including MHA, MQA, GQA, unequal sequence lengths, and padded head dimensions. This covers the failing backward combinations absent from the existing tests.

Reuse the shared causal and sliding-window predicate in fused dQ, retain the backward end_n bound, and mask logits before exponentiation. Add a low-level dropout and ALiBi regression test.

github-actions · 2026-10-23T03:10:13Z

🏷️ 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 6859 --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

The regression test must use the public wrapper and specify an explicit dtype for its ALiBi slopes.

2 open issues

What changed in this PR

This PR addresses Flash Attention CI failures by changing dQ backward masking and adding a focused regression test.

Changes:

Reuses attention_mask and masks invalid logits before exponentiation.
Adds gfx950 backward checks for causal attention with dropout and ALiBi, including a padded head dimension.

File	Description
`op_tests/triton_tests/attention/test_mha_dao_ai.py`	Adds two regression cases that compare gradients against a reference.
`aiter/ops/triton/_triton_kernels/flash_attn_triton_amd/bwd.py`	Replaces local dQ masking logic with the shared helper before exponentiation.

_{Review effort: Balanced}

Copilot

Copilot was unable to review this pull request because the user who requested the review has insufficient remaining quota.

Copilot

Copilot was unable to review this pull request because the user who requested the review has insufficient remaining quota.

Copilot

Copilot was unable to review this pull request because the user who requested the review has insufficient remaining quota.

Compute dK/dV and dQ in separately compiled phases while keeping the 128-row tile. Apply the causal probability mask after the exponential.

Copilot

Copilot was unable to review this pull request because the user who requested the review has insufficient remaining quota.

aiter-bot · 2026-10-29T14:47:29Z

NOT READY FOR HUMAN REVIEW

PR Self-Review

Revision 4c315368d3fa · 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

The PR corrects gfx950 causal backward gradients and prevents the tested native aborts. It masks invalid dQ logits before exponentiation, simplifies the causal phase bodies, gives dropout backward the forward RNG state, and separates windowed and unwindowed autotune entries. The previous behavior failed the FA2 correctness cases described below.

Validation and unverified scope

- On a local gfx950 node, `832eafce1` passed the complete ordered FA2 test file with `FLASH_ATTENTION_TRITON_AMD_AUTOTUNE=0`: 60,576 passed and 27,004 skipped. The 12 new regression cases and four previously failing FA cases also passed.
- Current head `4c315368d` passed 41 targeted Dao-AI tests with autotuning enabled. The complete FA2 file has not been rerun after the autotune-key change.
- The d56 performance result above is from the current head. Five other cases below the 10% target were seen only in a preparatory scratch sweep.
- MI250X and MI300X were not tested locally. The compiler lowering behind the earlier bad store address remains unidentified.

PR-specific question

The new op test checks for incorrect or nonfinite dQ, dK, and dV, or an abort, in gfx950 causal backward with dropout and ALiBi. Its dense 113×203 d59 case covers the main CI gradient failure. Existing Dao-AI backward tests do not combine these conditions, and the existing dropout-scores test checks forward output rather than gradients.

Changed op tests

`op_tests/triton_tests/attention/test_mha_dao_ai.py` adds a public-wrapper autograd comparison against a PyTorch reference for six shapes in both dense and varlen layouts, including MHA, MQA, GQA, unequal sequence lengths, and padded head dimensions. This covers the failing backward combinations absent from the existing tests.

PR state

WAIT — pull request is still a draft

Copilot findings

PASS — no active Copilot finding lacks a substantive author reply

Code and op-test hygiene

PASS — no committed artifact or conflict marker detected

Required CI

WAIT — Black: missing
WAIT — Ruff: missing
WAIT — Repository dependency: missing
PASS — Aiter Test Gate: passed
Required CI is incomplete; this PR is not ready for human review.

_{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 causal phases duplicate ALiBi, dropout, and gradient logic rather than reuse shared helpers, violating repository reuse rules and requiring future fixes in multiple places.

2 open issues

_{Review effort: Balanced}

micmelesse · 2026-11-04T15:39:19Z

Sorry @azaidy and @brunomazzottiamd, I am putting this in draft. I want to get #6502 merged first. I will ping you when this pr is ready.

brunomazzottiamd · 2026-11-04T15:41:42Z

Sorry @azaidy and @brunomazzottiamd, I am putting this in draft. I want to get #6502 merged first. I will ping you when this pr is ready.

No worries @micmelesse! I started reviewing this PR 10 minutes ago! I'll publish some minor nitpicks related to the test, address them at your judgment.

brunomazzottiamd

Sharing minor nitpicks.

brunomazzottiamd · 2026-11-04T15:27:13Z

+def test_mha_dao_ai_causal_dropout_alibi_backward(
+    dao_ai_impl, batch, seqlen_q, seqlen_k, head_dim, nheads_k, varlen
+):
+    from aiter.ops.triton.utils._triton.arch_info import get_arch


Can we import get_arch at module top level instead of inside a test function?

brunomazzottiamd · 2026-11-04T15:28:46Z

+    if get_arch() != "gfx950":
+        pytest.skip("gfx950 causal backward configuration regression")


Suggestion / non-blocker: Implement the skip with an annotation.

@pytest.mark.skipif(get_arch() != "gfx950", reason="gfx950 causal backward configuration regression")

brunomazzottiamd · 2026-11-04T15:36:39Z

+    for actual, expected in ((dq, dq_ref), (dk, dk_ref), (dv, dv_ref)):
+        assert torch.isfinite(actual).all()
+        torch.testing.assert_close(actual, expected, atol=1e-2, rtol=1e-2)


Can you please add messages to this assertions, saying which tensor is failing (dq, dk, dv)? I think this improve troubleshooting in case of failures.

Fix gfx950 Flash Attention backward with shared mask

28618e7

Reuse the shared causal and sliding-window predicate in fused dQ, retain the backward end_n bound, and mask logits before exponentiation. Add a low-level dropout and ALiBi regression test.

micmelesse changed the title ~~[Triton/Gluon] Fix gfx950 Flash Attention backward with shared mask~~ [Triton/Gluon] Fix Flash Attention main ci Oct 23, 2026

micmelesse marked this pull request as ready for review October 23, 2026 04:14

micmelesse requested review from a team and a balanced review from Copilot October 23, 2026 04:14

github-actions Bot added the Triton/Gluon label Oct 23, 2026

Copilot started reviewing on behalf of micmelesse October 23, 2026 04:15

View session

Copilot AI reviewed Oct 23, 2026

View reviewed changes

Comment thread op_tests/triton_tests/attention/test_mha_dao_ai.py Resolved

Comment thread op_tests/triton_tests/attention/test_mha_dao_ai.py Resolved

micmelesse marked this pull request as draft October 23, 2026 15:02

micmelesse force-pushed the micmelesse/fix-fa-main-ci branch from 7ca2392 to 1da3992 Compare October 27, 2026 17:29

micmelesse marked this pull request as ready for review October 27, 2026 17:59

Copilot AI balanced review requested due to automatic review settings October 27, 2026 17:59

Copilot AI reviewed Oct 27, 2026

Copilot AI balanced review requested due to automatic review settings October 27, 2026 19:36

Copilot AI reviewed Oct 27, 2026

Copilot AI balanced review requested due to automatic review settings October 27, 2026 20:28

Copilot AI reviewed Oct 27, 2026

zufayu requested review from a team and k50112113 October 28, 2026 02:34

zufayu assigned k50112113 Oct 28, 2026

micmelesse marked this pull request as draft October 28, 2026 14:52

Address Copilot comments

da7cb76

micmelesse force-pushed the micmelesse/fix-fa-main-ci branch from 21bc0f3 to da7cb76 Compare October 28, 2026 17:58

zufayu removed request for a team and k50112113 October 29, 2026 00:02

Simplify gfx950 causal Flash Attention backward

6029a41

Compute dK/dV and dQ in separately compiled phases while keeping the 128-row tile. Apply the causal probability mask after the exponential.

micmelesse marked this pull request as ready for review October 29, 2026 14:46

Copilot AI balanced review requested due to automatic review settings October 29, 2026 14:46

Copilot AI reviewed Oct 29, 2026

micmelesse marked this pull request as draft October 29, 2026 14:47

micmelesse added 3 commits October 29, 2026 15:46

Restore causal bounds and launch backward phases together

becb128

Run causal backward phases in each program

832eafc

Include causal window in backward autotune key

4c31536

micmelesse marked this pull request as ready for review November 2, 2026 20:00

Copilot AI balanced review requested due to automatic review settings November 2, 2026 20:00

Copilot started reviewing on behalf of micmelesse November 2, 2026 20:01

View session

Copilot AI reviewed Nov 2, 2026

View reviewed changes

Comment thread aiter/ops/triton/_triton_kernels/flash_attn_triton_amd/bwd.py Resolved

Comment thread aiter/ops/triton/_triton_kernels/flash_attn_triton_amd/bwd.py Resolved

micmelesse requested a review from azaidy November 2, 2026 21:58

azaidy assigned azaidy and brunomazzottiamd and unassigned k50112113 Nov 4, 2026

micmelesse marked this pull request as draft November 4, 2026 15:39

brunomazzottiamd reviewed Nov 4, 2026

View reviewed changes

micmelesse mentioned this pull request Nov 4, 2026

[Triton/Gluon] [CI] Streamline Flash Attentions api #6502

Merged

zufayu removed the request for review from azaidy November 5, 2026 00:03

		if get_arch() != "gfx950":
		pytest.skip("gfx950 causal backward configuration regression")

Conversation

micmelesse commented Oct 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Motivation

Changes

Performance (if applicable)

Testing

Documentation

Dependencies

Breaking Changes

PR Self-Review

Reviewed revision

Behavior and intent

Validation and unverified scope

PR-specific question

Changed op tests

Uh oh!

github-actions Bot commented Oct 23, 2026

🏷️ CI Guide

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

🟡 Changes recommended

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Uh oh!

aiter-bot commented Oct 29, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

NOT READY FOR HUMAN REVIEW

PR Self-Review

Author answers

Behavior and intent

Validation and unverified scope

PR-specific question

Changed op tests

PR state

Copilot findings

Code and op-test hygiene

Required CI

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

🟡 Changes recommended

Uh oh!

Uh oh!

Uh oh!

micmelesse commented Nov 4, 2026

Uh oh!

brunomazzottiamd commented Nov 4, 2026

Uh oh!

brunomazzottiamd left a comment

Choose a reason for hiding this comment

Uh oh!

brunomazzottiamd Nov 4, 2026

Choose a reason for hiding this comment

Uh oh!

brunomazzottiamd Nov 4, 2026

Choose a reason for hiding this comment

Uh oh!

brunomazzottiamd Nov 4, 2026

Choose a reason for hiding this comment

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

micmelesse commented Oct 23, 2026 •

edited

Loading

aiter-bot commented Oct 29, 2026 •

edited

Loading