flash-attn-triton

FlashAttention-2 · forward + backward · Triton

The score matrix
never lands.

Attention at long context is slow because of memory movement, not because of the maths. These kernels do the same arithmetic as ordinary attention — more of it, in the backward pass — and still win, purely by never writing the N×N score matrix to GPU main memory.

algorithm
FlashAttention-2, forward and backward
result
exact — not an approximation
precision
fp16, bf16
head dim
32, 64, 128
masking
causal and full, mask applied on the diagonal only
tile sizes
autotuned over BLOCK_M, BLOCK_N, warps, stages
target
NVIDIA SM 80+, CUDA PyTorch, Triton
licence
MIT
Ordinary attention 128 MiB of score matrix written to HBM, per head, at 8K context in fp16 — then read back, rewritten, read again.
This kernel 0 B of score matrix written to HBM. It keeps one 32 KiB tile on chip and throws it away. Only L, one float per row, survives.

Run it

Watch the inner loop

Every square is one BLOCK_M×BLOCK_N tile of the score matrix. Amber means the tile exists right now, on chip. Everything already computed is gone — it was never stored. Click any row to follow that program's softmax state.

Fig. 1 forward tile schedule · _attn_fwd
4

keys / values — inner loop, BLOCK_N at a time →

queries — one program per BLOCK_M

in SRAM now computed, discarded not reached skipped by causal mask on the diagonal — the only tiles carrying a mask program resident on an SM

The argument

Why skipping HBM is the whole speedup

Ordinary attention writes S = QKᵀ/√d to HBM, reads it back to run softmax, writes P, and reads it once more to multiply by V. That is Θ(N²) traffic across the slowest link in the machine, on top of the Θ(Nd) the inputs actually need.

The hierarchy is lopsided enough that this dominates everything. On an A100, SRAM moves data roughly ten times faster than HBM, and there is about 192 KB of it per SM. A 8192×8192 fp16 score matrix is 128 MiB. It cannot live on chip, so every element makes several round trips over the slow link.

Tiling keeps the working set small enough to stay on chip for the entire inner loop. Counting HBM accesses rather than FLOPs: standard attention is Θ(N² + Nd), the tiled version is Θ(N²d²/M) where M is SRAM size. With d = 64 and M ≈ 100 KB, d²/M is about 1/25.

Note — the counterintuitive part

The backward pass does more arithmetic than the ordinary version, because it recomputes P from the saved L instead of reading it back. It is still faster. At these shapes attention is memory-bound, so trading FLOPs for HBM traffic is a trade worth making.

Table 1GPU memory hierarchy · A100
LevelSizeBandwidth
SRAM (on chip)~192 KB / SM~19 TB/s
HBM (main memory)40–80 GB~1.5–2 TB/s
Table 2HBM accesses
ImplementationComplexity
Standard attentionΘ(N² + Nd)
FlashAttentionΘ(N²d²/M)

Work it out

What a shape actually costs

Nothing measured here — this is just the arithmetic of the tensor sizes, which is enough to show where the quadratic term takes over.

Fig. 2footprint calculator
score matrix, ordinary attention— B × H × N × N
everything this kernel holds— Q, K, V, O plus L — all linear in N
ratio— how many times larger the score matrix alone is
ordinary attention runs out at—  

The catch there isn't

Tiling softmax is exact algebra

Softmax is not local: the denominator needs every score in the row, and a tile only ever sees part of one. The fix is to carry running statistics and correct them.

Keep m, the largest score seen so far, and l, the running sum of exp(score − m). When a new tile arrives with a bigger max, rescale everything already accumulated by α = exp(m_old − m_new). That single factor corrects the sum and the output accumulator together.

Every step is an algebraic identity, so after the last tile acc / l is the same value ordinary attention produces — up to the floating point rounding the reference incurs too. That is why the test suite compares against fp32 PyTorch with dtype tolerances rather than approximation tolerances. The readout in the visualiser above runs this comparison live.

FlashAttention-2 adds three things on top, all implemented here: defer the normalisation to one divide after the loop, put queries in the outer loop so the accumulator stays in registers, and split the grid over the sequence dimension so long context with small batch still fills the SMs. Full derivation in DERIVATION.md.

Fig. 3the update, per tile
# running state, resident in registers
m_i = -inf        # biggest score so far
l_i = 0           # running sum of exp(s - m_i)
acc = 0           # output, not yet divided

# for each K/V tile:
m_ij  = max(m_i, max(qk))
p     = exp2(qk - m_ij)
alpha = exp2(m_i - m_ij)      # the correction
l_i   = l_i * alpha + sum(p)
acc   = acc * alpha + p @ v
m_i   = m_ij

# once, after the loop (this is the FA-2 part)
acc = acc / l_i
L   = m_i + log(l_i)      # N floats, not N x N

exp2 rather than exp: it is a single hardware instruction, so log2(e) is folded into the scale once, before the loop.

Scope

What is and isn't in here

Forward kernel, FA-2 style — normalise at the end, Q in the outer loop, grid split over the sequencedone
Backward kernels for dQ, dK, dV, rebuilding P from the saved Ldone
Causal masking, with the mask applied only on the diagonal tiledone
fp16 and bf16done
Sequence lengths that aren't a multiple of the tile sizedone
Autotuning over BLOCK_M / BLOCK_N / warps / stagesdone
autograd.Function plus a drop-in FlashSelfAttention moduledone
Tests against fp32 PyTorch across shapes, dtypes and maskingdone
Benchmark harness — speed, TFLOP/s, peak memory, OOM pointdone
MQA/GQA, dropout, ALiBi, sliding window, variable-length batchesnot yet
Hopper / FlashAttention-3 pathout of scope

Measurements

Benchmarks

This table is generated from docs/results.json, which the benchmark harness writes. It is empty until someone runs it on real hardware — there are no numbers here from a machine the code wasn't run on.

Table 3bench/bench_attention.py

reading results.json…

Source

Where things live

kernels/fwd.pyforward kernel, autotuned kernels/bwd.pydelta pre-pass, then dK/dV and dQ interface.pyautograd.Function, flash_attn, module reference.pyplain PyTorch attention, SDPA, FLOP count bench/speed and memory measurement tests/correctness and memory scaling docs/DERIVATION.mdthe maths, in full
using it
# needs an NVIDIA GPU (SM 80+), CUDA PyTorch, Triton.
# Triton does not run on Windows -- Linux, WSL2 or Colab.
pip install -r requirements.txt && pip install -e .

# then
from flash_attn_triton import flash_attn, FlashSelfAttention

out = flash_attn(q, k, v, causal=True)   # (B, H, N, D), fp16/bf16
attn = FlashSelfAttention(embed_dim=1024, num_heads=16, causal=True)

Sources

Papers this came from

The algorithm

  • FlashAttention: Fast and Memory-Efficient Exact Attention with IO-AwarenessDao, Fu, Ermon, Rudra, Ré · NeurIPS 2022 · arXiv:2205.14135
  • FlashAttention-2: Faster Attention with Better Parallelism and Work PartitioningDao · ICLR 2024 · arXiv:2307.08691
  • FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precisionShah et al. · NeurIPS 2024 · arXiv:2407.08608 · Hopper only, not implemented

Why tiling softmax stays exact

  • Online normalizer calculation for softmaxMilakov, Gimelshein · 2018 · arXiv:1805.02867
  • Self-attention Does Not Need O(n²) MemoryRabe, Staats · 2021 · arXiv:2112.05682

Why memory movement is the bottleneck

  • Data Movement Is All You Need: A Case Study on Optimizing TransformersIvanov et al. · MLSys 2021 · arXiv:2007.00072
  • Roofline: An Insightful Visual Performance Model for Multicore ArchitecturesWilliams, Waterman, Patterson · CACM 2009

The compiler

  • Triton: An Intermediate Language and Compiler for Tiled Neural Network ComputationsTillet, Kung, Cox · MAPL 2019