Case study: FlashAttention

Attention is where tokens meet

• A transformer block maps [n × d] → [n × d]: one row per token
• LayerNorm, projections, MLP: each row on its own
• Attention is the only op where one token reads another
• Two matrix multiplies with a softmax in between:

\[ O \;=\; \mathrm{softmax}\!\left(\frac{QK^\top}{\sqrt{d_k}} + \text{mask}\right) V \]

• That chain, GEMM → softmax → GEMM, is what this case study is about

Queries, keys and values

• Three ordinary GEMMs per token row:
  • Query Q = X·Wq: what this token is looking for
  • Key K = X·Wk: what this token offers
  • Value V = X·Wv: what it hands over
• Multi-head: split d into h heads of d_k = d / h columns
  • a reshape to [h × n × dk], no arithmetic
  • Llama-3-8B: d = 4096, h = 32, d_k = 128
• From here on: one head, [n × dk] matrices

The same computation at real size

Llama-3-8B shapes, one layer, n = 8,192 tokens, FP16

• S and P are 64× larger than Q, K, V, O and grow as n²
• Two GEMMs: 4·h·n²·d_k ≈ 1.1 TFLOP per layer, on tensor cores
• Scale, mask, softmax: ≈ 1 FLOP per byte moved
• Is attention limited by FLOPs or by bytes?

FlashAttention: the idea

• What it is: exact attention in one GPU kernel that never stores S or P
• Why: done step by step, S and P (n × n per head) go out to GPU memory (HBM) and back
  • HBM is large but slow; the GPU waits on bytes, not arithmetic
• Tiling: split Q, K, V into blocks that fit in on-chip memory
• Fusion: QKᵀ, mask, softmax and ·V for a block in one kernel
• Obstacle: softmax needs the whole row → online softmax

Dao, Fu, Ermon, Rudra, Ré, “FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness”, NeurIPS 2022.

Problem: softmax needs the whole row

\[ O_i \;=\; \sum_j p_{ij}\, v_j, \qquad p_{ij} = \frac{e^{\,s_{ij} - m_i}}{\ell_i}, \qquad m_i = \max_j s_{ij}, \qquad \ell_i = \sum_j e^{\,s_{ij} - m_i} \]

• A tile brings in only some of row i's keys
• Two things are unknown until the last tile:
  • the row max \(m_i\): a later tile may hold a bigger score
  • the row sum \(\ell_i\): so we cannot divide yet
• Standard softmax therefore needs the full row of S in memory

Online softmax: divide once, at the end

\(\ell_i\) is the same for every term of the row, so it can be pulled out of the sum:

\[ O_i \;=\; \sum_j \frac{e^{\,s_{ij} - m_i}}{\ell_i}\, v_j \;=\; \frac{1}{\ell_i} \underbrace{\sum_j e^{\,s_{ij} - m_i}\, v_j}_{o_i} \]

Keep running \(m, \ell, o\) on chip, updated per tile; nothing is divided:

\[ m \leftarrow \max(m, \max_{j \in \text{tile}} s_j), \quad \alpha = e^{\,m_{\text{old}} - m}, \quad \ell \leftarrow \alpha\,\ell + \sum_{j\in\text{tile}} e^{\,s_j - m}, \quad o \leftarrow \alpha\, o + \sum_{j\in\text{tile}} e^{\,s_j - m}\, v_j \]

• \(\alpha\) fixes terms made with an old max: \(e^{s - m_{\text{old}}}\, e^{m_{\text{old}} - m_{\text{new}}} = e^{s - m_{\text{new}}}\)
• After the last tile, \(\ell\) is the full row sum: \(O = o / \ell\), exactly

Online normalizer: Milakov & Gimelshein, “Online normalizer calculation for softmax”, 2018.

Takeaways

• Attention is limited by bytes, not FLOPs: S and P dominate
• Fusing across a reduction needs an algebraic rewrite
• Same maths, same answer: far faster, and memory grows as n, not n²
• Compilers must model the memory hierarchy, not just arithmetic