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

[RichardChamberlain1]

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

[RichardChamberlain1]

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*