Flash Attention

Dao et al., "FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness", NeurIPS 2022

Learning Objectives

After studying this section you will be able to:

1. Explain why standard attention is bottlenecked by memory bandwidth, not arithmetic, and quantify the gap between HBM and SRAM on modern GPUs.
2. Derive the online softmax identities that allow exact softmax computation in a single pass without storing the full $N \times N$ score matrix.
3. Walk through a complete tiled computation of our running example "the cat sat on the mat" and verify that the output is identical to Chapter 1.
4. Explain how Flash Attention achieves $O(N)$ memory and $3\text{--}8\times$ wall-clock speedup while computing mathematically identical results.
5. Describe how Flash Attention integrates with multi-head attention, KV-cache, causal masking, and modern systems like PyTorch's scaled_dot_product_attention.

The Real Problem

By Chapter 12, we have seen attention variants that change the mathematics — masking, sparsity, linear kernels. Flash Attention is fundamentally different: it does NOT change what is computed. The output is mathematically identical to standard scaled dot-product attention from Chapter 1 — identical up to floating-point precision, since the reordered arithmetic may differ by machine epsilon (~$10^{-16}$). What it changes is how the computation is executed on hardware.

The problem it solves is simple to state: standard attention is slow because it moves too much data, not because it does too much math.

The GPU Memory Hierarchy

Modern GPUs have a two-level memory hierarchy. Understanding it is essential to understanding why Flash Attention works.

The key ratio: SRAM is roughly 10\u00d7 faster than HBM but ~4,000\u00d7 smaller (80 GB vs ~20 MB on A100). Every byte read from or written to HBM costs approximately 40\u00d7 more time than an arithmetic operation. This gap is called the memory wall.

The Memory Wall in Standard Attention

Standard attention (Chapter 1) executes these steps, each requiring a full round-trip to HBM:

1. Read $Q$ and $K$ from HBM → compute $S = QK^T / \sqrt{d_k}$ → write $S$ ($N \times N$) back to HBM
2. Read $S$ from HBM → compute $P = \text{softmax}(S)$ → write $P$ ($N \times N$) back to HBM
3. Read $P$ and $V$ from HBM → compute $O = PV$ → write $O$ to HBM

Total HBM access: $O(N^2 + Nd)$ reads/writes. For a sequence of $N = 8192$ tokens with $d = 128$, the score matrix alone is $8192^2 = 67{,}108{,}864$ elements — 256 MB in float32. This matrix is written, read, softmaxed, written again, read again. The GPU spends most of its time waiting for data, not computing.

The core insight: Standard attention is memory-bound, not compute-bound. The GPU has enough arithmetic throughput to compute attention much faster than it can read and write the intermediate matrices. Flash Attention eliminates these intermediates entirely.

Exact vs Approximate: A Key Distinction

Chapter 12 explored methods like Linformer and Performer that reduce attention's complexity by changing the mathematics — using low-rank projections or kernel approximations to avoid computing the full $N \times N$ score matrix. These methods produce different outputs from standard attention, trading accuracy for speed.

Flash Attention takes a fundamentally different approach: it produces mathematically identical output to standard attention (identical up to floating-point rounding, typically differing by less than $10^{-16}$). The innovation is entirely in how the computation is scheduled on hardware. This distinction is critical because Flash Attention can be used as a drop-in replacement with no accuracy trade-off — every model that uses standard attention can switch to Flash Attention and get the exact same training dynamics and outputs, just faster.

The Story Behind Flash Attention

Tri Dao's IO-Aware Insight

In 2022, Tri Dao (Stanford, advised by Christopher Ré) published "FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness." The paper's key contribution was not a new attention formula but a new way to think about the problem: treat attention as an IO problem, not a compute problem.

Dao observed that the standard attention algorithm was designed for mathematical clarity, not hardware efficiency. It creates a clean pipeline (scores → softmax → output) but this pipeline requires materializing the $N \times N$ score matrix in HBM. His insight: if we can fuse all three steps into a single kernel that processes data in tiles small enough to fit in SRAM, we eliminate all intermediate HBM writes.

The challenge: softmax is a non-decomposable operation. Unlike addition or max, the softmax of a row depends on all elements in that row (because of the denominator $\sum_j e^{s_j}$). You cannot compute the softmax of the first half of a row, then the second half, and combine them — at least not with the naive formulation.

The Online Softmax Trick

The breakthrough came from adapting the online softmax algorithm (Milakov & Gimelshein, 2018). The idea: maintain running statistics — the current maximum $m$, the current normalizer $l$, and the current unnormalized output $O$ — and update them incrementally as each tile of keys is processed. When the maximum changes (because a new tile has a larger score), all previously accumulated values are rescaled to be consistent with the new maximum.

This rescaling is the mathematical trick that makes tiling work. It preserves exact numerical equivalence while allowing computation to proceed one small tile at a time, entirely within SRAM.

Mathematical Definition

The Standard Attention Formula

Flash Attention computes exactly the same function as Chapter 1:

$\text{Output} = \text{softmax}\!\left(\frac{QK^T}{\sqrt{d_k}}\right) \cdot V$

where $Q, K \in \mathbb{R}^{N \times d_k}$ and $V \in \mathbb{R}^{N \times d_v}$. The difference is entirely in the execution strategy.

Online Softmax Equations

For each query row $i$, we process the keys in tiles of size $T$. Let the current K-tile cover columns $j_0$ to $j_1$. We maintain three running statistics:

• $m$: the running maximum score seen so far (scalar per query)
• $l$: the running softmax denominator (scalar per query)
• $O$: the running unnormalized output (vector per query)

For each new tile, compute the tile scores $s_t = Q_i \cdot K_t^T / \sqrt{d_k}$, then update:

$m_{\text{new}} = \max(m_{\text{old}},\; \max_j(s_{t,j}))$

$\alpha = e^{m_{\text{old}} - m_{\text{new}}}$(rescaling factor)

$l_{\text{new}} = l_{\text{old}} \cdot \alpha + \sum_j e^{s_{t,j} - m_{\text{new}}}$

$O_{\text{new}} = O_{\text{old}} \cdot \alpha + \sum_j e^{s_{t,j} - m_{\text{new}}} \cdot V_j$

After all tiles are processed, the final output is: $\text{output}_i = O_{\text{final}} / l_{\text{final}}$

Why Online Softmax Works

The key identity is that rescaling preserves the ratio. Suppose after tile 1 we have computed $l_1 = \sum_j e^{s_j - m_1}$. When tile 2 arrives with a new maximum $m_2 > m_1$, every previously computed exponential $e^{s_j - m_1}$ needs to become $e^{s_j - m_2}$. The multiplicative correction is:

$e^{s_j - m_2} = e^{s_j - m_1} \cdot e^{m_1 - m_2}$

This is a single scalar multiplication that corrects all previously accumulated values at once. The factor $\alpha = e^{m_1 - m_2}$ is always $\leq 1$ (since $m_2 \geq m_1$), so it never causes overflow.

Why this matters: Without online softmax, computing softmax requires two passes over the data — one to find the max, one to compute exp and normalize. Online softmax does it in a single pass by correcting previous values when the max changes. This is what allows the entire attention computation to be fused into one kernel.

Step-by-Step Worked Example

Let's trace Flash Attention for query row 0 ("The") using our running example. With $N = 5$ and $T = 2$, we process 3 K-tiles: $K[0{:}2]$, $K[2{:}4]$, $K[4{:}5]$. The running statistics start at: $m = -\infty$, $l = 0$, $O = [0,0,0,0]$.

K-Tile 1 of 3: $Q[0{:}2] \times K[0{:}2]$

Score "The" against keys ["The", "cat"]:

$s[0,0] = Q_{\text{The}} \cdot K_{\text{The}} / 2 = (1 \cdot 0 + 0 \cdot 1 + 1 \cdot 0 + 0 \cdot 1) / 2 = 0.000$

$s[0,1] = Q_{\text{The}} \cdot K_{\text{cat}} / 2 = (1 \cdot 1 + 0 \cdot 0 + 1 \cdot 1 + 0 \cdot 0) / 2 = 1.000$

$m_{\text{new}} = \max(-\infty,\; \max(0.000, 1.000)) = 1.000$

$\text{exp\_scores} = [e^{0.000 - 1.000},\; e^{1.000 - 1.000}] = [0.3679,\; 1.0000]$

Since $m_{\text{old}} = -\infty$, $\alpha = 0$ (first tile, no previous values to rescale):

$l = 0 \cdot 0 + 0.3679 + 1.0000 = 1.3679$

$O = [0,0,0,0] \cdot 0 + 0.3679 \cdot V_{\text{The}} + 1.0000 \cdot V_{\text{cat}} = [0.3679,\; 1.0000,\; 0.0000,\; 0.0000]$

K-Tile 2 of 3: $Q[0{:}2] \times K[2{:}4]$

Score "The" against keys ["sat", "on"]:

$s[0,2] = 0.500$ (The vs sat), $s[0,3] = 0.500$ (The vs on)

$m_{\text{new}} = \max(1.000,\; \max(0.500, 0.500)) = 1.000$ (unchanged!)

Since the max did not change, $\alpha = e^{1.000 - 1.000} = 1.000$. No rescaling needed:

$\text{exp\_scores} = [e^{0.500 - 1.000},\; e^{0.500 - 1.000}] = [0.6065,\; 0.6065]$

$l = 1.3679 \times 1.000 + 0.6065 + 0.6065 = 2.5809$

$O = [0.3679,\; 1.0000,\; 0.0000,\; 0.0000] \times 1.000 + 0.6065 \cdot V_{\text{sat}} + 0.6065 \cdot V_{\text{on}}$

$= [0.3679,\; 1.0000,\; 0.6065,\; 0.6065]$

K-Tile 3 of 3: $Q[0{:}2] \times K[4{:}5]$

Score "The" against key ["mat"]:

$s[0,4] = 0.750$, $m_{\text{new}} = \max(1.000, 0.750) = 1.000$ (still unchanged)

$\alpha = 1.000$, $\text{exp} = e^{0.750 - 1.000} = 0.7788$

$l = 2.5809 + 0.7788 = 3.3597$

$O = [0.3679,\; 1.0000,\; 0.6065,\; 0.6065] + 0.7788 \times [0.5,\; 0.5,\; 0.5,\; 0.5]$

$= [0.7573,\; 1.3894,\; 0.9959,\; 0.9959]$

Final Normalisation

$O_{\text{The}} = [0.7573,\; 1.3894,\; 0.9959,\; 0.9959] \;/\; 3.3597$

$= [0.2254,\; 0.4135,\; 0.2964,\; 0.2964]$ — identical to standard attention (Chapter 1).

Interactive: Tiled Computation Explorer

Step through all 9 tiles of Flash Attention. Watch how the running statistics ($m$, $l$, $O$) evolve and how the current tile block moves across the score matrix. The highlighted cells show what is currently in SRAM — only a small 2×2 block at a time.

Full Attention Weights and Output

Flash Attention produces the same weights and output as standard attention. The weight matrix is never explicitly computed in Flash Attention — it exists only implicitly through the tiled accumulation — but mathematically the result is identical.

Attention Weight Matrix — Standard = Flash ($5 \times 5$)

Output Matrix — Standard = Flash ($5 \times 4$)

Maximum difference from standard attention: $5.55 \times 10^{-17}$ — floating-point machine epsilon. The outputs are mathematically identical.

Backward Pass and Recomputation

The forward pass discussion above is only half the story. For training, we need gradients. Standard attention stores the full $N \times N$ attention weight matrix $P$ during the forward pass so it is available for backpropagation. Flash Attention cannot store this matrix — that is the whole point — so how does the backward pass work?

The answer is recomputation (Algorithm 2 in the paper). During the forward pass, Flash Attention stores only three things per query row: the output $O$, the row-wise maximum $m$, and the row-wise normalizer $l$. During the backward pass, when gradients with respect to the attention weights are needed, the algorithm recomputes the attention matrix block-by-block from $Q$, $K$, $V$ (which are already stored for backpropagation), using the saved $m$ and $l$ to reconstruct the correct softmax values for each tile.

This recomputation strategy is arguably the paper's most important practical contribution. It is what makes Flash Attention viable for training, not just inference. The extra arithmetic cost of recomputation is modest (one additional pass over the data) and is far outweighed by the massive reduction in HBM access. In practice, the reduced memory traffic more than compensates for the extra FLOPs, resulting in faster training overall.

Why this matters: Without the recomputation trick, Flash Attention would save memory during inference but still need the full $N \times N$ matrix for training. The backward pass algorithm is what makes Flash Attention a complete training solution, not just an inference optimization.

Causal Masking Integration

For decoder models (GPT, Llama, etc.), causal masking prevents tokens from attending to future positions (Chapter 3). The paper shows how this integrates naturally with tiling:

• Skip entire tiles: When all K-positions in a tile are strictly after all Q-positions (the tile is entirely above the diagonal in the attention matrix), the entire tile can be skipped — no computation, no memory access.
• Partial masking: For tiles that straddle the diagonal boundary, masking is applied within the tile by setting future-position scores to $-\infty$ before the softmax computation. The online softmax handles these masked values naturally.

This tile-level skipping provides an additional speedup for causal attention beyond the IO benefits. Roughly half the tiles can be skipped entirely (those in the upper-triangular region of the block grid), reducing effective computation by nearly $2\times$ for causal models compared to bidirectional attention.

Dropout in Tiled Computation

Standard attention applies dropout to the attention weight matrix $P$, which requires storing the full $N \times N$ dropout mask. Flash Attention avoids this by storing only the pseudorandom number generator seed for each tile block. During both forward and backward passes, the dropout mask is regenerated on-the-fly from the seed, keeping memory usage at $O(1)$ per tile rather than $O(N^2)$ for the full mask.

Block-Sparse Flash Attention

Section 5 of the paper introduces block-sparse Flash Attention, which combines tiling with predetermined sparsity patterns. Given a block-sparsity mask that specifies which tiles to compute and which to skip, the algorithm simply avoids loading and computing masked-out tiles. The IO complexity improves to:

$O\!\left(\frac{N^2 d^2}{M} \cdot s\right)$where $s$ is the fraction of non-zero blocks

With 50% sparsity, this provides roughly $2\times$ additional speedup over dense Flash Attention while maintaining exact attention for the non-masked positions. This bridges the gap between exact and sparse attention methods — you get the IO efficiency of Flash Attention combined with the computational savings of sparsity.

Applications Across Domains

Long-Context NLP

Flash Attention enabled the jump from 2K to 100K+ context windows. Modern large language models — including GPT-4, Claude, Gemini, and Llama — are widely reported to use Flash Attention variants, as the technique has become a standard building block for long-context transformers. Without it, the $O(N^2)$ memory of standard attention makes sequences beyond ~8K tokens impractical on 80 GB GPUs.

Vision Transformers

High-resolution images tokenized into patches produce very long sequences (a 1024×1024 image with 16×16 patches = 4096 tokens). Flash Attention makes ViTs practical at these resolutions without the memory overhead of storing $4096^2$ attention matrices per head per layer.

Scientific Sequence Modeling

Protein sequences (up to ~30K amino acids) and genomic sequences (millions of base pairs) require long-range attention. AlphaFold 2 (2021) used its own memory-efficient chunking strategy before Flash Attention existed, but subsequent protein structure models (OpenFold, ESMFold, and AlphaFold 3) benefit directly from Flash Attention-style tiling for their pair representation attention layers.

Code Generation

Codex, StarCoder, and similar models process entire files (10K+ tokens). Flash Attention allows these models to attend over the full file context during both training and inference, enabling cross-function reasoning that would be impossible with shorter context windows.

Connection to Modern Systems

Flash Attention 2

Dao (2023) released Flash Attention 2 with two major improvements: (1) swapping the inner/outer loop order so the outer loop iterates over Q tiles (not K tiles), reducing non-matmul FLOPs by ~50%, and (2) better work partitioning across GPU thread blocks and warps. Flash Attention 2 achieves 50–73% of the theoretical peak FLOP/s on A100 GPUs, versus 25–40% for Flash Attention 1.

Flash Attention 3

Flash Attention 3 (Shah et al., 2024) targets Hopper architecture GPUs (H100). It exploits asynchronous data movement with the Tensor Memory Accelerator (TMA), warp specialization for overlapping computation and data loading, and FP8 low-precision support. It achieves 75% of the H100's theoretical peak throughput.

KV-Cache Integration

During autoregressive generation (Chapters 3, 5, 6), each new token only needs to compute attention scores against all previous keys/values. Flash Attention tiles work naturally with the KV-cache: the K and V tiles are simply loaded from the cache rather than recomputed. The online softmax approach handles the growing sequence length without ever materializing the full attention matrix.

PyTorch scaled_dot_product_attention

Since PyTorch 2.0, torch.nn.functional.scaled_dot_product_attention() automatically dispatches to Flash Attention when available. It selects the optimal backend (Flash Attention, Memory-Efficient Attention, or math fallback) based on input shapes, device, and dtype. In production, you never call Flash Attention directly — PyTorch does it for you.

Complexity Analysis

The paradox: Flash Attention has the same arithmetic complexity but is much faster. The speedup comes entirely from reducing HBM access. The IO complexity improves from $O(N^2 + Nd)$ to $O(N^2 d^2 / M)$. With typical SRAM sizes, this is a large constant-factor improvement.

Memory savings are even more dramatic. Standard attention stores the $N \times N$ score matrix, which at $N = 8192$ is 256 MB per head per layer. Flash Attention stores only the running statistics ($O(N)$ per head) plus the current tile ($O(T^2)$ in SRAM). This reduction enables training with much longer sequences on the same hardware.

Formal IO Model and Lower Bound

The paper defines a formal IO model of computation (Definition 1): given SRAM of size $M$ words and HBM of unbounded size, the cost of an algorithm is measured by the number of HBM accesses (reads and writes). Under this model, Flash Attention achieves IO complexity $O(N^2 d^2 / M)$, improving over standard attention's $O(N^2 d + Nd)$ HBM accesses.

More importantly, Proposition 3 proves this is asymptotically optimal: no exact attention algorithm can achieve fewer than $\Omega(N^2 d^2 / M)$ HBM accesses in the worst case. This means Flash Attention is not just faster — it is provably the best possible IO-aware exact attention algorithm (up to constant factors).

Benchmark Results from the Paper

The paper reports concrete speedups on standard benchmarks:

• BERT-large training: 15% end-to-end speedup over the MLPerf 1.1 speed record.
• GPT-2 training: Up to 3\u00d7 speedup compared to HuggingFace and Megatron-LM implementations.
• Long Range Arena: Flash Attention enables training on sequence lengths up to 16K on a single GPU, achieving new state-of-the-art on the Path-X task (sequence length 16,384) — a task previously unsolved by standard Transformers due to memory constraints.

These results demonstrate that the theoretical IO improvements translate directly to practical speedups. The gains are largest for longer sequences, where the $O(N^2)$ memory overhead of standard attention dominates runtime.

Python Implementation

A complete class-based implementation. The forward() method implements Flash Attention with online softmax tiling. The standard_attention() method provides the baseline for verification. Click any line to see the exact values flowing through the computation.

       d0   d1   d2   d3
The  1.0  0.0  1.0  0.0
cat  0.0  2.0  0.0  1.0
sat  1.0  1.0  1.0  0.0
on   0.0  0.0  1.0  1.0
mat  1.0  0.0  0.0  1.0

       d0   d1   d2   d3
The  0.0  1.0  0.0  1.0
cat  1.0  0.0  1.0  0.0
sat  1.0  1.0  0.0  0.0
on   0.0  0.0  1.0  1.0
mat  1.0  0.0  0.5  0.5

       d0   d1   d2   d3
The  1.0  0.0  0.0  0.0
cat  0.0  1.0  0.0  0.0
sat  0.0  0.0  1.0  0.0
on   0.0  0.0  0.0  1.0
mat  0.5  0.5  0.5  0.5

       The     cat     sat      on     mat
The  0.000   1.000   0.500   0.500   0.750
cat  1.500   0.000   1.000   0.500   0.250
sat  0.500   1.000   1.000   0.500   0.750
on   0.500   0.500   0.000   1.000   0.500
mat  0.500   0.500   0.500   0.500   0.750

[[1.000], [1.500], [1.000], [1.000], [0.750]]

       The     cat     sat      on     mat
The  0.3679  1.0000  0.6065  0.6065  0.7788
cat  1.0000  0.2231  0.6065  0.3679  0.2865
sat  0.6065  1.0000  1.0000  0.6065  0.7788
on   0.6065  0.6065  0.3679  1.0000  0.6065
mat  0.7788  0.7788  0.7788  0.7788  1.0000

       The     cat     sat      on     mat
The  0.1095  0.2976  0.1805  0.1805  0.2318
cat  0.4026  0.0898  0.2442  0.1481  0.1153
sat  0.1519  0.2505  0.2505  0.1519  0.1951
on   0.1903  0.1903  0.1154  0.3137  0.1903
mat  0.1892  0.1892  0.1892  0.1892  0.2430

       d0      d1      d2      d3
The  0.2254  0.4135  0.2964  0.2964
cat  0.4602  0.1475  0.3018  0.2058
sat  0.2495  0.3481  0.3481  0.2495
on   0.2854  0.2854  0.2106  0.4089
mat  0.3108  0.3108  0.3108  0.3108

       d0   d1   d2   d3
The  1.0  0.0  1.0  0.0
cat  0.0  2.0  0.0  1.0
sat  1.0  1.0  1.0  0.0
on   0.0  0.0  1.0  1.0
mat  1.0  0.0  0.0  1.0

       d0   d1   d2   d3
The  0.0  1.0  0.0  1.0
cat  1.0  0.0  1.0  0.0
sat  1.0  1.0  0.0  0.0
on   0.0  0.0  1.0  1.0
mat  1.0  0.0  0.5  0.5

       d0   d1   d2   d3
The  1.0  0.0  0.0  0.0
cat  0.0  1.0  0.0  0.0
sat  0.0  0.0  1.0  0.0
on   0.0  0.0  0.0  1.0
mat  0.5  0.5  0.5  0.5

       d0   d1   d2   d3
The  0.0  0.0  0.0  0.0
cat  0.0  0.0  0.0  0.0
sat  0.0  0.0  0.0  0.0
on   0.0  0.0  0.0  0.0
mat  0.0  0.0  0.0  0.0

    d0   d1   d2   d3
  0.0  1.0  0.0  1.0   The
  1.0  0.0  1.0  0.0   cat

    d0   d1   d2   d3
  1.0  0.0  0.0  0.0   The
  0.0  1.0  0.0  0.0   cat

    d0   d1   d2   d3
  1.0  0.0  1.0  0.0   The
  0.0  2.0  0.0  1.0   cat

       The    cat
The  0.000  1.000
cat  1.500  0.000

       sat     on
The  0.500  0.500
cat  1.000  0.500

       mat
The  0.750
cat  0.250

       The    cat
The  0.3679  1.0000
cat  1.0000  0.2231

       sat     on
The  0.6065  0.6065
cat  0.6065  0.3679

       d0      d1      d2      d3
The  0.2254  0.4135  0.2964  0.2964
cat  0.4602  0.1475  0.3018  0.2058
sat  0.2495  0.3481  0.3481  0.2495
on   0.2854  0.2854  0.2106  0.4089
mat  0.3108  0.3108  0.3108  0.3108

       d0   d1   d2   d3
The  1.0  0.0  1.0  0.0
cat  0.0  2.0  0.0  1.0
sat  1.0  1.0  1.0  0.0
on   0.0  0.0  1.0  1.0
mat  1.0  0.0  0.0  1.0

       d0   d1   d2   d3
The  0.0  1.0  0.0  1.0
cat  1.0  0.0  1.0  0.0
sat  1.0  1.0  0.0  0.0
on   0.0  0.0  1.0  1.0
mat  1.0  0.0  0.5  0.5

       d0   d1   d2   d3
The  1.0  0.0  0.0  0.0
cat  0.0  1.0  0.0  0.0
sat  0.0  0.0  1.0  0.0
on   0.0  0.0  0.0  1.0
mat  0.5  0.5  0.5  0.5

       The     cat     sat      on     mat
The  0.1095  0.2976  0.1805  0.1805  0.2318
cat  0.4026  0.0898  0.2442  0.1481  0.1153
sat  0.1519  0.2505  0.2505  0.1519  0.1951
on   0.1903  0.1903  0.1154  0.3137  0.1903
mat  0.1892  0.1892  0.1892  0.1892  0.2430

       d0      d1      d2      d3
The  0.2254  0.4135  0.2964  0.2964
cat  0.4602  0.1475  0.3018  0.2058
sat  0.2495  0.3481  0.3481  0.2495
on   0.2854  0.2854  0.2106  0.4089
mat  0.3108  0.3108  0.3108  0.3108

       d0      d1      d2      d3
The  0.2254  0.4135  0.2964  0.2964
cat  0.4602  0.1475  0.3018  0.2058
sat  0.2495  0.3481  0.3481  0.2495
on   0.2854  0.2854  0.2106  0.4089
mat  0.3108  0.3108  0.3108  0.3108

1import numpy as np
2import math
3
4class FlashAttention:
5    """
6    Flash Attention (Dao et al., 2022)
7    IO-aware exact attention that produces identical output to
8    standard scaled dot-product attention but minimizes HBM
9    reads/writes by processing Q, K, V in tiles, using online
10    softmax to avoid materializing the N*N score matrix.
11    """
12
13    def __init__(self, d_k: int, tile_size: int = 2):
14        self.d_k = d_k
15        self.tile_size = tile_size
16        self.scale = math.sqrt(d_k)
17
18    def standard_attention(self, Q, K, V):
19        """Standard O(N^2) attention for verification."""
20        scores = Q @ K.T / self.scale
21        max_s = np.max(scores, axis=-1, keepdims=True)
22        exp_s = np.exp(scores - max_s)
23        weights = exp_s / np.sum(exp_s, axis=-1, keepdims=True)
24        output = weights @ V
25        return weights, output
26
27    def forward(self, Q, K, V):
28        """
29        Flash Attention forward pass -- tiled online softmax.
30        Output is IDENTICAL to standard_attention().
31        """
32        N, D = Q.shape
33        D_v = V.shape[1]
34        T = self.tile_size
35
36        m_run = np.full(N, -np.inf)
37        l_run = np.zeros(N)
38        O_run = np.zeros((N, D_v))
39
40        for j0 in range(0, N, T):
41            j1 = min(j0 + T, N)
42            K_t = K[j0:j1]
43            V_t = V[j0:j1]
44
45            for i0 in range(0, N, T):
46                i1 = min(i0 + T, N)
47                Q_t = Q[i0:i1]
48
49                s_t = Q_t @ K_t.T / self.scale
50                m_new = np.maximum(m_run[i0:i1], s_t.max(axis=1))
51                e_t = np.exp(s_t - m_new[:, None])
52                rescale = np.exp(m_run[i0:i1] - m_new)
53
54                l_run[i0:i1] = l_run[i0:i1] * rescale + e_t.sum(axis=1)
55                O_run[i0:i1] = O_run[i0:i1] * rescale[:, None] + e_t @ V_t
56                m_run[i0:i1] = m_new
57
58        return O_run / l_run[:, None]
59
60    def explain(self, Q, K, V, tokens, query_idx=0):
61        """Print a tile-by-tile trace for one query token."""
62        N = Q.shape[0]
63        T = self.tile_size
64        t = tokens[query_idx]
65        i = query_idx
66
67        m_val = -np.inf
68        l_val = 0.0
69        o_val = np.zeros(V.shape[1])
70
71        print(f"\n=== Flash trace for '{t}' (row {i}, T={T}) ===")
72        print(f"Q[{i}] = {Q[i]}")
73
74        tile_num = 0
75        for j0 in range(0, N, T):
76            j1 = min(j0 + T, N)
77            tile_num += 1
78            K_t = K[j0:j1]
79            V_t = V[j0:j1]
80
81            scores = Q[i:i+1] @ K_t.T / self.scale
82            s = scores[0]
83            m_new = max(m_val, float(s.max()))
84
85            print(f"\n--- Tile {tile_num}: K[{j0}:{j1}] ---")
86            for idx in range(j0, j1):
87                print(f"  score({t},{tokens[idx]}) = {s[idx-j0]:.4f}")
88
89            e = np.exp(s - m_new)
90            rescale = np.exp(m_val - m_new) if np.isfinite(m_val) else 0.0
91
92            l_val = l_val * rescale + float(e.sum())
93            o_val = o_val * rescale + e @ V_t
94            m_val = m_new
95
96            print(f"  m={m_new:.4f}, l={l_val:.4f}")
97            print(f"  O_run = {np.round(o_val, 4)}")
98
99        final = o_val / l_val
100        print(f"\nFinal: O[{t}] = {np.round(final, 4)}")
101
102
103# -- Shared Example (same Q, K, V as every chapter) --
104tokens = ["The", "cat", "sat", "on", "mat"]
105
106Q = np.array([
107    [1.0, 0.0, 1.0, 0.0],   # The
108    [0.0, 2.0, 0.0, 1.0],   # cat
109    [1.0, 1.0, 1.0, 0.0],   # sat
110    [0.0, 0.0, 1.0, 1.0],   # on
111    [1.0, 0.0, 0.0, 1.0],   # mat
112])
113
114K = np.array([
115    [0.0, 1.0, 0.0, 1.0],   # The
116    [1.0, 0.0, 1.0, 0.0],   # cat
117    [1.0, 1.0, 0.0, 0.0],   # sat
118    [0.0, 0.0, 1.0, 1.0],   # on
119    [1.0, 0.0, 0.5, 0.5],   # mat
120])
121
122V = np.array([
123    [1.0, 0.0, 0.0, 0.0],   # The
124    [0.0, 1.0, 0.0, 0.0],   # cat
125    [0.0, 0.0, 1.0, 0.0],   # sat
126    [0.0, 0.0, 0.0, 1.0],   # on
127    [0.5, 0.5, 0.5, 0.5],   # mat
128])
129
130# -- Run --
131fa = FlashAttention(d_k=4, tile_size=2)
132
133std_weights, std_output = fa.standard_attention(Q, K, V)
134print("Standard Attention Weights:")
135print(np.round(std_weights, 4))
136
137flash_output = fa.forward(Q, K, V)
138print("\nFlash Attention Output:")
139print(np.round(flash_output, 4))
140
141diff = np.abs(flash_output - std_output).max()
142print(f"\nMax diff: {diff:.2e}")
143
144fa.explain(Q, K, V, tokens, query_idx=0)

PyTorch Implementation

The same algorithm in PyTorch. Note the differences from NumPy: .size() instead of .shape, .unsqueeze() instead of [:, None], dim= instead of axis=, and s_t.max(dim=1).values instead of s_t.max(axis=1). In production, use torch.nn.functional.scaled_dot_product_attention() which calls the optimized CUDA kernel.

tensor([[1.0, 0.0, 1.0, 0.0],
        [0.0, 2.0, 0.0, 1.0],
        [1.0, 1.0, 1.0, 0.0],
        [0.0, 0.0, 1.0, 1.0],
        [1.0, 0.0, 0.0, 1.0]])

tensor([[0.0, 1.0, 0.0, 1.0],
        [1.0, 0.0, 1.0, 0.0],
        [1.0, 1.0, 0.0, 0.0],
        [0.0, 0.0, 1.0, 1.0],
        [1.0, 0.0, 0.5, 0.5]])

tensor([[1.0, 0.0, 0.0, 0.0],
        [0.0, 1.0, 0.0, 0.0],
        [0.0, 0.0, 1.0, 0.0],
        [0.0, 0.0, 0.0, 1.0],
        [0.5, 0.5, 0.5, 0.5]])

tensor([[0.000, 1.000],
        [1.500, 0.000]])

tensor([[0.2254, 0.4135, 0.2964, 0.2964],
        [0.4602, 0.1475, 0.3018, 0.2058],
        [0.2495, 0.3481, 0.3481, 0.2495],
        [0.2854, 0.2854, 0.2106, 0.4089],
        [0.3108, 0.3108, 0.3108, 0.3108]])

1import torch
2import torch.nn as nn
3import math
4
5class FlashAttention(nn.Module):
6    """
7    Flash Attention in PyTorch (Dao et al., 2022)
8
9    Algorithmic demonstration of IO-aware tiling.
10    In production, use torch.nn.functional.scaled_dot_product_attention()
11    which calls the optimized CUDA kernel directly.
12    """
13
14    def __init__(self, d_model: int, tile_size: int = 2):
15        super().__init__()
16        self.d_model = d_model
17        self.tile_size = tile_size
18        self.scale = math.sqrt(d_model)
19
20    def forward(
21        self,
22        Q: torch.Tensor,
23        K: torch.Tensor,
24        V: torch.Tensor,
25    ) -> torch.Tensor:
26        """
27        Args:
28            Q: (N, d_model)
29            K: (N, d_model)
30            V: (N, d_v)
31        Returns:
32            output: (N, d_v) -- identical to standard attention
33        """
34        N = Q.size(0)
35        D_v = V.size(1)
36        T = self.tile_size
37
38        m_run = torch.full((N,), float("-inf"), device=Q.device)
39        l_run = torch.zeros(N, device=Q.device)
40        O_run = torch.zeros(N, D_v, device=Q.device)
41
42        for j0 in range(0, N, T):
43            j1 = min(j0 + T, N)
44            K_t = K[j0:j1]
45            V_t = V[j0:j1]
46
47            for i0 in range(0, N, T):
48                i1 = min(i0 + T, N)
49                Q_t = Q[i0:i1]
50
51                s_t = Q_t @ K_t.T / self.scale
52                m_new = torch.maximum(m_run[i0:i1], s_t.max(dim=1).values)
53                e_t = torch.exp(s_t - m_new.unsqueeze(1))
54                rescale = torch.exp(m_run[i0:i1] - m_new)
55
56                l_run[i0:i1] = l_run[i0:i1] * rescale + e_t.sum(dim=1)
57                O_run[i0:i1] = O_run[i0:i1] * rescale.unsqueeze(1) + e_t @ V_t
58                m_run[i0:i1] = m_new
59
60        return O_run / l_run.unsqueeze(1)
61
62
63# -- Shared example --
64tokens = ["The", "cat", "sat", "on", "mat"]
65
66Q = torch.tensor([
67    [1.0, 0.0, 1.0, 0.0],
68    [0.0, 2.0, 0.0, 1.0],
69    [1.0, 1.0, 1.0, 0.0],
70    [0.0, 0.0, 1.0, 1.0],
71    [1.0, 0.0, 0.0, 1.0],
72])
73K = torch.tensor([
74    [0.0, 1.0, 0.0, 1.0],
75    [1.0, 0.0, 1.0, 0.0],
76    [1.0, 1.0, 0.0, 0.0],
77    [0.0, 0.0, 1.0, 1.0],
78    [1.0, 0.0, 0.5, 0.5],
79])
80V = torch.tensor([
81    [1.0, 0.0, 0.0, 0.0],
82    [0.0, 1.0, 0.0, 0.0],
83    [0.0, 0.0, 1.0, 0.0],
84    [0.0, 0.0, 0.0, 1.0],
85    [0.5, 0.5, 0.5, 0.5],
86])
87
88fa = FlashAttention(d_model=4, tile_size=2)
89
90with torch.no_grad():
91    flash_out = fa(Q, K, V)
92    std_out = torch.softmax(Q @ K.T / fa.scale, dim=-1) @ V
93
94print("Flash Output:")
95print(flash_out.round(decimals=4))
96print(f"\nMax diff: {(flash_out - std_out).abs().max():.2e}")

Key Takeaways

1. Same math, different execution. Flash Attention produces mathematically identical output to standard scaled dot-product attention (differing only at floating-point precision). The $N \times N$ attention matrix is never stored — only $T \times T$ tile scores at a time.
2. IO-awareness is the insight. The bottleneck in standard attention is memory bandwidth (HBM reads/writes), not arithmetic. Flash Attention fuses all operations into a single kernel that keeps intermediates in SRAM.
3. Online softmax enables tiling. The running statistics $(m, l, O)$ with rescaling preserve exact softmax semantics while processing keys in arbitrarily small tiles.
4. Memory drops from $O(N^2)$ to $O(N)$. This enables training with 10–100× longer sequences on the same GPU hardware.
5. Wall-clock speedup of 3–8× despite doing the same number of arithmetic operations. The speedup comes from eliminating HBM round-trips.
6. Universal adoption. Flash Attention has become the de facto standard for transformer training and inference, integrated into PyTorch via scaled_dot_product_attention() and adopted across virtually all major LLM frameworks.

Exercises

1. Tile size impact: Modify the code to use tile_size=1 (process one key at a time). Verify the output is still identical. How does the number of tiles change? What happens to the number of rescaling operations?
2. Rescaling trace: Trace Flash Attention for query "cat" (row 1). At which tile does the running maximum first change from its initial tile-1 value? What is the rescaling factor $\alpha$?
3. Memory calculation: For $N = 16384$, $d = 128$, float16, calculate: (a) the size of the full $N \times N$ score matrix in GB, (b) the size of Flash Attention's running statistics per head, (c) the memory savings factor.
4. Causal masking: How would you integrate the causal mask from Chapter 3 into Flash Attention? Hint: for tiles where the Q-row index is less than the K-column range, the entire tile can be skipped. For boundary tiles, apply $-\infty$ masking within the tile scores.
5. Multi-head integration: Flash Attention operates per-head. If a model has $h = 32$ heads, how would you parallelize Flash Attention across heads? What is the peak SRAM usage per head?

References

1. Dao, T., Fu, D. Y., Ermon, S., Rudra, A., & Ré, C. (2022). FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. NeurIPS 2022.
2. Dao, T. (2023). FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning. arXiv:2307.08691.
3. Shah, J., Bikshandi, G., Zhang, Y., Thakkar, V., Ramani, P., & Dao, T. (2024). FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision. arXiv:2407.08608.
4. Milakov, M. & Gimelshein, N. (2018). Online normalizer calculation for softmax. arXiv:1805.02867.
5. Rabe, M. N. & Staats, C. (2022). Self-attention Does Not Need $O(n^2)$ Memory. arXiv:2112.05682.