FlashAttention

Every modern language model — GPT, LLaMA, Gemini — is built on the same core operation: self-attention. For each token in a sequence, attention computes a weighted blend of all the other tokens, letting every word "look at" every other word at once.

The catch is the cost. With $N$ tokens you need an $N \times N$ matrix of scores — quadratic in both time and memory. For short sequences that is fine. For sequences of thousands of tokens, that matrix alone can fill and overflow the fast on-chip memory of a GPU.

In 2022, Tri Dao and colleagues at Stanford published FlashAttention: an algorithm that computes the exact same result as standard attention while never materializing the full $N \times N$ matrix in fast memory. The trick is not an approximation. It is an IO-aware kernel — a piece of code that reasons about where data lives (slow HBM vs. fast SRAM) and moves it as little as possible.

The demo below visualises FlashAttention's core idea: tile the attention computation so that each tile fits in fast on-chip memory (SRAM), accumulate the result with an online softmax, then discard the tile before loading the next one.

<!-- {{c_html_intro}} -->
<div class="controls">
  <label>{{label_seq_len}} <strong id="nval">8</strong>
    <input type="range" id="nslider" min="4" max="16" step="4" value="8">
  </label>
  <label>{{label_tile_size}} <strong id="tval">4</strong>
    <input type="range" id="tslider" min="2" max="8" step="2" value="4">
  </label>
</div>
<div class="mode-btns">
  <button id="btn-standard" class="active" type="button">{{btn_standard}}</button>
  <button id="btn-flash" type="button">{{btn_flash}}</button>
</div>
<canvas id="canvas" width="480" height="280"></canvas>
<div id="stats" class="stats"></div>
<div class="action-btns">
  <button id="btn-run" type="button">{{btn_run}}</button>
  <button id="btn-reset" type="button" class="ghost">{{btn_reset}}</button>
</div>
<p class="caption" id="caption">{{caption_idle}}</p>

/* {{c_css_intro}} */
* { box-sizing: border-box; }
body { font-family: system-ui, sans-serif; color: #222; margin: 0; }
.controls { display: flex; gap: 1.4rem; flex-wrap: wrap; margin-bottom: .6rem; }
label { font-size: .85rem; color: #444; display: flex; align-items: center; gap: .5rem; }
input[type=range] { width: 90px; }
.mode-btns { display: flex; gap: .5rem; margin-bottom: .5rem; }
button { font: 600 13px system-ui; padding: .4rem .85rem; border: 1px solid #1d3557;
         background: #1d3557; color: #fff; border-radius: 7px; cursor: pointer; }
button.ghost { background: #fff; color: #1d3557; }
button.active { background: #e63946; border-color: #c0303c; }
canvas { display: block; border: 1px solid #cdd9e3; border-radius: 8px;
         background: #f8fafc; max-width: 100%; }
.stats { font-size: .85rem; margin: .4rem 0; min-height: 1.3em; color: #1d3557; font-weight: 600; }
.action-btns { display: flex; gap: .5rem; margin-top: .4rem; }
.caption { font-size: .82rem; color: #555; margin: .5rem 0 0; line-height: 1.5; }

// Code not found

Notice what never happens: the full $N \times N$ matrix of attention scores is never written to slow GPU memory (HBM). Standard attention writes it once and reads it back; FlashAttention rewrites each tile entirely in fast SRAM and throws it away. The final output is the same — bit-for-bit identical — because the online softmax accumulates running statistics that let it rescale on the fly.

Standard benchmarks measure FLOPs. But on modern GPUs, arithmetic is cheap — memory bandwidth is the real bottleneck.

• Standard attention reads the $Q$, $K$, $V$ matrices, writes the $N \times N$ score matrix $S$ to HBM, reads it back to apply softmax, writes the softmax result, reads it again to multiply by $V$. Total HBM accesses: $O(N^2)$.
• FlashAttention tiles the computation. For a tile block size $B_r \times B_c$, it keeps the current tile entirely in SRAM, accumulates partial softmax statistics, and writes only the output $O$ back to HBM. Total HBM accesses: $O(N^2 d / M)$, where $d$ is head dimension and $M$ is SRAM size — sublinear in $N$ relative to the naïve approach when $M$ is large.
• Recomputation on the backward pass: FlashAttention does not store the $N \times N$ attention matrix for backpropagation. Instead it recomputes it from the tiles during the backward pass. This trades extra FLOPs for drastically less memory — and FLOPs are cheap, HBM reads are not.
• The result is exact: unlike sparse or approximate attention methods, FlashAttention produces the numerically identical output. No information is lost; only the order of operations changes.

Tri Dao's 2022 paper showed 2–4× wall-clock speedup over standard PyTorch attention on A100 GPUs, and 5–20× less memory for long sequences. The 2023 FlashAttention-2 paper pushed it further with better parallelism, reaching up to 9× speedup. This is what allows modern LLMs to handle context windows of 100k+ tokens — it is the same reason transformer training became practical at scale.

FlashAttention is now a default ingredient in virtually every serious transformer implementation:

• LLM training: GPT-4, LLaMA 2/3, Mistral, and most large-scale models use FlashAttention to cut training costs significantly. Without it, training on sequences longer than a few thousand tokens would require impractically large GPU clusters.
• Long-context inference: 100k-token context windows (Anthropic Claude, GPT-4 Turbo) are only feasible because attention no longer blows up GPU memory at those lengths.
• Protein and molecular structure: AlphaFold-style models fold thousands of residues; the same IO-aware trick lets them run on hardware that would otherwise run out of memory.
• Video and multimodal models: treating video frames as tokens means sequences of tens of thousands; FlashAttention makes the attention pass tractable.
• Speculative decoding and fine-tuning: techniques like LoRA and QLoRA that fine-tune large models on consumer GPUs depend heavily on the memory savings FlashAttention provides.

The insight generalises: whenever you have a computation that produces a large intermediate matrix that is immediately consumed and discarded, tiling and recomputation can replace it with a series of passes over fast memory. The matrix multiplication literature has known this for decades; FlashAttention applied the idea to the attention operator specifically.

FlashAttention is a reminder that the bottleneck is rarely where you think it is. The naive analysis says attention is slow because of $O(N^2)$ multiplications. The real analysis says attention is slow because those multiplications force $O(N^2)$ round-trips between slow and fast memory — and that those round-trips can be slashed without losing a single bit of accuracy.

Tiling, online softmax, and recomputation on the backward pass are not exotic tricks. They are classical ideas from high-performance computing, applied precisely to the right operation at the right moment. The payoff is enormous: context windows have grown from a few thousand tokens to hundreds of thousands, and transformer models that once required dozens of GPUs now fit on a handful.

The deeper lesson: when you understand the memory hierarchy, you can redesign algorithms that look fixed — and sometimes the redesign changes everything.