Attention
Self-attention’s naive implementation is memory-bound for a specific, avoidable reason: computing softmax(QKT/√d)V requires materializing the full N×N score matrix, writing it out to HBM, then reading it straight back for the softmax and the second matmul. None of that traffic does useful arithmetic, and an A100 has only 192KB of on-chip SRAM per streaming multiprocessor, across 108 of them, to avoid it with.
FlashAttention avoids ever writing that matrix to HBM. It loops over blocks of K and V in an outer loop, loading each block into SRAM once, then loops over blocks of Q in an inner loop, computing softmax incrementally across blocks, the same tiling idea from matrix multiplication applied to a running reduction instead of a running sum, and fusing every step of the attention computation into one kernel so intermediate results never leave the chip. The payoff is up to a 7.6× speedup on the attention computation itself for GPT-2, up to 9× fewer HBM accesses than the standard implementation, and, end to end, a 15% wall-clock speedup training BERT-large over the MLPerf 1.1 record and a 3× speedup on GPT-2.
The original kernel still reached only 25 to 40% of a GPU’s theoretical peak FLOPs/s, and FlashAttention-2 traced the shortfall to work partitioning rather than the algorithm: for some problem shapes too few thread blocks were active at once, hurting occupancy, and shared-memory traffic between warps within a block was higher than it needed to be. Rebalancing that work, splitting the sequence dimension across more thread blocks and dividing work between warps to cut shared-memory communication, without changing the underlying algorithm, roughly doubled throughput over the original kernel, reaching 50 to 73% of theoretical peak FLOPs/s on an A100 and, in end-to-end GPT-style training, up to 225 TFLOP/s per A100, 72% model FLOPs utilization.
On Hopper the ceiling moved again: FlashAttention-2, tuned for the A100, achieves only 35% utilization on an H100, and FlashAttention-3 closes the gap with the same asynchronous machinery the previous section’s GEMM kernels use. Warp specialization overlaps computation with TMA data movement, block-wise matmuls are interleaved with the softmax so the tensor cores are not idle while exponentials run on slower units, and block quantization with incoherent processing exploits Hopper’s hardware FP8 support. The result is a 1.5 to 2.0× speedup on H100, up to 740 TFLOP/s at FP16, 75% utilization, close to 1.2 PFLOP/s at FP8, and 2.6× lower numerical error than a baseline FP8 attention.
All three kernels handle the case where a full sequence is processed at once, training or a single prefill pass. Decoding, where the KV cache grows by one token per step, needs the same IO-aware framing applied to a much smaller, much more frequent matmul, which is why inference engines run a separate, differently tuned attention kernel at decode time rather than reusing the training kernel unchanged.
In serving stacks that split has a name: FlashInfer, a library and kernel generator for inference, ships separate optimized kernels for prefill, decode, and mixed batching, behind unified attention, GEMM, and mixture-of-experts APIs that select among backends including FlashAttention-2 and 3, cuDNN, CUTLASS, and TensorRT-LLM. Its attention kernels operate directly on the paged and ragged KV-cache layouts the next chapters’ engines use, add cascade attention to share the KV cache of common prefixes across requests, and stay compatible with CUDA Graphs and torch.compile for low-latency serving.