FlashAttention: The IO-Aware Trick That Made Long Context Cheap
Deep Dives 9 min read advanced

FlashAttention: The IO-Aware Trick That Made Long Context Cheap

FlashAttention is an IO-aware, exact attention algorithm from 2022 that avoids writing the full N-by-N attention matrix to slow GPU HBM. Using tiling, an online-softmax running-statistics trick, kernel fusion, and recomputation, it cuts memory from O(N^2) to O(N) and delivered up to 7.6x speedups. FlashAttention-2 reached ~70% of A100 peak FLOPs; FlashAttention-3 (2024) exploits Hopper asynchrony and FP8 to hit ~840 TFLOPs BF16 (~75% H100 utilization). It now powers PyTorch, vLLM, and long-context serving.

Aisha Patel
Aisha Patel
Aug 1, 2026

Every time you paste a 100,000-token document into a chatbot and get an answer in seconds, you are benefiting from a 2022 paper that most people using AI have never heard of. FlashAttention did not make models smarter. It made attention — the single most expensive operation in a transformer — dramatically cheaper to compute, and in doing so it quietly unlocked the long-context era we now take for granted.

This is a deep look at what FlashAttention actually does, why the "obvious" way of computing attention was wasting most of your GPU, and how three generations of the algorithm pushed hardware utilization from roughly a third of a chip's potential to three-quarters of it.

The problem: attention is memory-bound, not compute-bound

The standard attention computation looks innocent. Given queries Q, keys K, and values V, you compute:

S = Q Kᵀ        # scores:      N × N
P = softmax(S)  # weights:     N × N
O = P V         # output:      N × d

where N is the sequence length and d is the head dimension. The trouble is that intermediate N × N matrix. For a sequence of 8,000 tokens, that is 64 million numbers — per head, per layer. Memory cost grows with the square of sequence length, which is exactly why long context was historically so painful.

But the deeper problem is subtler and it is about where those numbers live. A modern GPU has two very different kinds of memory:

  • HBM (high-bandwidth memory) — the big pool, 40–80 GB on an A100, but "only" about 1.5–2 TB/s of bandwidth.
  • SRAM (on-chip) — tiny, around 20 MB, but roughly an order of magnitude faster.

Here is the mismatch that defines the whole problem. An A100 can do 312 TFLOPS of FP16 math but only move about 2 TB/s through HBM. The arithmetic units are so fast that for an operation like attention, the GPU spends most of its time waiting on memory, not computing. The standard implementation makes this worse by writing the full S matrix to HBM, reading it back to compute the softmax, writing P back out, and reading it again for the final multiply. Attention becomes memory-bound: bottlenecked by data movement, not by floating-point throughput.

You can buy a faster GPU and barely move the needle, because the FLOPs were never the constraint.

The insight: never materialize the big matrix

FlashAttention, introduced by Tri Dao and colleagues in May 2022, is built on a deceptively simple idea: compute the exact same attention output without ever writing the full N × N matrix to HBM. It is not an approximation. The numbers that come out are bit-for-bit the standard result. What changes is the memory choreography.

Three techniques make this work.

1. Tiling

Instead of computing all of S at once, FlashAttention splits Q, K, and V into blocks small enough to fit in SRAM. It loads a block of queries and a block of keys/values, computes their partial scores on-chip, and moves on. The giant intermediate matrix never exists in slow memory — only small tiles do, briefly, in fast memory.

2. Online softmax

Tiling creates a puzzle: softmax needs to normalize across an entire row, but tiling only ever sees one block of that row at a time. The solution is the online softmax — a running-statistics trick. As each new key block arrives, the algorithm keeps a running maximum and a running sum, and rescales the partial output it has accumulated so far. When the last block is processed, the result is mathematically identical to having seen the whole row at once. This is the mathematical heart of the algorithm, and it is what makes "exact attention without the full matrix" possible.

3. Kernel fusion and recomputation

The whole sequence — matmul, softmax, matmul — runs inside a single fused GPU kernel, so intermediates stay in SRAM instead of round-tripping to HBM. And rather than storing the attention matrix for the backward pass, FlashAttention simply recomputes it during backpropagation. Recomputation sounds wasteful, but recomputing in fast memory is cheaper than reading from slow memory. This is the counterintuitive core of IO-aware design: doing more math to do less memory movement is a net win.

Why it worked so well

The payoff shows up in both speed and memory:

  • Memory drops from O(N²) to O(N). Because the full matrix never materializes, memory now grows linearly with sequence length instead of quadratically. This is the single change that made long context economically viable.
  • HBM accesses fall sharply. The paper shows FlashAttention needs O(N²d²M⁻¹) HBM accesses (where M is SRAM size) versus Ω(Nd + N²) for standard attention — a large reduction for realistic d and M.
  • Real speedups. The original paper reported up to 7.6× faster attention on GPT-2-scale workloads, with wall-clock training speedups across BERT and GPT models.

Crucially, none of this touched model quality. FlashAttention is a systems optimization, not a modeling change. You get the identical model, trained and served faster and on longer sequences.

Three generations of getting closer to the metal

The idea proved so foundational that it kept evolving with the hardware.

FlashAttention-2 (2023) rethought how work is partitioned across a GPU's thread blocks and warps, cutting non-matmul overhead and improving parallelism. The result was roughly 2× faster than the original, reaching around 70% of the A100's theoretical peak FLOPs — extraordinary for an operation that used to leave most of the chip idle.

FlashAttention-3 (July 2024) was rebuilt for NVIDIA's Hopper (H100) architecture and leaned into features the earlier versions couldn't exploit:

  • Asynchrony — overlapping computation with data movement using the Tensor Memory Accelerator (TMA) and warp specialization, so the chip loads the next tile while computing the current one.
  • Interleaving matmul and softmax operations so the two never fully block each other.
  • FP8 low precision with block quantization to keep accuracy acceptable.

The numbers are striking. In BF16, FlashAttention-3 reaches up to 840 TFLOPs/s — about 75% H100 utilization, up from roughly 35% with the previous approach. In FP8, it approaches 1.2 PFLOPs/s, while keeping numerical error about 2.6× smaller than a naive FP8 attention baseline. That utilization jump — from a third of the chip to three-quarters — is money: it is the difference between renting one H100 and renting two.

A concrete intuition for the online softmax

If one part of FlashAttention feels like magic, it is the claim that you can compute a softmax without ever seeing the whole row. Here is the intuition, stripped of notation.

A softmax needs two things about the row: the maximum value (subtracted for numerical stability) and the sum of the exponentials (the normalizer). Both seem to require the full row. But you can maintain them incrementally. As each new block of scores arrives, you check whether it contains a value larger than your running maximum. If it does, you have been under-scaling everything so far, so you retroactively correct the accumulated sum and the accumulated output by a single multiplicative factor, then fold in the new block.

The key property is that this correction is exact, not approximate — it produces the identical number the full-row softmax would have. You are trading a large memory footprint for a little extra bookkeeping arithmetic, and on a GPU that trade is wildly favorable. This same running-statistics pattern shows up elsewhere in high-performance computing precisely because it lets you stream data that is too big or too slow to hold all at once.

What FlashAttention does *not* do

It is worth being precise about the limits, because FlashAttention is sometimes oversold. It does not change the asymptotic compute of attention — you still do roughly O(N²) floating-point work, because every query still attends to every key. What drops to linear is memory, not FLOPs. For truly enormous sequences you still eventually want algorithmic sparsity (the approach DeepSeek and others take with sparse attention) on top of FlashAttention's systems-level wins.

It is also exact, which is a feature but also a boundary: it will not give you the accuracy-for-speed trade that approximate methods offer. And its gains are hardware-specific — the whole point is to match a particular GPU's memory hierarchy, which is why each new NVIDIA architecture needed a new FlashAttention rewrite rather than a free ride.

Where it lives now

You almost certainly use FlashAttention without invoking it by name. It is the default behind PyTorch's scaled_dot_product_attention, it powers inference in vLLM and the Hugging Face stack, and it is standard equipment in training pipelines across the industry. When a model advertises a 128K or 1M context window, FlashAttention (or a descendant of its IO-aware philosophy) is almost always part of what makes that window affordable to serve.

It also reframed how the field thinks about performance. Before FlashAttention, "make attention faster" usually meant "approximate it" — sparse patterns, low-rank tricks, linear attention. FlashAttention showed that you could keep attention exact and still win enormously, simply by respecting the memory hierarchy of the hardware you run on. That lesson — count your memory accesses, not just your FLOPs — now shapes GPU kernel design far beyond attention.

The Bottom Line

FlashAttention is one of those rare optimizations that changed what was possible without changing what the model does. By refusing to write the giant attention matrix to slow memory — and instead tiling, streaming a running softmax, and fusing everything into one kernel — it turned attention from a memory-bound bottleneck into something that can nearly saturate modern GPUs. The long-context models we now treat as normal are standing on this work. If you remember one idea from it, make it this: on today's hardware, the scarce resource is memory bandwidth, and the biggest wins come from moving less data, not doing less math.