FlashAttention

FlashAttention is an I/O-aware GPU kernel for scaled dot-product attention . It computes exactly the same output as a textbook attention implementation — softmax(QK^T / sqrt(d_k)) V — but reorganizes the operations to never write the N x N attention matrix to GPU main memory. By tiling the inputs into blocks that fit in fast on-chip SRAM and fusing the softmax across blocks, it cuts attention’s memory traffic by an order of magnitude and unlocks long-context inference at scale.

The bottleneck it solves

Naive attention has three steps: compute , apply softmax row-wise to get , then multiply to produce the output . The intermediate matrix (and ) is — quadratic in sequence length. At and fp16, that’s 128 MB per head per layer, written to HBM (slow GPU global memory) and read back twice.

The compute is fast; the memory traffic is the wall. Modern GPUs offer ~200 FLOPS per byte loaded from HBM. Naive attention nowhere near saturates the FLOPS — it spends most of its time waiting for memory.

The trick: tiling and online softmax

FlashAttention partitions , , into blocks small enough to fit in SRAM. For each pair of (Q-block, K-block), it computes a partial block in SRAM, runs the relevant portion of the softmax against the running normalizer, multiplies by the V-block, and accumulates the result. The full attention matrix never materializes in HBM — only the input/output tensors do.

The technically interesting part is the softmax. Softmax requires the row’s full sum of exponentials in the denominator, which seems to require seeing the whole row at once. FlashAttention uses an online softmax trick: process blocks left-to-right, keeping a running max and running sum . When a new block arrives with a larger max, rescale the prior accumulators by the exponential of the difference. This produces the exact softmax output, computed incrementally, with total memory.

What it actually delivers

On A100 GPUs, FlashAttention 2 achieves around 70% of the theoretical peak FLOPS for attention — versus around 30-40% for naive PyTorch attention, which is bandwidth-bound. End-to-end, this translates to 2-4x faster training and 2-3x faster inference for transformer models, with the gains larger at longer contexts.

The bigger story is what it makes possible. Long-context training (32K, 128K, 1M tokens) is essentially impossible without FlashAttention or a close variant — the memory and bandwidth costs would be untenable. Every modern long-context LLM is trained and served with some flavor of FlashAttention or its descendants (xFormers, PyTorch SDPA, the various TensorRT-LLM kernels).

How it composes with the rest of the stack

FlashAttention stacks cleanly with grouped-query-attention (the kernel doesn’t care how many K/V heads there are), with kv-cache (the prefill phase uses FlashAttention; the decode phase uses a specialized variant called Flash-Decoding), and with quantization. Most production inference engines — vLLM, TensorRT-LLM, SGLang — use FlashAttention as the default attention kernel and only fall back to naive attention when the input doesn’t fit the kernel’s constraints.

FlashAttention is the canonical example of how modern ML performance is dominated by memory hierarchy rather than raw compute. The math of attention hasn’t changed since 2017; reordering when each tensor lives in which memory level was worth a 2-4x speedup across the entire field.