Most discussions of transformer cost start with arithmetic: attention scales quadratically with sequence length, so doubling the context quadruples the work. That is true, but it is not the reason attention is slow on real hardware. On a modern accelerator the multiplications are comparatively cheap. What costs time is moving the numbers between the GPU’s main memory and the small, fast memory that sits next to the arithmetic units.
FlashAttention, introduced in 2022 by Dao et al. as FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness, is built around that observation. It does not approximate attention, drop tokens or sparsify the score matrix. It reorganises the computation so that the large intermediate matrices are never written to main memory at all: the same output, far less data movement. Variants of it are now the default attention path in mainstream training and serving stacks.
Attention is memory-bound, not compute-bound
Consider what a single attention head does for a sequence of length N. It computes an N × N matrix of scores, applies a softmax across each row, then multiplies the resulting probabilities by the values. A naive implementation materialises two full N × N matrices — the raw scores and the post-softmax probabilities — in memory.
The arithmetic is worth doing yourself. At a sequence length of 8,192 in 16-bit precision, one such matrix is 8,192 × 8,192 × 2 bytes ≈ 134 MB — for a single head. With 32 heads that is about 4 GB of intermediate data for one attention layer in one forward pass, all of it written out and read back, with the softmax step touching it twice.
Meanwhile the matrix multiplications are exactly the dense, regular work that tensor cores handle well. Arithmetic throughput has grown faster than memory bandwidth, so the kernel spends most of its wall-clock time waiting for transfers. Adding FLOPs to such a kernel costs almost nothing; removing memory traffic is what makes it faster.
Two levels of memory, one bottleneck
The hardware picture is simple. High-bandwidth memory (HBM) is the large pool quoted on the spec sheet, tens of gigabytes. On-chip SRAM — the shared memory and registers attached to each streaming multiprocessor — is tiny by comparison: the FlashAttention paper puts an A100 at 40 GB of HBM against roughly 20 MB of on-chip SRAM in total, with the on-chip memory roughly an order of magnitude faster.
An N × N score matrix at long context does not fit in 20 MB, so a standard implementation has no choice but to stage it through HBM, and every round trip is a stall. The design question becomes: can we produce the same result while keeping every intermediate inside the fast memory?
Three ideas that make it work
- Tiling. Split the queries, keys and values into blocks small enough that a tile of the score matrix fits in SRAM. Load a block of queries, stream blocks of keys and values past it, and compute the scores, their softmax contribution and the output accumulation inside fast memory. Q, K and V are read once; the output is written once.
- Online softmax. Softmax normally needs the whole row: subtract the row maximum for numerical stability, then divide by the sum of exponentials. Tiling means you only ever see part of a row. The fix is to carry a running maximum and a running sum alongside the partial output, rescaling the accumulated result whenever a new tile raises the maximum. The algebra works out exactly.
- Recomputation. The backward pass needs the attention probabilities to compute gradients. Rather than keeping them from the forward pass — which would reintroduce the memory cost — FlashAttention recomputes them from Q and K using the saved softmax statistics. That is more arithmetic in exchange for less traffic, the right trade when arithmetic is the abundant resource.
Put together, the kernel is an outer loop over query blocks and an inner loop over key and value blocks, with the running statistics and the partial output held in fast memory. The output is exact — the same attention, not a cheaper approximation of it — which is what makes it safe to apply to an already-trained model.
What changed in versions 2 and 3
The first version parallelised across the batch and the attention heads, which leaves GPUs underused when the batch is small and the sequence is long — precisely the long-context regime people care about. FlashAttention-2 (Dao, 2023) added parallelism over the sequence dimension, reduced the non-matmul work around the softmax and improved how work is partitioned between warps. The paper reports roughly 2× over the original, with attention reaching 50–73% of the theoretical maximum matmul throughput on an A100.
FlashAttention-3 (Shah et al., 2024) targets Hopper-class hardware, exploiting asynchronous tensor-core instructions and the tensor memory accelerator through warp specialisation: some warps do nothing but fetch the next tile while others compute on the current one, forming a producer–consumer pipeline. The paper reports 1.5–2.0× over FlashAttention-2 on H100 in FP16, and higher throughput again in FP8.
Each generation extracts more overlap between data movement and arithmetic, and each is tied more tightly to a particular GPU architecture. Related work such as ring attention extends the same tiling logic across devices, splitting one sequence over several GPUs.
Turning it on, and when you do not get it
In practice this is rarely a code-writing exercise. PyTorch dispatches scaled_dot_product_attention to a fused backend when the inputs and hardware allow it, and Hugging Face Transformers exposes an attention implementation argument when a model is loaded. The work is in the conditions, not the call:
- Hardware generation. The flash-attn project requires Ampere or newer for FlashAttention-2, and its version 3 kernels target Hopper. Older cards fall back to a slower path.
- Dtype and shape constraints. Fused kernels typically require half or bfloat16 inputs and particular head dimensions. Step outside them and you silently get the reference implementation.
- Masking and bias. Causal and padding masks are supported; arbitrary additive attention biases often are not, and a custom mask is a common reason a model quietly loses the fast path.
- Verify, do not assume. Profile the kernel that actually ran. A change upstream can disable fusion without any error message.
It also composes with the rest of the stack. Memory saved on activations is memory available for the KV cache, which is why these kernels and the paged cache described in serving LLMs on your own GPUs are usually deployed together, and why hybrid attention designs assume an efficient exact kernel as their baseline.
What this means for organisations training on their own hardware
The lesson generalises beyond attention. When you budget a training run or a fine-tune, the first question is not how many FLOPs the model needs but where the bytes are going; a kernel waiting on HBM will not go faster on a card with more tensor cores. Before reaching for architectural changes that alter model quality, confirm that the exact, well-optimised kernels are actually being used — throughput with no accuracy trade-off, and often the difference between a long-context experiment being affordable or not.
It also changes how you read hardware requirements. Because these kernels are tied to specific GPU generations and precisions, the gap between an older and a newer card is wider than the headline specifications suggest. AMD Telecom trains and fine-tunes deep-learning models and large language models on its own in-house GPUs.
If you are planning a training run or assessing what your infrastructure can realistically support, our AI Model Training and AI Model Development services are a good place to start — contact us to discuss your requirements.