FlashAttention

In short: FlashAttention makes transformer attention fast not by changing the math, but by restructuring the computation so it reads and writes far less to slow high-bandwidth memory (HBM), tiling the work into on-chip SRAM and fusing every step into one GPU kernel.

If you have ever profiled a transformer and watched the GPU sit at 30% utilization while attention dominates the runtime, you have met the problem FlashAttention was built to solve. The surprise, when Tri Dao and colleagues published it in 2022, was that standard attention is not compute-bound at all on modern accelerators. It is memory-bound. The arithmetic is cheap; shuttling intermediate tensors in and out of GPU memory is what kills you.

Why attention is memory-bound, not compute-bound

The scaled dot-product attention for a sequence of length N and head dimension d computes S = QKᵀ, then P = softmax(S), then O = PV. The intermediate matrices S and P are each N×N. For a sequence of 4,096 tokens that is roughly 16.7 million elements — per head, per layer. A naive implementation materializes all of S and P in HBM, writes them, reads them back for the softmax, writes again, reads once more for the PV product.

The costs are wildly asymmetric. On an NVIDIA A100, on-chip SRAM runs around 19 TB/s while HBM delivers roughly 1.5–2 TB/s — an order of magnitude gap. The GPU can do the FLOPs in a blink; it spends most of its time waiting on round trips to HBM.

⚡ Standard attention's HBM traffic scales as O(N²). FlashAttention cuts the extra reads/writes of the N×N score matrix entirely, so HBM accesses scale as O(N²d²/M) where M is the SRAM size — in practice a 5–20× reduction in memory traffic and 2–4× wall-clock speedups.

The core idea: tiling and kernel fusion

FlashAttention never builds the full N×N score matrix. Instead it splits Q, K, and V into blocks (tiles) sized to fit in SRAM, loads one tile set at a time, computes a partial attention result entirely on-chip, and fuses the matmul, softmax, and output accumulation into a single CUDA kernel. There is no intermediate write of S or P back to HBM at all.

Concretely: for each block of queries, the kernel iterates over blocks of keys/values. For each key block it computes the local scores, updates a running softmax, and accumulates into the output — all in registers and shared memory. Only the final O tile (size N×d, not N×N) is written to HBM.

The clever bit: online (streaming) softmax

Softmax is the obstacle to tiling. A normal softmax needs the global maximum and the global sum of exponentials over the whole row before it can normalize — but tiling only ever sees one block of scores at a time. FlashAttention borrows the online softmax trick: it keeps a running max m and running denominator ℓ, and when a new block arrives it rescales the previously accumulated output by exp(m_old − m_new) before adding the new contribution.

This makes the result numerically identical to the full softmax (it is exact, not an approximation) while requiring only O(N) extra memory for the statistics rather than O(N²) for the probability matrix. The backward pass uses the same idea, recomputing the needed blocks on the fly rather than storing them — trading a little extra compute for a large memory saving, which is exactly the right trade on a memory-bound kernel.

Numbers that matter in practice

? On GPT-2-scale training, the original paper reported end-to-end speedups around 3× and the ability to push context lengths to 16K and beyond on the same hardware, because peak memory grows linearly in N rather than quadratically.

FlashAttention-2 (2023) improved parallelism and work partitioning across GPU warps, and FlashAttention-3 (2024) exploited the asynchrony and FP8 support of NVIDIA Hopper (H100) GPUs, reaching a large fraction of peak throughput. The technique is now the default attention backend in PyTorch's scaled_dot_product_attention, vLLM, and most production inference stacks.

Why hardware engineers care

FlashAttention is a textbook case of the lesson every VLSI and systems engineer learns eventually: the roofline model rules everything. When your kernel sits below the memory-bandwidth ceiling, buying more FLOPs does nothing — you have to move less data. The same instinct drives dataflow choices in AI accelerators, where SRAM scratchpad sizing, operand reuse, and on-chip buffering decide whether a design hits its throughput target. If you are designing or verifying such accelerators, tools like Synopsys and Cadence high-level synthesis and memory-compiler flows, and Siemens' verification suite, are where these bandwidth-vs-area trade-offs get pinned down in silicon.

FlashAttention also pairs naturally with inference-side memory tricks. For autoregressive decoding, pairing it with KV-cache optimization is what makes long-context serving affordable, and both sit alongside quantization in the broader toolkit of transformer inference optimization. If you want the silicon-design side of this story, our partners at Avecas VLSI design services work on exactly these accelerator memory subsystems.

Want to build intuition for IO-aware kernels, the roofline model, and modern accelerator design? Explore the AI Hardware track on our CourseTron courses page.

Frequently asked questions

Does FlashAttention change the attention output?

No. It is mathematically exact — the online softmax produces bit-comparable results (within normal floating-point tolerance) to a naive implementation. It only changes how the computation is scheduled across memory, not what is computed.

Is FlashAttention only useful for long sequences?

It helps most at longer sequence lengths where the N² score matrix dominates, but even at moderate lengths the kernel fusion removes several HBM round trips, so you still see gains. The benefit grows as context length increases.

Do I need to write CUDA to use it?

Rarely. It is already wired into PyTorch, vLLM, Hugging Face Transformers, and most frameworks. You typically just enable the FlashAttention backend; writing the kernel yourself is only needed for research or custom hardware.

How is this different from a KV cache?

A KV cache avoids recomputing keys and values for past tokens during decoding; FlashAttention changes how a single attention computation moves data through the memory hierarchy. They solve different bottlenecks and are used together.