FlashAttention

FlashAttention is a technique that makes AI language models faster and less memory-hungry, without changing their answers. It reorganises one heavy calculation, attention, so that the graphics card fetches data from its slow main memory far less often. That lets longer documents fit on the same hardware and speeds up training. It is built into PyTorch, a common framework for AI models.

What problem does it solve?

Attention compares every token in a text with every other token, so its time and memory grow with the square of the text length. A standard implementation writes the whole table of those comparisons to the GPU's main memory (HBM) and reads it back. On an Nvidia A100, that memory holds 40 to 80 GB at 1.5 to 2.0 TB/s; the on-chip SRAM holds only 192 KB per processor block, but is about an order of magnitude faster (Dao et al., 2022). According to the paper, most operations in a transformer are held up by memory access.

How does it work?

FlashAttention, by Tri Dao and colleagues (May 2022), is IO-aware: it treats reads and writes between the GPU's memory levels as the real cost.

• Tiling: it splits the queries, keys and values into blocks that fit in SRAM, computes attention block by block and builds up the softmax step by step, so the full table never goes to main memory.
• Recomputation: for training, it stores only small normalisation values and recomputes the attention blocks during the backward pass, which is faster than reading a stored table.
• One kernel: all attention steps run in a single GPU program.

The result is exact attention, with extra memory that grows linearly with the text length. The number of calculations still grows with the square. In the paper, FlashAttention trained GPT-2 up to three times as fast as the Hugging Face implementation and 1.8 times as fast as Megatron, BERT-large 15% faster than the MLPerf 1.1 record, and used up to 20 times less memory than exact attention baselines.

Which versions exist?

• FlashAttention-2 (July 2023) divides the work better over the GPU: about twice as fast, at 50 to 73% of the A100's theoretical maximum. PyTorch 2.2 (January 2024) built it into its standard attention function.
• FlashAttention-3 (July 2024) targets Nvidia's H100: 1.5 to 2 times faster than version 2 in the forward pass, up to 740 TFLOPs/s in FP16.
• FlashAttention-4 (March 2026) targets Blackwell GPUs such as the B200: up to 1,613 TFLOPs/s in BF16.

The repository also offers FlashAttention-2 for AMD GPUs through ROCm. Fine-tuning tools such as LLaMA-Factory can switch it on with a setting.

Example

In March 2026, the developers of the vLLM Semantic Router, which classifies requests before they reach a language model, described the problem on an AMD MI300X (Liu et al.). With standard attention, one classifier needed about 1.5 GB for the attention table at 8K tokens, about 6 GB at 16K and about 24 GB at 32K: every doubling of the text quadrupled the memory. Next to the language model on the same GPU, that did not fit. A FlashAttention operator made attention memory grow linearly and made texts of 8K to 32K tokens possible.