DL0056 FlashAttention

Written by

Explain FlashAttention. Why can it compute exact attention faster while using less memory?

Answer

FlashAttention is an exact, IO-aware implementation of scaled dot-product attention. Instead of materializing the full $N\times N$ score and probability matrices in high-bandwidth memory (HBM), it loads blocks of $Q$ , $K$ , and $V$ into fast on-chip SRAM, computes attention block by block, and maintains online softmax statistics. Tiling reduces expensive HBM reads and writes, while recomputation during the backward pass can be cheaper than storing large intermediates. The mathematical result matches standard attention up to normal floating-point differences; the speedup comes from changing the execution schedule, not from approximating attention.

Comparison of standard attention and FlashAttention data movement through HBM and SRAM.

Figure 1: IO comparison showing why avoiding N×N intermediate writes makes FlashAttention faster and more memory efficient.

(1) IO Awareness: The kernel is organized around the GPU memory hierarchy so most score computation and softmax updates happen in SRAM.
(2) Online Softmax: A running row maximum and normalization sum allow each $K,V$ tile to update the output without retaining previous score blocks.
(3) Exact Result: All query-key interactions are evaluated; causal masking, dropout, and backward gradients are fused into specialized kernels rather than approximated.

Mathematical Formulation:
$m_i^{(t)}=\max\!\left(m_i^{(t-1)},\max_j S_{ij}^{(t)}\right)$
$\ell_i^{(t)}=e^{m_i^{(t-1)}-m_i^{(t)}}\ell_i^{(t-1)}+\sum_j e^{S_{ij}^{(t)}-m_i^{(t)}}$

Where:

$t$ indexes key/value tiles, $i$ indexes a query row, and $j$ indexes keys inside tile $t$ .
$S_{ij}^{(t)}=Q_iK_j^T/\sqrt d$ is the scaled score between query row $Q_i$ and key row $K_j$ , with head width $d$ .
$m_i^{(t)}$ is the running row maximum after tile $t$ , initialized with $m_i^{(0)}=-\infty$ .
$\ell_i^{(t)}$ is the running softmax denominator, initialized with $\ell_i^{(0)}=0$ ; the exponential factors rescale earlier partial sums when the maximum changes.