Skip to content

FlashAttention Explained: How It Makes Transformers Faster on GPUs

FlashAttention computes exact attention without writing the huge attention matrix to GPU memory, using tiling and an online softmax. How it works, why it matters for long contexts, and what changed in FlashAttention 2, 3 and 4.

4 min read

FlashAttention Explained cover image from TensorCore.si

Short answer

FlashAttention is an algorithm that computes exactly the same result as standard attention while never writing the full attention matrix to GPU memory. It splits queries, keys and values into tiles that fit in fast on-chip memory, computes softmax piece by piece with an online softmax, and recomputes values in the backward pass instead of storing them. Memory use grows linearly with sequence length instead of quadratically, and attention runs several times faster. FlashAttention-4 (2026) reaches about 1,600 TFLOPS in BF16 on NVIDIA Blackwell, about 71% of peak.

The problem with standard attention

Attention computes softmax(Q × Kᵀ) × V. For a sequence of N tokens, Q × Kᵀ is an N × N matrix for every attention head. A standard implementation writes that matrix to GPU memory, reads it back for the softmax, writes the result, and reads it again to multiply by V.

At long contexts the matrix becomes huge. For one 32,768-token sequence with 32 heads in 16-bit, it is about 69 GB for a single layer. Even when it fits, moving it in and out of HBM takes far longer than the math, so the Tensor Cores wait on memory.

How FlashAttention works

  1. Tiling: split Q, K and V into blocks that fit in the shared memory and registers of each streaming multiprocessor.
  2. Compute on chip: for each block of queries, loop over blocks of keys and values, multiplying them with Tensor Cores without leaving the chip.
  3. Online softmax: softmax normally needs the whole row's maximum and sum. FlashAttention keeps a running maximum and sum, rescaling partial results as new blocks arrive, so the final answer is exact.
  4. Write once: only the final output, N × head dimension, goes back to HBM.
  5. Recompute in training: the backward pass recomputes the attention scores block by block instead of storing the N × N matrix, trading a little extra math for much less memory traffic.

The result is "IO-aware" attention: it is designed around the cost of moving data, which is the real bottleneck on modern GPUs. See how a GPU works for the memory hierarchy it exploits.

FlashAttention versions

VersionYearMain GPUsWhat changed
FlashAttention2022Ampere (A100)Tiling, online softmax and recomputation; linear memory in sequence length
FlashAttention-22023Ampere, Ada, HopperBetter work split across thread blocks and warps, and parallelism over sequence length; about twice as fast as version 1
FlashAttention-32024Hopper (H100)Uses Hopper's asynchronous Tensor Cores and Tensor Memory Accelerator, overlaps softmax with matrix math, adds FP8; up to about 740 TFLOPS in FP16, 75% of H100 peak
FlashAttention-42026Blackwell (B200)Pipelines matrix math and softmax using Blackwell's tensor memory, computes some exponentials in software, and is written in a Python kernel language; up to about 1,613 TFLOPS in BF16, 71% of peak

The FlashAttention-4 authors note a trend: Tensor Core throughput has grown faster than the units that compute exponentials and than shared memory bandwidth, so softmax and data movement increasingly limit attention. NVIDIA's Blackwell Ultra doubled its exponential units for the same reason.

Why it matters

  • Long contexts: without FlashAttention, 100,000-token and million-token context windows would be impractical.
  • Training speed: attention is a large share of transformer compute, and faster attention shortens training runs.
  • Inference: faster prefill for long prompts and less memory per request. See prefill vs decode.
  • Exactness: unlike sparse or approximate attention, it changes no results apart from normal floating-point rounding.

Where you are already using it

You rarely call FlashAttention directly. PyTorch's scaled_dot_product_attention can use a FlashAttention kernel automatically, Hugging Face Transformers can load models with it, and inference engines such as vLLM, SGLang and TensorRT-LLM use FlashAttention or similar fused attention kernels. NVIDIA's cuDNN library has adopted many of the same techniques.

Frequently asked questions

Is FlashAttention an approximation?

No. It computes exact attention. Results differ from a standard implementation only by tiny floating-point rounding differences.

Which GPUs support FlashAttention?

FlashAttention-2 supports NVIDIA Ampere, Ada and Hopper GPUs, FlashAttention-3 targets Hopper, and FlashAttention-4 targets Blackwell. AMD maintains versions for Instinct GPUs through ROCm.

Does FlashAttention reduce the KV cache?

No. It removes the temporary attention matrix, not the stored keys and values. Reducing the KV cache needs other techniques such as grouped-query attention or KV cache quantization. See KV cache explained.

Who created FlashAttention?

Tri Dao and colleagues, first at Stanford in 2022. Later versions were developed with collaborators at Princeton, Together AI, Colfax Research, Meta and NVIDIA.

Sources and further reading

Think you understand it now? Prove it.

Solve real Tensor Core math in the TensorCore.si Challenge and earn a certificate with your name.

Take the challenge