SystemsFlashAttention
Flash Attention Principles
PremiumThrough interactive visualizations, gain a deep understanding of Flash Attention's core techniques: the memory bottleneck, Online Softmax, and tiled matrix multiplication.
CookLLM DocsThrough interactive visualizations, gain a deep understanding of Flash Attention's core techniques: the memory bottleneck, Online Softmax, and tiled matrix multiplication.
Before we dive into the code implementation of Flash Attention, we must first answer a foundational question: why is the standard attention formula still not fast enough on modern GPUs?
First, we need to establish an extremely important concept: in the GPU architecture, all computation (such as matrix addition and multiplication) must happen in the SRAM (Shared Memory) that sits close to the cores.
This means that even if your VRAM (HBM) is as large as 80GB, the data must first be "moved" into the SRAM, which is only tens of MB in size, before it can be processed by the compute cores.
This is premium content. Please log in to access the full article.
In the "naive" implementations of deep learning frameworks such as PyTorch, the Attention computation is split into multiple independent operators (Ops). This leads to a serious efficiency problem:
GPU moves and from HBM to SRAM and computes the score matrix .# The IO nightmare of the standard Attention implementation
def standard_attention(Q, K, V):
# 1. HBM -> SRAM (compute) -> HBM (store S)
S = Q @ K.T
# 2. HBM (read S) -> SRAM (compute) -> HBM (store P)
P = softmax(S)
# 3. HBM (read P) -> SRAM (compute) -> HBM (store O)
O
This repeated I/O round trip of "move in -> compute -> evict -> move back" is the biggest performance killer.
You might ask: how much impact can storing back to HBM and reading it back again really have?
| Storage type | Capacity example (A100) | Bandwidth | Speed metaphor |
|---|---|---|---|
SRAM (shared memory) | ~20 MB | ~19 TB/s | F1 race car 🏎️ |
HBM (VRAM) | 40~80 GB | ~1.5 TB/s | ordinary sedan 🚗 |
SRAM's bandwidth is typically more than about 10 times higher than HBM's.
Because of this speed gulf, if the algorithm keeps moving intermediate data back and forth between HBM and SRAM, an awkward situation arises:
The GPU's powerful compute cores spend most of their time "idling," waiting in agony for data to be delivered from the slow HBM.
This state is called IO-bound, meaning that the compute capability is held back by the memory transfer speed.
To get a quantitative feel for it: in standard Attention, the time spent reading and writing the intermediate matrices and far exceeds the time spent actually performing the matrix multiplication computation.
Since SRAM is so fast, why not just put the entire Attention matrix in SRAM, finish computing, and be done? This is where rigid physical and economic constraints come into play:
SRAM has an extremely low storage density, which makes its cost extremely high. According to relevant materials (such as the background cited in the FlashAttention paper):
80GB capacity SRAM memory could cost as much as $13,000 (an order-of-magnitude estimate).HBM of the same capacity costs only about $2,000.In real hardware, an A100's HBM can reach 80GB, but SRAM is usually only 192 KB / SM (per streaming multiprocessor). Because of this capacity limit, you cannot cram the entire attention matrix into SRAM all at once. As the sequence length grows, the size of the intermediate matrix explodes at .
The core logic of Flash Attention lies right here: since SRAM is expensive and small but fast, while HBM is large and cheap but slow, we must abandon the fantasy of "reading and writing everything at once."
We need to introduce two core ideas:
SRAM.SRAM, complete the , Softmax, and multiplication by in one go, writing only the final result back to HBM in the last step.In this way, we never even generate (nor write to HBM) that huge intermediate matrix.
| Method | HBM read/write volume | Complexity |
|---|---|---|
Standard Attention | as the sequence gets longer, IO explodes | |
Flash Attention | linear growth, hugely saving bandwidth |
SRAMGPU has no choice but to "evict" this giant matrix from SRAM and write it back to the slow HBM.GPU has to go back to HBM and move the it just stored back into SRAM again.