Add flex_attention (score_mod / mask_mod) on the generic flash-attention kernel by RichardChamberlain1 · Pull Request #931 · ROCm/FlyDSL

RichardChamberlain1 · 2026-07-30T11:46:51Z

Summary

Ports PyTorch's flex_attention per-element hooks onto FlyDSL's existing dense f16/bf16 flash-attention forward kernel (Phase 1: dense, forward-only, correctness-first).

score_mod(score, b, h, q_idx, kv_idx) -> score and mask_mod(b, h, q_idx, kv_idx) -> bool as compile-time callables over fx scalars, inlined into the kernel MLIR and specialized per mod (keyed into the JIT cache by callable identity). No-mod path compiles byte-identical to today.
New kernels/attention/flex_attention.py: public flydsl_flex_attention(...) + built-in alibi_score_mod / sliding_window_mask_mod / causal_mask_mod.
Hooks live in GenericSoftmaxHelper (flash_attn_utils.py): apply_score_mod (PyTorch scale-then-mod semantics) runs before the mask; mask_mod extends apply_kv_mask. Mods threaded through the builder (incl. auto-tile + pad-mask recursion sites) and into traits.cache_tag.
Numerical-safety guards extended from causal-only to mask_mod: the c_neg_floor clamp and the epilogue 1/l reciprocal guard, so fully-masked rows yield 0/-inf (matching torch) instead of NaN. Fused gpfetch path is disabled when a mod is present.
scripts/run_benchmark.sh: flex_attention op entry with a 14-row shape sweep (seq-len ladder x all cases + GQA), mirroring the flash_attn benchmark pattern. test_flex_attention.py main() prints flash_attn-style aligned rows.

Test plan

pytest tests/kernels/test_flex_attention.py — 37 cases pass on MI300X (gfx942): no-mod parity, alibi, sliding-window, causal-via-mask; bf16/f16; MHA/GQA/D=64; odd/multi-tile seqlens (250/384/1024); cross-attention incl. fully-masked rows; LSE checks.
No regression: test_flash_attn_fwd.py LSE-dense (20 cases) pass — shared softmax/epilogue unaffected.
bash scripts/run_benchmark.sh --only flex_attention — all 14 rows emit parsed TFLOPS.
CI on the target runner.

Out of scope (future phases): block sparsity (BlockMask), backward, tensor-capturing mods, gfx950 dualwave/fp8 flex paths.

🤖 Generated with Claude Code

…ttention kernel Port PyTorch flex_attention's per-element hooks onto the dense f16/bf16 forward kernel: score_mod(score,b,h,q,kv) transforms each logit and mask_mod(b,h,q,kv) keeps/drops via where(mask,score,-inf). Mods are compile-time callables specialized per kernel and keyed into the JIT cache by identity, so the no-mod path is unchanged. - New kernels/attention/flex_attention.py: public flydsl_flex_attention() + built-in alibi_score_mod / sliding_window_mask_mod / causal_mask_mod. - Thread score_mod/mask_mod through the builder (incl. auto-tile and pad-mask dispatch recursion sites) and into traits/cache_tag; disable the fused gpfetch path when a mod is present. - Apply score_mod (scale-then-unscale to match PyTorch's qk*sm_scale semantics) before the mask hook; extend the -inf clamp and epilogue reciprocal guard to mask_mod builds so fully-masked rows yield 0 instead of NaN. - Tests (tests/kernels/test_flex_attention.py) and a run_benchmark.sh op entry. Verified on MI300X: 37 flex cases pass (no-mod parity, alibi, sliding-window, causal-via-mask; odd/multi-tile seqlens; cross-attention incl. fully-masked rows) with no regression in test_flash_attn_fwd.py LSE-dense. Co-Authored-By: Claude <noreply@anthropic.com>

…pe sweep Make the flex_attention benchmark consistent with the flash_attn one: - test_flex_attention.py main(): print a flash_attn-style aligned row (GPU header + "config shape | St | MaxErr MinCos | Time(us) TFLOPS") mirroring _fmt_result/_fmt_extra_normal_row, instead of the terse "TFLOPS=.. TB/s=.." line. run_benchmark.sh still parses TFLOPS via its flash-attn table regex. Also dedup: import _acc_metric/_flops from test_flash_attn_fwd instead of copying them (drops an unused F import). - run_benchmark.sh: expand DEFAULT_FLEX_ATTENTION_SHAPES from 3 rows to a curated 14-row sweep (seq-len ladder 2048/4096/8192 x all flex cases + two GQA rows); flex cases only, no cross-kernel baselines. Verified on MI300X: 37 correctness tests pass; run_benchmark.sh --only flex_attention parses real TFLOPS for every row. Co-Authored-By: Claude <noreply@anthropic.com>

Reformat test_flex_attention.py with black (line-length 120) to pass the "Check Python Code Style" CI gate on PR #931. Formatting-only: black explodes multi-arg calls one-per-line and normalizes whitespace. ruff already clean. 37 tests still pass on MI300X. Co-Authored-By: Claude <noreply@anthropic.com>

RichardChamberlain1 · 2026-07-30T14:21:51Z

FlyDSL flex_attention Performance Report

Run 20260730_150830 · samples=12 · iters=100 · warmup=10 · dtype=bf16

mi300 (SMC-SC-DI10-33.dh144.dcgpu)

GPU: AMD MI300 (gfx942), CU count 304, pciChipId 0x74a1 · roofline peak 1307 TFLOPS BF16 (full device)
Config: B=2, H=32, Hkv=32, D=128, bf16 · shapes S=2048 / 4096 / 8192

* causal cases run causal=True (half the KV work, causal-counted TFLOPS) — not comparable vs dense no_mod.

Shape S=2048 (B=2, H=32, Hkv=32, D=128)

case	lat µs (mean±std)	CV %	TFLOPS (mean±std)	vs no_mod
no_mod	352.2±1.9	0.6%	390.3±2.2	baseline
alibi	395.5±0.9	0.2%	347.5±0.8	+12.3%
sliding_window	383.4±2.1	0.5%	358.5±2.0	+8.9%
causal_via_mask	373.3±2.5	0.7%	368.2±2.5	+6.0%
prod_flash_attn	353.0±0.7	0.2%	389.3±0.8	+0.2%
torch_sdpa	464.9±2.2	0.5%	295.6±1.4	+32.0%
aiter	309.6±9.3	3.0%	444.2±13.4	−12.1%
prod_flash_attn_causal	283.3±2.8	1.0%	242.7±2.4	causal*
aiter_causal	175.4±2.2	1.3%	392.1±5.0	causal*

Shape S=4096 (B=2, H=32, Hkv=32, D=128)

case	lat µs (mean±std)	CV %	TFLOPS (mean±std)	vs no_mod
no_mod	1308.1±3.8	0.3%	420.3±1.2	baseline
alibi	1478.3±3.3	0.2%	371.9±0.8	+13.0%
sliding_window	1416.7±2.7	0.2%	388.1±0.7	+8.3%
causal_via_mask	1387.7±1.9	0.1%	396.2±0.5	+6.1%
prod_flash_attn	1309.8±3.3	0.3%	419.7±1.1	+0.1%
torch_sdpa	1706.5±3.3	0.2%	322.1±0.6	+30.5%
aiter	1156.9±2.2	0.2%	475.2±0.9	−11.6%
prod_flash_attn_causal	824.3±16.9	2.1%	333.7±6.6	causal*
aiter_causal	615.9±1.2	0.2%	446.4±0.8	causal*

Shape S=8192 (B=2, H=32, Hkv=32, D=128)

case	lat µs (mean±std)	CV %	TFLOPS (mean±std)	vs no_mod
no_mod	4677.9±8.2	0.2%	470.1±0.8	baseline
alibi	5276.7±15.5	0.3%	416.7±1.2	+12.8%
sliding_window	4975.1±24.2	0.5%	442.0±2.2	+6.4%
causal_via_mask	4951.8±4.0	0.1%	444.1±0.4	+5.9%
prod_flash_attn	4706.2±6.1	0.1%	467.3±0.6	+0.6%
torch_sdpa	7095.7±9.7	0.1%	309.9±0.4	+51.7%
aiter	4181.1±9.2	0.2%	525.9±1.2	−10.6%
prod_flash_attn_causal	3084.2±72.7	2.4%	356.7±8.5	causal*
aiter_causal	2422.5±5.8	0.2%	453.9±1.1	causal*

A from-scratch attention forward written on FlyDSL's CuTe-style layout API (make_tiled_mma / make_fragment_{A,B,C} / fx.gemm / fx.copy), independent of the legacy raw-MFMA flash_attn_generic path. Models MMA/pipeline structure on hgemm_layout_gfx950 and softmax numerics on kernels/norm/softmax_kernel. Per (batch, head, q-tile) workgroup: Q resident, KV loop with flash-attention online softmax (running m_i/l_i + O rescale), the QK-C-fragment -> PV-A-fragment bridge via LDS, and GEMM2 P@V. Verified vs torch SDPA on gfx950 (MI350): 8/8 cases pass (single/multi KV-tile, batch, multi-head, multi-q-tile, D=64/128, bf16/f16), all cos=1.0, max_err <= 1.2e-3. Phase 0 (dense forward, no flex mods). Constraints, all enforced in make_flex_attn_param: block_m=32, block_n/head_dim/seqlen_kv multiples of 32, MFMA 16x16x32 (16x16x16 hits an fx.gemm lowering bug on this build), V host-transposed. score_mod/mask_mod hooks, block_m>32, in-kernel V transpose, and LDS pipelining are follow-ups. Co-Authored-By: Claude <noreply@anthropic.com>

Restructure the layout-API flash-attention kernel to use a composable PipelineScheduler that manages the KV-loop stages (LoadKV, ReadKV, GEMM1, Softmax, BridgeP, GEMM2) via Wire declarations and cluster-based execution. Currently runs at force_depth=1 (monolithic, no decomposition). Performance: 111 → 227 TFLOPS at S=8192 on MI350 (gfx950), up from 5% to ~10% of peak. Key optimizations in this commit: - Vectorized global→LDS DMA via BufferCopyLDS128b (replacing scalar element-by-element copy) - LDS double-buffering with async prefetch (overlaps DMA with compute) - Generalized per-slot softmax row map (block_m unlocked to 32/64/128) Pipeline fixes for the non-staggered path: - Prologue: split LoadKV/ReadKV with s_waitcnt between them - lds_ring_slots: always ≥2 for double-buffered DMA (was =depth, broke at depth=1 by clobbering the current tile's LDS buffer with prefetch) - Epilogue drain: only at depth>1 when decomposition is active - Last-tile: always skip LoadKV in the cluster body New file: kernels/attention/pipeline.py — stage-based composable software pipeline framework with Wire/PipelineStage/PipelineScheduler, supporting monolithic and decomposed (multi-slot) stage execution, cluster assignment, prologue/main-loop/epilogue emission, and stagger infrastructure. Co-Authored-By: Claude <noreply@anthropic.com>

…gress) Pipeline scheduler (pipeline.py): - Removed unused pipeline builder code and standalone stage classes (-387 lines) - Removed all stage-name checks (LoadKV) — scheduler is now fully generic, using position-based prime stage identification (_prime_fn, _is_prime) - Clean prologue/main-loop/epilogue: depth=1 runs all stages synchronously per tile, depth>=2 primes first stage ahead and prefetches - lds_ring_slots always >=2 for double-buffered DMA - Epilogue skips last sub-stage of decomposed stages (prevents double-count) Kernel (flex_attention_layout_gfx950.py): - pipe_depth configurable from host wrapper (default=1) - Decomposed softmax: _start produces frag_P (prev v_p * corr), rescales l_i and O; _finish sums current v_p into l_i. Swapped decompose order so SoftmaxStart is in C1 (before BridgeP in C2) - V carry for depth>=2 via _gemm2_d1/_gemm2_d2 staticmethod dispatch - Hand-coded epilogue for depth>=2 (pipeline epilogue drain WIP) Status: depth=1 passes all 10 tests (124-227 TFLOPS). Depth=2 compiles and runs but produces partial NaN — V carry + fastmath -inf interaction under investigation. Proven-correct Python simulation exists. Co-Authored-By: Claude <noreply@anthropic.com>

Correct decomposed softmax ordering, lagged P@V bridging, epilogue drain, and separate loop-carried fragments; add pipe_depth to JIT kernel names and d2 layout correctness tests. Co-authored-by: Cursor <cursoragent@cursor.com>

Add pipeline hooks for entry handoffs, partial waitcnt policies, and sched_after; split BridgeP/Gemm2 substages without in-stage barriers; emit flash-style sched_group_barrier pairs on dual-wave softmax and GemM2 PV; extend layout tests and pd1/pd2/ps2 benchmark compare. Co-authored-by: Cursor <cursoragent@cursor.com>

…adahead. Tune memory-cluster vmcnt for overlapping K/V DMA, drop redundant in-cluster sync_after on flex substages with targeted lgkm waits, and add optional lighter C1→C2 boundary sync plus hoisted ReadK after C3 for staggered depth-2 emit.

…lding. Refactor depth-2 multi-tile path to use pipeline memory clusters with hand C1/C3, readahead, and ring LDS slots; add dualwave main-body and layout adapters for future j+=2 scheduling. Tests updated for Skv32/pd2 edge cases. Co-authored-by: Cursor <cursoragent@cursor.com>

…rformance for buffering approch Co-authored-by: Cursor <cursoragent@cursor.com>

…lex_attention # Conflicts: # scripts/run_benchmark.sh

…n MFMA operand When BufferCopy128b loads a 128-bit value (vector<8xbf16>) but the 16x16x16 MFMA intrinsic expects a 64-bit operand (vector<4xi16>), the emitAtomCallSSA bitcast was invalid (128 bits != 64 bits), causing an LLVM assertion: "Invalid cast!". Replace the direct bitcast with a width-aware matchWidth helper that detects when the source is wider than the target, bitcasts to the target element type at full width, then extracts the low slice via vector.extract_strided_slice. Same-width bitcasts are unchanged. This unblocks mma_k=16 for attention kernels, enabling block_n=64 with matching QK C / PV A fragment sizes for register P-bridge and flash-style permlane32_swap reductions. Co-Authored-By: Claude <noreply@anthropic.com>

Fix MFMA 16x16x16 bf16/f16 lowering crash when BufferCopy128b loads 128-bit values that feed 64-bit MFMA operands. The MLIR canonicalizer folds extract_strided_slice + bitcast chains back to the wider source, producing invalid width-changing bitcasts (e.g. i128 → vector<4xi16>) that LLVM rejects. Three layers of defense: 1. emitAtomCallSSA matchWidth (CDNA3/CDNA4 MmaAtom.cpp): handles the SSA path with vector extract_strided_slice when source is wider than the MFMA operand. 2. ConvertAtomCallToSSAForm narrowToMmaWidth: narrows register values to match the MMA atom's expected operand width at pass 06. 3. fly-fix-bitcast-width pass (new): runs before the canonicalizer, inserts llvm.freeze on narrowing extract_strided_slice results whose source traces back to a wider integer type, blocking the canonicalizer from folding the chain into an invalid bitcast. Verified correct for mma_k=32 block_n=32 (no regression), mma_k=16 block_n=32, and mma_k=16 block_n=64. Co-Authored-By: Claude <noreply@anthropic.com>

…attention

Major restructuring of the flex attention kernel from MFMA 16x16x32 to MFMA 32x32x16, matching flash attention's architecture. Key changes: - MFMA 32x32x16 bf16: each thread holds 16 M-values at one N-column - Column reduction softmax: per-thread max/sum with no cross-lane ds_swizzle ops in the main loop. One permlane32_swap on max to combine row-halves. The per-column exp scaling is numerically safe for normalized inputs (Q/K magnitudes ~0.1) and the O/l ratio cancels per-lane differences. - Register P bridge: P scores stay in registers for the PV GEMM (no LDS round-trip, no s_barrier for P). Works because mma_m==mma_n gives matching C→A fragment element count. - Double-buffered DMA: next tile prefetched during current compute. - Dual-wave stagger: num_groups>=2 enables phase-shifted overlap between memory and compute clusters. - Dynamic scf.for loop: eliminates register spilling from compile-time unrolling of large tile counts. - K LDS swizzle: SwizzleType.get(3,3,3) for bank-conflict reduction. - Q pre-scaled by 1/sqrt(D) to keep raw scores small for the column reduction path. - waves_per_eu=2 hint for CU packing. Verified correct (cos>0.999) at S=64..4096 with realistic input magnitudes. Performance: 507 TFLOPS at S=2048 with num_groups=8. Co-Authored-By: Claude <noreply@anthropic.com>

…FLOPS) Eliminates the LDS write + barrier + LDS read round-trip for the P (softmax output) bridge between QK and PV GEMMs, and adds K LDS swizzle for bank-conflict-free reads. Result: 286 → 428 TFLOPS. Key changes: - Swap QK GEMM to K=A, Q=B so C fragment M-rows = score indices, N-cols = query. This enables register-only C→B packing since C and B share lane→column mapping. - Pack P into MFMA B operand via cvt_pk_bf16_f32 (bf16) or f32→f16 truncation. No cross-lane movement — purely register-local. - Hand-roll PV MFMA: V loaded as A operand, P packed as B, raw mma_atom_call_ssa. - Remove host V pre-transpose (v.permute). V stays in BSHD format; the kernel tiles V into compact [block_n, 32] D-chunk sub-tiles in LDS during DMA. - O accumulator as 4 raw v16f32 (one per D-chunk) with M=D, N=query layout. Manual O store maps element e → D-position 8*(e//4)+e%4+4*(lane//32). - Softmax simplified: npair=1 always for 32x32 (each lane has 16 score values at 1 query column — exact per-row softmax via permlane32 only, no shuffle_xor). - K LDS Swizzle(3,3,3) with DMA global-address compensation via crd2idx through the swizzled layout (self-inverse XOR swizzle). Co-Authored-By: Claude <noreply@anthropic.com>

RichardChamberlain1 and others added 4 commits July 29, 2026 17:37

Merge branch 'main' into rchamber/flex_attention

a6c786d

RichardChamberlain1 marked this pull request as draft July 30, 2026 14:23

RichardChamberlain1 and others added 14 commits August 3, 2026 16:37

Replaced static loop which was causing register growth and hurting pe…

3cee3e2

…rformance for buffering approch Co-authored-by: Cursor <cursoragent@cursor.com>

Merge branch 'main' of https://github.com/ROCm/FlyDSL into rchamber/f…

af0b7bd

…lex_attention # Conflicts: # scripts/run_benchmark.sh

Merge branch 'rchamber/add_16x16x16_MFMA_support' into rchamber/flex_…

214222d

…attention

Navigation Menu

Uh oh!

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

Add flex_attention (score_mod / mask_mod) on the generic flash-attention kernel#931
RichardChamberlain1 wants to merge 18 commits into
mainfrom
rchamber/flex_attention

RichardChamberlain1 commented Jul 30, 2026

Uh oh!

RichardChamberlain1 commented Jul 30, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Uh oh!

Conversation

RichardChamberlain1 commented Jul 30, 2026

Summary

Test plan

Uh oh!

RichardChamberlain1 commented Jul 30, 2026

FlyDSL flex_attention Performance Report

mi300 (SMC-SC-DI10-33.dh144.dcgpu)

Shape S=2048 (B=2, H=32, Hkv=32, D=128)

Shape S=4096 (B=2, H=32, Hkv=32, D=128)

Shape S=8192 (B=2, H=32, Hkv=32, D=128)

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant