上一代的问题：Standard attention materializes the full N×N attention matrix in HBM — each forward pass reads and writes the huge matrix from slow GPU memory multiple times: write QKT to HBM, read it for softmax normalization, write the normalized matrix back, read it for the V multiplication. This causes O(N2) memory and bandwidth cost, making attention the dominant bottleneck for long sequences. Prior work focused on reducing FLOPs (sparse/linear attention) but ignored the IO bottleneck entirely.
这代改了什么：Tiling + Online Softmax — split Q, K, V into blocks that fit entirely in fast on-chip SRAM (~20 MB, 19 TB/s). Each block computes the attention scores and weighted sums on-chip without ever writing the full N×N matrix to HBM. A novel online softmax algorithm accumulates the correct normalization incrementally, block by block, using only two running statistics kept in registers. Only the final output O is written back to HBM.
效果：2–4× end-to-end training speedup (7.6× on the attention kernel alone). Memory scales linearly with sequence length instead of quadratically. Enabled 128K context in Llama 3, Multi-head Latent Attention (MLA), Grouped Query Attention (GQA), and virtually all modern long-context LLMs. Won the 2024 NeurIPS Test of Time Award.
The fundamental insight of FlashAttention is that the GPU has a heterogeneous memory hierarchy with vastly different bandwidths and capacities, and standard attention is entirely bottlenecked by the slowest layer.
| Memory Level | Size | Bandwidth | Speed Ratio |
|---|---|---|---|
| HBM (High Bandwidth Memory) | 40–80 GB (A100/80) | ~1.5–2.0 TB/s | 1× |
| L2 Cache | ~40 MB | ~4–8 TB/s | ~3–5× |
| SRAM (Shared Memory / Registers) | ~20 MB (192 KB per SM × 108 SMs) | ~19 TB/s | ~10–12× |
Standard attention does the following dance with HBM:
For d=64, n=1024, the attention matrix is 8 MB — it just barely fits in SRAM. But at n=8192 it's 512 MB, and at n=128K it's 128 GB, far exceeding even HBM capacity. This means standard attention fails entirely for long sequences on a single GPU, forcing approximations.
The core innovation is decomposing the attention computation into blocks that fit in SRAM, combined with an online softmax that computes the correct normalization without ever seeing all values at once.
The Q matrix is split into row blocks Qi (size Br × d). The K and V matrices are split into column blocks Kj, Vj (size Bc × d). The outer loop iterates over Q blocks; the inner loop over K/V blocks.
For each block iteration:
Standard softmax for a row x of length N:
The problem: computing the correct normalization requires knowing the global max and sum, but we only see one block at a time. Online softmax solves this by maintaining two running statistics and correcting when a larger max is found:
When mnew > mold (a larger maximum appears in the current block), all previously accumulated contributions are rescaled by the factor emold - mnew, which is < 1. This ensures the final output is bitwise identical to standard softmax. The two running statistics m and d are stored in registers, never written to HBM.
FlashAttention provides a rigorous IO complexity bound that proves its optimality.
Let M be the SRAM size. The block sizes are chosen as:
Total HBM accesses for FlashAttention:
Compare with standard attention:
For typical values (d=64, M≈100 KB per thread block):
Importantly, the paper proves that this bound is asymptotically optimal: no exact attention algorithm (that computes the same mathematical function) can achieve asymptotically fewer HBM accesses for all SRAM sizes M. The proof uses a reduction from the matrix multiplication IO lower bound.
The backward pass of standard attention needs the full softmax output matrix P (N×N) to compute gradients. Storing P for backprop would require O(N2) HBM — exactly what FlashAttention avoids.
FlashAttention's solution: recompute on the fly. During the forward pass, only the softmax normalization statistics (m, d) are stored — two vectors of length N, total O(N) storage. During the backward pass:
The recomputation adds approximately 1/3 more FLOPs (recomputing the attention matrix in the backward pass) but saves an enormous amount of HBM bandwidth — not storing N×N can be the difference between fitting in VRAM and OOM. In practice, on GPU architectures, the extra compute is cheap because SRAM bandwidth is so much higher than HBM bandwidth.
FlashAttention naturally extends to block-sparse attention, where a sparsity mask indicates which blocks of the attention matrix are non-zero and should be computed. The tiled structure means skipping a block is as simple as not loading Kj/Vj for that iteration.
With a carefully designed sparsity pattern (e.g., local windows plus global tokens), block-sparse FlashAttention achieves 2–4× further speedup while maintaining better-than-chance accuracy on tasks where full attention is impractical. This was used to achieve 63.1% accuracy on Path-256 (64K tokens) — the first model to beat random chance on this benchmark.
| Model / Task | Sequence Length | Speedup | Notes |
|---|---|---|---|
| BERT-large training | 512 | 15% | Faster than MLPerf 1.1 record |
| GPT-2 training | 1,024 | 3× | vs HuggingFace / Megatron-LM baseline |
| Long Range Arena avg | 1,000–4,000 | 2.4× | Across all 5 tasks |
| Attention kernel only | variable | 7.6× | Measured on A100 |
| Path-X (first Transformer above chance) | 16,384 | — | 61.4% accuracy |
| Path-256 | 65,536 | — | 63.1% accuracy (block-sparse) |
The impact goes beyond raw speed. FlashAttention's linear memory scaling means that training with 128K context (and more recently, 1M+ context) is feasible on a single node. Every major LLM post-2022 uses some variant of FlashAttention or its descendants.
# Conceptual FlashAttention forward pass (simplified)
def flashattention(Q, K, V, Br, Bc):
# Q: (n, d), K/V: (n, d)
# Br, Bc: block sizes for Q and K/V
n = Q.shape[0]
O = zeros(n, d)
m = full(n, -inf) # running max per row
d = zeros(n) # running sum per row
for i in range(0, n, Br):
Qi = Q[i:i+Br]
Oi = O[i:i+Br]; mi = m[i:i+Br]; di = d[i:i+Br]
for j in range(0, n, Bc):
Kj = K[j:j+Bc]; Vj = V[j:j+Bc]
Sij = Qi @ Kj.T # (Br, Bc) on SRAM
mij = max(Sij, axis=1) # row max of this block
mi_new = max(mi, mij) # update running max
Pij = exp(Sij - mi_new[:, None]) # safe softmax numerator
di_new = di * exp(mi - mi_new) + sum(Pij, axis=1)
# Rescale accumulated output and add new contribution
Oi = Oi * exp(mi - mi_new)[:, None] + Pij @ Vj
mi = mi_new; di = di_new
O[i:i+Br] = Oi / di[:, None] # final normalization
m[i:i+Br] = mi; d[i:i+Br] = di
return O
上一篇（问题来源）：Before FlashAttention, the attention efficiency literature was dominated by sparse attention (Child et al. 2019, BigBird, Longformer, ETC) and linear attention (Katharopoulos et al. 2020, Performer, Linformer). These works focused on reducing the FLOP count of attention by approximating the full softmax attention with cheaper alternatives. However, none of them addressed the memory access pattern — on modern GPUs, the standard attention implementation is already compute-bound for small N but memory-bound for large N, and reducing FLOPs does not reduce HBM accesses for the attention matrix that is still materialized.
下一篇（影响扩散）：
被谁采用：Every major transformer library and model since 2023:
引用/采用场景：
FlashAttention is, in retrospect, one of those ideas that feels obvious once explained. The fundamental problem was well-known: attention is O(N2) in both compute and memory, and GPUs have a fast but tiny on-chip memory. The entire research community was chasing FLOP reduction while leaving the real bottleneck — HBM bandwidth — untouched.
The lesson is profound for ML systems research: optimize for the real bottleneck, not the theoretical one. On modern hardware, memory bandwidth, not FLOPs, is almost always the binding constraint. This is the same principle that made cuBLAS successful (tiling matrix multiplies for cache), the same principle behind operator fusion in XLA/TorchDynamo, and the same principle behind flash attention for other primitives (FlashFFTConv, FlashDecoding).
Three things make FlashAttention a landmark paper:
In the broader arc of transformer research, FlashAttention sits alongside the original transformer paper and the GPT-2 scaling paper as a foundational infrastructure contribution. It didn't invent a new architecture or a new capability; it removed a bottleneck that was silently limiting the entire field. The NeurIPS 2024 Test of Time Award is well-deserved.