Walk through the FlashAttention algorithm: what does it change vs standard attention, why does that change matter for GPU performance, and what evolves between v1, v2, and v3?
FlashAttention keeps the n×n matrix out of HBM by tiling Q/K/V into SRAM and using online softmax. v1 introduced tiling, v2 added sequence dim parallelism, v3 went Hopper async + FP8. Same math, 2-4x faster.
Standard attention is like solving a huge crossword by repeatedly carrying the whole grid back and forth between your desk and a shelf across the room. FlashAttention does it section by section at your desk and never brings the whole grid out, so you get exactly the same answers without wasting time on the walking. The desk space is small but the trip to the shelf is the slow part, and the section by section trick avoids it almost entirely. Each new version (v2, v3) is a smarter way of organizing the same trick to fit the latest desks (GPUs).
Concept explanation~2 min read
Everything you need to truly understand this topic: intuition, mechanics, step by step explanation, code, formulas, and worked example. Click to expand.
Concept explanation~2 min read
Everything you need to truly understand this topic: intuition, mechanics, step by step explanation, code, formulas, and worked example. Click to expand.
FlashAttention is the most important attention kernel optimization of the last few years. It makes long context training and inference tractable on modern GPUs without changing the math at all, and the lesson it teaches about I/O-aware kernel design now governs how nearly every transformer is served in production.
Understanding it requires understanding GPU memory hierarchy: tensor cores want data far faster than HBM can deliver it, and attention's natural implementation produces an n×n intermediate that has to go through HBM repeatedly. FlashAttention says: do not write that intermediate to HBM at all. The sections below walk the bandwidth bound failure mode of standard attention, the tiling and online softmax mechanics that make v1 work, the asymmetry between forward and backward, the v2 parallelism redesign, the Hopper-specific v3 path, and what FlashAttention is not.
The HBM-bandwidth problem in standard attention
On A100, FP16 tensor core peak is ~312 TFLOPs but HBM bandwidth is only ~2 TB/s. The arithmetic intensity works out to ~150 FLOPs per byte: any kernel that reads or writes more bytes than it has FLOPs to spend on them stalls the tensor cores.
Where the time goes. Standard attention produces an n×n matrix (4096×4096 = 16M entries per head per layer at n=4k) that is written to HBM after QK^T, read back for softmax, written again, then read once more for the V multiply. Those reads and writes operate at very low arithmetic intensity, far below the 150 FLOPs/byte that would saturate the chip.
The matmuls finish in a fraction of the time the I/O takes; the GPU is idle waiting on memory. Optimizing the matmul kernel further is wasted effort. The right fix is to eliminate the HBM round trips on the n×n matrix entirely.
Attention is bandwidth bound on every modern accelerator. The score matrix is the offending data structure.
Situations where this technique stops working.
2–4 min · Everything important, quickly.
| Aspect | Standard attention | FlashAttention v1 | FlashAttention-v2 | FlashAttention v3 |
|---|---|---|---|---|
| Year | 2017 | 2022 | 2023 | 2024 |
| Materializes n×n in HBM? | Yes | No | No | No |
| Memory complexity | O(n²) | O(n) | O(n) | O(n) |
| Parallelism | Batch × heads | Batch × heads | Batch × heads × seq | Batch × heads × seq + async |
| Target hardware | Any | Ampere+ | Ampere+ | Hopper (H100) |
| Precision | FP16/BF16 | FP16/BF16 | FP16/BF16 | FP16/BF16/FP8 |
Real products, models, and research that use this idea.
- PyTorch's torch.nn.functional.scaled_dot_product_attention dispatches to FlashAttention-v2 kernels on supported GPUs by default.
- vLLM and TensorRT-LLM both use FlashAttention style kernels in their decode path, often combined with PagedAttention for KV cache.
What an interviewer would ask next. Try answering before peeking at the approach.
QHow does the backward pass avoid storing the n×n attention matrix?
Forward stores only Q, K, V plus per-row m and l. Backward recomputes attention tile by tile from those during backprop. Trade: extra FLOPs (cheap on tensor cores) for huge memory savings (expensive in HBM).
Red flags & common mistakes
The phrases that signal junior thinking. Click to expand.
Red flags & common mistakes
The phrases that signal junior thinking. Click to expand.
Treating FlashAttention as an approximation, or attributing the speedup to a faster matmul kernel rather than to eliminating HBM round trips.
60 second bullets to scan on the way to the call.
Why attention is bound by memory bandwidth on modern GPUs
Role of tiling Q, K, V into SRAM blocks
Primary sources. Browse if you want the original framing.
Same topic, related formats. Practice these next.