Foundation Lab

FlashAttention: IO-Aware Attention

Enabled training with 100K+ context windows—without this, long context is impractical

Concept 68 of 100EfficiencyPhase 6

#68FlashAttnEfficiency

key equation\text{Memory: } O(N) \text{ vs } O(N^2)

Selected Foundation Object

Keep the equation fixed; move through the evidence.

Concept 68 of 100FlashAttnEfficiency / Phase 6: Modern efficiency & inference

Current question

GPU memory hierarchy matters: SRAM (fast, small) vs HBM (slow, large)

\text{Memory: } O(N) \text{ vs } O(N^2)

PredictionCommit before the demo.

Ask what should change when the equation is manipulated, then let the visualization test that expectation.

EvidenceCompare local witness and source.

Use the runnable panel, the key equation, and canonical papers as separate forms of evidence for the same object.

InvariantName what survives notation changes.

The useful learning product is the reusable mechanism you can carry into another model, paper, or engineering tradeoff.

Next moveContinue in the newer notebook.

This atlas page keeps the working demo; the domain notebook carries the fuller Intuition -> Math -> Code -> Demo sequence.

Enabled training with 100K+ context windows—without this, long context is impractical
Foundational for modern LLMs: used in LLaMA, GPT-4, Claude, etc.
Shows that algorithmic innovation can beat hardware by understanding memory hierarchy

What is still poorly explained in textbooks and papers:

GPU memory hierarchy matters: SRAM (fast, small) vs HBM (slow, large)
Materializing N×N attention wastes memory bandwidth—the bottleneck, not FLOPs
Online softmax: you can compute softmax in one pass by tracking running max and sum

If you want intuition first, start with the key equation and the visualization. Come back here for the full walkthrough.

Key Equation

\text{Memory: } O(N) \text{ vs } O(N^2)

Standard attention materializes $O(N^2)$ intermediate matrices:

\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d}}\right) V

FlashAttention fuses operations, keeping data in SRAM:

Memory: $O(N)$ instead of $O(N^2)$
Speed: 2-4× faster than standard attention

Dao et al.2022NeurIPS

Dao2023arXiv

Choose the next question to carry this object forward.