FlashAttention

An I/O-aware, exact attention algorithm that minimizes memory reads/writes to the GPU, making training and inference of Transformers significantly faster and more memory-efficient.

How it works

Standard attention materialises a full N×N attention score matrix in GPU HBM (slow off-chip memory). FlashAttention tiles the attention computation into blocks that fit in the fast on-chip SRAM, computing and discarding partial results without ever writing the full N×N matrix to HBM. The final output and softmax statistics are accumulated on-chip. This reduces HBM reads/writes from O(N²) to O(N), resulting in 2-4x faster attention and O(N) rather than O(N²) memory for the attention step.

Why it matters

FlashAttention is one of the most impactful systems-level optimisations in modern LLM development. Its memory savings are what have made training with 100K+ token contexts practical. It has been integrated into PyTorch, JAX, and virtually every major training and inference framework. FlashAttention 2 and 3 have continued to push throughput higher. It demonstrates that careful systems optimisation — not just algorithmic innovation — can deliver order-of-magnitude improvements in practical AI performance.