Rung 3 — Triton FlashAttention-2¶

Aligned phase: Phase 3 (attention + fusion) Ship by: end of M6 Effort: ~4 weeks, ~40 hrs Signal: mid-high. Writing a working FA2 kernel in Triton is a real credential.

The artifact spec¶

Implement FlashAttention-2 forward pass in Triton, benchmark against torch.nn.functional.scaled_dot_product_attention (SDPA) and the official flash-attn package. Then write it up.

Minimum scope:

Forward pass only. Backward is a stretch goal. Inference eng cares about forward.
Causal + non-causal. Both masks.
FP16 and BF16. Test both.
Head dims 64 and 128. These cover Llama, Qwen, Mistral.
Batched. Batch=1, 8, 32. Seq=512, 2048, 8192, 16384.

Stretch goals (do at least one):

Split-KV variant for the decode case (Flash-Decoding).
MQA/GQA-aware version (broadcast K/V across query heads).
Comparison vs FA3 on H100 rented for one hour.

Benchmark table (the money shot)¶

impl	shape (B,H,N,D)	dtype	causal	ms	TFLOPs	% of FA2 official
torch SDPA	8,32,2048,128	bf16	yes	…	…	…
flash-attn v2	8,32,2048,128	bf16	yes	…	…	100%
ours (Triton)	8,32,2048,128	bf16	yes	…	…	70-90%

Target: 70-90% of official FA2 on your card. The official FA2 is hand-tuned CUDA; getting into that range with Triton is the credible outcome.

Repo structure¶

triton-fa2/
├── README.md
├── kernels/
│   ├── fa2_fwd.py         # main Triton kernel
│   ├── fa2_fwd_causal.py  # if you want a separate causal version
│   └── fa2_split_kv.py    # stretch
├── bench/
│   ├── bench.py
│   ├── shapes.py          # the sweep config
│   └── plot.py
├── tests/
│   ├── test_correctness.py  # vs SDPA to 1e-2 in fp16, 5e-3 in bf16
│   └── test_shapes.py
├── results/
│   ├── raw.csv
│   └── hero.png
└── notes/
    ├── derivation.md      # your online softmax + tiling walkthrough
    └── roofline.md        # arithmetic intensity analysis

Where to post¶

GitHub: github.com/<you>/triton-fa2. Pin.
Blog: “FlashAttention-2 in Triton: 82% of official FA2 in 300 lines.” Line count in title is a good hook.
r/MachineLearning: submit. Sunday morning ET.
r/LocalLLaMA: if you did the MQA/GQA variant, cross-post here — they care about serving.
GPU MODE Discord #triton and #flash-attention: share. Ask if anyone wants to review.
Twitter/X: thread with the derivation image + hero plot. Tag @tri_dao @thom_wolf.
LinkedIn: first LinkedIn post is acceptable here. Short: “Wrote FA2 in Triton, hit 82% of official. Repo: link. Blog: link.”

What signals it sends¶

“I understand online softmax and the FA tiling.”
“I can write Triton at a level higher than the intro tutorial.”
“I can profile and reason about achieved TFLOPs and memory bandwidth.”
“I have paid the FA tax” — every inference eng study asks about it.

Past examples to study¶

Official Triton FA2 tutorial — https://triton-lang.org/main/getting-started/tutorials/06-fused-attention.html. Start here. Understand every line, then rewrite from memory.
ThunderKittens FA2 — github.com/HazyResearch/ThunderKittens. 100 lines, 30% faster than official. Read the code, then blog b1.
xformers memory_efficient_attention — github.com/facebookresearch/xformers/tree/main/xformers/ops. Reference impl.
Umer Adil GPU MODE lecture 14 — practitioner’s guide to Triton on YouTube.
Tri Dao’s Flash Attention repo — github.com/Dao-AILab/flash-attention. Reference C++/CUTLASS impl.

Common mistakes¶

Skipping the numerical stability check. The -inf masking + exp(x - max) combo is where FP16 bugs hide. Compare to fp32 reference not fp16 reference.
Wrong block sizes. Tune BLOCK_M, BLOCK_N per SKU. What works on 3090 (BLOCK_M=128, BLOCK_N=32) is not optimal on 4090 (BLOCK_M=128, BLOCK_N=64). Use triton.autotune or a small sweep.
Not benchmarking at real shapes. People care about (B=1, H=32, N=8192) decode-like and (B=8, H=32, N=2048) prefill-like. Benchmark both.
Ignoring GQA. Modern models are GQA. If you don’t handle it your kernel is unusable for Llama-3 etc. At minimum stub the case.
No autotuner cache warmup in bench. First call includes JIT compile. Warm up, then measure.
Overclaiming. If you’re at 60% of official, say 60%. Do not round up. The community will notice.

Success criteria¶

Passes correctness against SDPA on all tested shapes
≥60% of official FA2 for at least one shape
Hero plot in blog
ncu profiling section: DRAM throughput, tensor-core util, achieved occupancy
Blog < 3000 words
Something upstream contributor (Tri, Horace, Mark) sees it (comment, retweet, or star)

If ≥5 of 6 check, you are ready for rung 4 or 5. Most people ship rung 5 (mini-engine) before rung 4 (quant bake-off) because rung 5 unlocks rung 6 (first PR).