Lensa ML
Lensa ML

FlashAttention

Step 1 of 6: The Memory Hierarchy

Large, slow HBM and small, fast SRAM

Data to read32 MB
Capacity (log scale), illustrative A100-class GPUSRAM, one SM: 192 KiBSRAM, all 108 SMs: 21 MBHBM: 80.0 GB10 KB1.0 MB100 MB10.0 GBdata: 32 MBBandwidth (linear scale)HBM 2 TB/sSRAM about 19 TB/s
Data size
32 MB
Read time from HBM
15.8 µs
At SRAM speed
does not fit
Bandwidth ratio
9.5×

32 MB only fits in HBM. Each full pass over it costs 15.8 µs of bandwidth, 1.5× the whole chip's SRAM.

Drag Data to read from 10 KB to 100 GB and watch where the dashed line crosses each memory level. An A100 does about 312 TFLOP/s of fp16 matrix math but reads HBM at about 2 TB/s, so a kernel needs roughly 150 FLOP per byte to keep the arithmetic units busy. Softmax, masking, and dropout do a few FLOP per element and are therefore bound by memory. For scale, one head's score matrix at 4,096 tokens is 33.6 MB in fp16, already larger than all on-chip SRAM.

Step 2 of 6: Standard Attention

The N×N matrix makes round trips to HBM

Sequence length N4,096
One head, d = 128, fp16: three kernels, each through HBMHBMQ1.0 MBK1.0 MBS34 MBP34 MBV1.0 MBO1.0 MBS = QKᵀP = softmax(S)O = PVkernels run on chip, with SRAMHBM traffic per head: 138 MBS and P reads/writes 97%Q, K, V, O 3%
S per head (fp16)
34 MB
S, 32 heads, one layer
1.1 GB
HBM traffic per head
138 MB
Layer time at 2 TB/s
2.21 ms

At N = 4,096, one head's S is 34 MB, larger than all on-chip SRAM, and the 32 heads of one layer need 1.1 GB. Moving it dominates the 2.21 ms the layer spends on memory.

Drag Sequence length N from 256 to 32,768. A standard implementation runs separate kernels, and each writes its result to HBM for the next one to read. Per head that is about 4N² + 4Nd elements moved, while the arithmetic is about 4N²d FLOP, which the GPU could finish quickly. Masking and dropout add more passes over the same N×N matrix. Training also keeps P for the backward pass. For the 7B model at 4,096 tokens that means 1.07 GB per layer across its 32 heads.

Step 3 of 6: Tiling

Blocks of Q, K and V that fit on chip

Block size64 rows
Score matrix, 256 × 256 tokensK/V tiles →Q blocks ↓grey: computed, used, discardedSRAM of one SMQ block (64×128)16 KiBK tile (64×128)16 KiBV tile (64×128)16 KiBS tile (64×64)8.0 KiBO block (64×128)16 KiBused 72 KiB of 192 KiBfp16 tiles, d = 128; fp32 accumulators and statistics not counted
Tiles in S
4 × 4
Current tile
none
SRAM per tile step
72 KiB
S written to HBM
0 B, not 131 KB

The 256×256 score matrix splits into 16 tiles of 64×64. Press Play or Step to run the loop: one Q block stays on chip while K/V tiles stream past it.

Set Block size, then press Play or Step to run the loop. The paper picks K/V tiles of ⌈M/4d⌉ rows and Q blocks of min(⌈M/4d⌉, d) rows, where M is SRAM size. The original kernel put K/V tiles in the outer loop. FlashAttention-2 swapped the loops, as shown here, so each thread block owns one Q block and writes its output once. Matmul, mask, softmax, dropout, and the second matmul all run in one fused kernel. Real kernels also need fp32 accumulators, so 64 or 128 rows is typical.

Step 4 of 6: Online Softmax

A running max and sum, rescaled per chunk

Score x₅4.5
Chunk size
1.2x13.0x20.4x32.1x44.5x51.8x63.6x70.9x8m = −∞, ℓ = 0Bars: e^(x − m) / ℓ so far. Outlines: the full softmax.0.020.120.010.050.540.040.220.01numbers under the bars: full softmax over all 8 scores
Running max m
−∞
Running sum ℓ
0
Rescale factor
none yet
Max |online − full|
0/8 seen

Nothing is read yet. Step through the 4 chunks of 2 and watch m and ℓ update.

Drag Score x₅, pick a Chunk size, then Step through the chunks. The update is m_new = max(m, chunk max) and ℓ_new = ℓ · e^(m − m_new) + Σ e^(x − m_new). Subtracting the max keeps exponents small: fp16 overflows above 65,504, which is e^11.1. Milakov and Gimelshein (2018) published this single-pass normalizer for softmax. FlashAttention applies the same rescale factor to the running output, the unnormalized sum of weights times V, and divides by ℓ once the last K/V tile is done.

Step 5 of 6: IO Savings

Fewer HBM accesses, linear memory

Sequence length N8,192
SRAM per SM (M)192 KiB
Head dimension d
Extra memory per head (log scale)5122K8K32K128K1.0 KB1.0 MB1.0 GBsequence length N (log scale)standard: N² scoresFlashAttention: m and ℓ, O(N)HBM accessesstandardFlash
Standard, Nd + N²
68.2M
Flash, N²d²/M
11.2M
Fewer accesses by
6.1×
Extra memory per head
134 MB vs 66 KB

With M = 98,304 elements and d = 128, the illustrative ratio is about M / d² = 6.1× fewer HBM accesses, with K/V tiles of about 192 rows. Extra memory drops from 134 MB to 66 KB per head at N = 8,192.

Drag Sequence length N and SRAM per SM, and switch Head dimension d. Theorem 2 of the FlashAttention paper gives Θ(Nd + N²) HBM accesses for standard attention and Θ(N²d²/M) for FlashAttention. Both still grow with N², and the FLOP count does not go down. The backward pass even recomputes S from Q and K, adding arithmetic, yet runs faster because it reads less. The paper also shows that no exact attention algorithm can do asymptotically better across a range of SRAM sizes.

Step 6 of 6: Exact, and Everywhere

FlashAttention 1, 2 and 3

Version
Block size (exactness check)4 keys
202220232024FlashAttention-2, A100 (Ampere)Parallel over sequence length as well as heads;fewer non-matmul FLOPs, better warp splitReported: about 2× FA1, 50-73% of peak FLOP/sQuery 1's output, d = 8, random Q, K, V (illustrative)outline: standard attention; fill: tiled with online softmax
K/V tiles
4
Max |tiled − standard|
2.2e-16
Target GPU
A100 (Ampere)

Split into 4 tiles of up to 4 keys, the tiled output differs from standard attention by at most 2.2e-16, floating-point rounding only. FlashAttention-2 changes how the loop is scheduled on the A100 (Ampere), not this result.

Switch Version and drag Block size to rerun the exactness check. Approximate methods such as sparse or low-rank attention change the result to save compute, while every FlashAttention version reproduces standard attention up to rounding. In PyTorch 2.0 and later, torch.nn.functional.scaled_dot_product_attention picks a FlashAttention kernel when the dtype is fp16 or bf16 and the shapes allow it, falling back to other kernels otherwise. Hugging Face Transformers exposes it as attn_implementation="flash_attention_2".

What is FlashAttention?

FlashAttention is an exact attention algorithm that computes softmax(QKᵀ/√d)V in small tiles kept in on-chip GPU memory, so the full N×N score matrix is never written to slow off-chip memory. It gives the same output as standard attention. The gain comes from moving fewer bytes, not from doing less arithmetic.

Attention is the part of a transformer whose cost grows with the square of the sequence length, and on GPUs its standard form is limited by memory traffic. FlashAttention speeds up GPT-2 training about 3× at 1K tokens and cuts attention's extra memory from quadratic to linear in sequence length. Its kernels now sit behind PyTorch's scaled_dot_product_attention, Hugging Face Transformers, vLLM, and most LLM training stacks.

Key ideas

  • Standard attention is limited by memory traffic, because it writes and reads the N×N score matrix through HBM several times.
  • FlashAttention splits Q, K, and V into tiles that fit in on-chip SRAM and fuses all of attention into one kernel.
  • Online softmax keeps a running max and sum per row, so softmax can be computed one tile at a time and rescaled at the end.
  • HBM accesses drop from Θ(Nd + N²) to Θ(N²d²/M), and extra memory drops from O(N²) to O(N).
  • FlashAttention is exact: it returns the same output as standard attention, apart from floating-point rounding.
  • FlashAttention-2 improves parallelism and work partitioning, and FlashAttention-3 targets Hopper GPUs with asynchrony and FP8.

Common questions

Do I need to retrain a model to use FlashAttention?

No. FlashAttention computes the same function as standard attention, so existing weights work unchanged. Swapping the attention kernel is a runtime choice, often a single flag. Outputs can differ in the last bits because the additions happen in a different order.

Why is FlashAttention faster if it does more FLOPs?

On modern GPUs, standard attention spends most of its time waiting on memory, not doing arithmetic. FlashAttention recomputes the score matrix in the backward pass instead of storing it, which adds FLOPs but removes far more bytes of HBM traffic. Since memory was the bottleneck, total time goes down.

Does FlashAttention help during LLM decoding?

Prefill, where the whole prompt is processed at once, benefits the most because it builds large score matrices. During decode each step has one query per sequence, so the score matrix is a single row and the main cost is reading the KV cache. Variants such as Flash-Decoding split that read across many thread blocks to use the whole GPU.

Which GPUs support FlashAttention?

FlashAttention-2 runs on NVIDIA Ampere, Ada, and Hopper GPUs such as the A100, RTX 4090, and H100, and the original version also covers older Turing cards. FlashAttention-3 is written for Hopper. AMD GPUs have ROCm ports, and PyTorch falls back to other kernels where no FlashAttention kernel applies.

Related explainers

References

  1. Dao et al. (2022), "FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness"
  2. Dao (2023), "FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning"
  3. Shah et al. (2024), "FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision"
  4. Milakov and Gimelshein (2018), "Online normalizer calculation for softmax"

Last updated