Add flex_attention (score_mod / mask_mod) on the generic flash-attention kernel - #931

Draft

RichardChamberlain1 wants to merge 22 commits into

mainfrom

rchamber/flex_attention

Commits on Jul 29, 2026

[feat] Add flex_attention (score_mod/mask_mod) on the generic flash-attention kernel

RichardChamberlain1
and
claude
committed
7f03cfb

Commits on Jul 30, 2026

[test] flex_attention: align bench output with flash_attn, expand shape sweep

RichardChamberlain1
and
claude
committed
a9697a5
[test] flex_attention: apply black formatting (CI style gate)

RichardChamberlain1
and
claude
committed
903cfff
Merge branch 'main' into rchamber/flex_attention
RichardChamberlain1
authored
a6c786d

Commits on Aug 3, 2026

[feat] Add flash-attention forward on the layout API (gfx950)

RichardChamberlain1
and
claude
committed
12f4e18
[feat] Layout attention: pipeline scheduler + vectorized DMA (gfx950)

RichardChamberlain1
and
claude
committed
a48a507
[wip] Pipeline scheduler cleanup + depth=2 decomposed softmax (in progress)

RichardChamberlain1
and
claude
committed
94c9454
Fix depth-2 layout attention pipeline for single- and multi-tile KV.

root
and
cursoragent
committed
59690e2

Commits on Aug 4, 2026

Improve layout attention pipeline scheduling and depth-2 stage wiring.

root
and
cursoragent
committed
92db086
Optimize pd2 layout pipeline barriers, waits, and inter-tile ReadK readahead.
root
committed
3004116

Commits on Aug 5, 2026

Experimental: pd2 flex layout manual cluster emit and dualwave scaffolding.

root
and
cursoragent
committed
a631e86

Commits on Aug 7, 2026

Replaced static loop which was causing register growth and hurting performance for buffering approch

root
and
cursoragent
committed
3cee3e2

Commits on Aug 13, 2026

Merge branch 'main' of https://github.com/ROCm/FlyDSL into rchamber/flex_attention
root
committed
af0b7bd

Commits on Aug 14, 2026

Fix MFMA 16x16x16 bf16/f16 lowering crash when copy atom is wider than MFMA operand

root
and
claude
committed
aa04ad2
[MFMA] Add 16x16x16 bf16/f16 support with fly-fix-bitcast-width pass

root
and
claude
committed
fc1da14
Merge branch 'rchamber/add_16x16x16_MFMA_support' into rchamber/flex_attention
root
committed
214222d
Switch flex attention to MFMA 32x32x16 with column reduction softmax

root
and
claude
committed
1ac1122

Commits on Aug 17, 2026

Remove LDS P bridge + K swizzle + V transpose reads (448 TFLOPS)

root
and
claude
committed
6ecdb33

Commits on Aug 18, 2026

Optimize flex attention: vectorized O store, native exp2, dead code removal (523 TFLOPS)

root
and
claude
committed
8d24a97
Softmax pre-scale + native exp2 + DMA/LDS read reorder (592 TFLOPS)

root
and
claude
committed
963b519
Add flex attention score/mask modifiers with tile-range skipping

root
and
claude
committed
01c5820
Refactor flex mods into FlexMod class hierarchy + expand test coverage

root
and
claude
committed
95a1bcb