Flash Attention

Flash Attention (FlashAttention) is an algorithm and family of GPU kernels that compute exact attention for transformer models without ever writing the full attention matrix to the GPU's main memory. Standard self-attention over N tokens builds an N × N score matrix, writes it to high-bandwidth memory (HBM), and reads it back for the softmax and the product with the values, so most of attention's run time goes to moving data. Tri Dao and colleagues at Stanford introduced the method in the 2022 paper FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. It splits the work into tiles that fit in fast on-chip SRAM, so memory use grows linearly with sequence length and the output matches standard attention up to floating-point rounding. Its kernels run inside PyTorch and in large language model (LLM) inference engines such as vLLM.

How Flash Attention Works#

A GPU has a small amount of very fast memory next to its compute units and a large, slower memory off-chip. The FlashAttention paper describes an NVIDIA A100 with 40-80 GB of HBM at 1.5-2.0 TB/s and 192 KB of SRAM in each of its 108 streaming multiprocessors at an estimated 19 TB/s. Attention performs few operations per byte it reads, so its speed depends on how often data crosses between the two levels, and the method counts those transfers, which is what "IO-aware" means.

FlashAttention loads a tile of queries and then streams tiles of keys and values through SRAM. For each tile it computes the scores, updates a running maximum and running sum for every query row, rescales the partial output whenever the maximum grows, and adds the new tile's contribution. This online softmax gives the exact softmax result without seeing a whole row at once, and only the final output and two numbers per row go back to HBM. The backward pass recomputes score tiles from Q, K, and V instead of storing the N × N matrix, which adds operations and still runs faster because it reads far less memory. The same loop in PyTorch, processing all queries against one key-value tile at a time:

import torch

torch.manual_seed(0)
n, d, block = 4096, 64, 256
q, k, v = torch.randn(3, n, d).unbind(0)

# Standard attention builds the full n x n score matrix in memory
reference = torch.softmax(q @ k.T / d**0.5, dim=-1) @ v

# FlashAttention-style: stream key/value tiles, keeping a running max and sum per query (online softmax)
row_max, row_sum, out = torch.full((n, 1), -torch.inf), torch.zeros(n, 1), torch.zeros(n, d)
for start in range(0, n, block):
    scores = q @ k[start : start + block].T / d**0.5  # only an n x block tile exists at once
    new_max = torch.maximum(row_max, scores.max(dim=-1, keepdim=True).values)
    rescale, probs = torch.exp(row_max - new_max), torch.exp(scores - new_max)
    row_sum = row_sum * rescale + probs.sum(dim=-1, keepdim=True)
    out = out * rescale + probs @ v[start : start + block]
    row_max = new_max
out = out / row_sum

print(f"Largest score tile: {tuple(scores.shape)} vs full matrix: ({n}, {n})")
print(f"Max difference from standard attention: {(out - reference).abs().max().item():.1e}")

The largest score tile is 4,096 × 256, one-sixteenth of the 4,096 × 4,096 matrix, and the result differs from standard attention by at most 4.3e-07, the size of FP32 rounding error. A real kernel also splits the queries into tiles, one per GPU thread block, and fuses the whole loop into a single kernel so the tiles never leave SRAM.

FlashAttention-2, -3, and -4#

Each version targets a newer GPU generation. FlashAttention-2 (2023) reorganized how work is split across threads and ran about twice as fast as the original, reaching 50-73% of an A100's peak throughput. FlashAttention-3 (2024) uses the asynchronous Tensor Cores of the H100 and supports FP8, reaching 740 TFLOPs/s in FP16 and close to 1.2 PFLOPs/s in FP8. FlashAttention-4 (2026) targets Blackwell GPUs and reports up to 1,613 TFLOPs/s in BF16 on a B200, 71% of peak. The kernels ship in the open-source flash-attention library, which works on FP16 and BF16 inputs, and PyTorch's scaled_dot_product_attention selects a FlashAttention-2 kernel automatically when the GPU and inputs support it, falling back to memory-efficient or plain implementations otherwise. Meta's SAM 2, which runs in the Ultralytics package, computes attention in its Hiera image encoder and its memory attention through that function, so PyTorch routes those layers to FlashAttention whenever the hardware and input precision allow.

Flash Attention vs. Related Terms#

• Standard attention: The same mathematical operation, computed by writing the full score and probability matrices to HBM.
• Sparse attention: Scores only selected query-key pairs. The FlashAttention paper also introduced a block-sparse variant that skips empty tiles.
• Linear attention: Replaces the softmax with a feature map so cost grows linearly with length, which changes the model's outputs.
• PagedAttention: A memory manager that stores the KV cache in fixed-size pages during serving. Attention kernels, FlashAttention included, read keys and values from those pages.
• Ring attention: Splits a very long sequence across devices, each computing blockwise attention on its shard while key-value blocks pass around a ring.

Impact on Training and Long Contexts#

Because memory grows linearly with sequence length, FlashAttention made long context windows practical to train. The original paper reported a 15% end-to-end speedup for BERT-large at sequence length 512 over the MLPerf 1.1 training record, a 3× speedup for GPT-2 at sequence length 1,024, and 2.4× on the Long Range Arena benchmark. The same savings apply in vision, where a Vision Transformer (ViT) or diffusion transformer turns high-resolution images and video into long token sequences. Longer sequences also improved accuracy: the paper reported the first transformers to beat chance on Path-X, a task that asks whether two points in a 128 × 128 image, flattened into a sequence of 16,384 pixels, are connected by a path.