Advanced · Systems

FlashAttention

Visual diagram · (in preparation) · Math · (in preparation) · Worked example · 3 difficulty levels.

TL;DR. An IO-aware exact attention algorithm that tiles Q, K, V into SRAM blocks, cutting memory from O(n²) to O(n) and giving 2–4× speedups.

Technical Definition

An IO-aware exact attention algorithm that tiles Q, K, V into SRAM blocks, cutting memory from O(n²) to O(n) and giving 2–4× speedups.

How it works

FlashAttention (Dao et al., 2022) fuses the attention computation into a single GPU kernel that streams tiles of the input through fast on-chip SRAM and never materializes the full n×n attention matrix in HBM. It computes the exact softmax via online normalization. V2 and V3 push throughput further with better work partitioning and FP8 support, and it is now the default attention kernel in most LLM training stacks.

Related Concepts

• Transformer — An architecture that uses self-attention to process sequences in parallel, powering modern language models like GPT and BERT.
• Attention Mechanism — A technique that lets models dynamically focus on the most relevant parts of the input when producing each output element.
• KV-Cache — A memory optimization storing previously computed key-value pairs during autoregressive generation to avoid redundant computation.
• GPU — Graphics Processing Unit — a massively parallel processor that powers most modern AI workloads.