FlashAttention

Compute attention in tiles .

FlashAttention computes softmax attention in small tiles that fit in on-chip memory. It avoids storing the full score matrix in HBM, reducing memory traffic without approximating the attention calculation.

tiled computation · online softmax

Read the paper

query, key and valuein progress

Attention scores

Each token has a query, key and value vector. A query-key dot product gives a score for a pair of tokens. Softmax turns the scores into weights, which are used to combine the values. Select tokens to inspect the example scores.

Attention weights for the token “it”

Select a token

“it” is the query. Select a token to inspect its example attention weight. These are illustrative values after softmax.

scores → softmax → weighted sumin progress

Standard attention

First compute QKᵀ, then apply softmax to each row, then multiply the weights by V. An implementation that stores the full score and weight matrices moves O(N²) intermediate data through GPU memory. Step through the operations below.

Play the operations or select a step.

0/3

The

anima

didn’

cross

the

road

The

animal

didn’t

cross

the

road

A six-token example with bidirectional attention.

Rows · query positions
Columns · key positions
Cell · one query-key score

HBM and on-chip SRAMin progress

Storage and workspace

Think of HBM as a storeroom and on-chip SRAM as a small workbench. A standard implementation sends intermediate matrices between the two. FlashAttention processes small tiles at the workbench and writes the finished output back to storage.

Memory transfers: 0

HBM ↔ on-chip transfers

tokens

QKV

SRAM · workbench

On the workbench: S · all n² scores

HBM · storeroom

stored: Q, K, V

standard · move 1/6 · Compute the scores at the workbench. There is one entry for each query-key pair.

No transfer at this step

Play the sequence or step through each operation.

one full run ≈ 4N² + N·d entries transferred

running max, sum and outputin progress

Tiling and online softmax

FlashAttention processes blocks of Q, K and V. For each query row, it keeps a running maximum m, normalization sum ℓ, and output vector O. A larger maximum rescales the earlier contributions. This lets each score tile be discarded after its contribution is added.

tile 1/16 · stage 1/4

query block 1/4 · kv block 1/4 · stage 1: Load a query block into on-chip memory.

Each square represents a block of query-key scores. The highlighted tile is active. Faded tiles have been accumulated and discarded. The bands mark the current query and key blocks.

running state · row 5

m · biggest score so far…

ℓ · running softmax sum…

O · partial output…

The running maximum, sum and output are sufficient for online softmaxOnline softmax combines tiles using a running maximum, normalization sum and output. It gives the full softmax result without storing the whole row.. Completed score tiles are no longer needed.

compare data transfersin progress

Memory cost

Both methods compute softmax attention. FlashAttention reduces transfers to HBM by keeping intermediate tiles on-chip. Change the sequence length to compare the traffic estimates. The timing bars use a simplified model, not measured benchmarks.

Illustrative cost model · d = 128

sequence length8k

Change the sequence length to compare estimated transfers.

full score matrix / head

67.1Mentries

flash keeps on-chip

16Kat a time

hbm traffic reduction

3.7×

standard · hbm traffic138.4M elems

flash · hbm traffic37.7M elems

Estimated transfer and compute time

The model adds estimated data-transfer time to compute time. Coral shows transfersusing an assumed bandwidth of 2 TB/sAssumed HBM bandwidth: 2 TB/s, roughly an A100-class GPU.. Indigo shows computation at 312 TFLOP/sAssumed math throughput: 312 TFLOP/s of fp16 matmul, roughly an A100-class GPU.. This estimate holds compute work fixed and varies the transfer cost. Real kernels overlap work and have other overheads.

Memory transfers computing

standard193 µs

waiting 138 µs

computing 55 µs

72% of estimated time is transfers · moves 138.4M elements

flash2.1× faster93 µs

waiting 38 µs

computing 55 µs

41% of estimated time is transfers · moves 37.7M elements

FlashAttention reduces intermediate transfers while computing the same softmax attention. These timings are estimates, not measured GPU benchmarks.

Cost assumptions +

standard: 2N² + 4Ndd: the dimension of each head’s query/key/value vectors (128 here). → 138.4M elements moved

flash: 4Nd + 2N²d²/MM: the number of elements that fit in on-chip SRAM (64K here). → 37.7M elements moved

Changing operation order changes floating-point rounding. Results should agree within numerical tolerance, not necessarily bit for bit.

Original paper

FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness Dao et al. · 2022

Loading diagram…

Reference

← All guides

FlashAttention

FlashAttention computes softmax attention in tiles that fit in fast on-chip memory. It reduces GPU memory traffic by avoiding a full score matrix in high-bandwidth memory (HBM).

Overview

How it works

The attention equation stays the same. FlashAttention changes how the GPU schedules the work and stores intermediate results. It computes exact attention, up to floating-point rounding differences.

Details

The cost of intermediate matrices

For N tokens, the score matrix has N² entries per head. A standard implementation writes scores and attention weights to HBM between operations. Reading and writing those matrices can take a large share of the runtime.

Details

HBM and on-chip SRAM

HBM is the GPU’s large main memory. On-chip SRAM is much smaller and faster to access. FlashAttention chooses blocks of queries, keys and values whose temporary results fit on-chip.

Details

Processing one tile

Load a query block and a key/value block. Compute their scores and add their contribution to the output. Discard the scores, then process the next block.

Details

Online softmax

For each query row, keep the running maximum m, normalization sum ℓ, and partial output vector O. If a new tile has a larger maximum, rescale the previous contributions. These updates give the same softmax result without keeping every score.

Details

Numerical accuracy

FlashAttention does not drop tokens or approximate softmax. It changes the order of floating-point operations, so results can differ slightly from another implementation. Exact attention does not imply bit-for-bit equality.

Details

Performance limits

Tiling reduces transfers of intermediate data. Attention still computes O(N²) token pairs, so the computation remains quadratic in sequence length. Actual speed depends on the GPU, head dimensions, precision and kernel. The interactive timing bars are estimates.

Sources

FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness ↗Dao et al. · 2022