# Attention · Kernel optimization

<!-- https://learn-kernels.com/chapters/kernel-optimization/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** (A GPU's main off-chip memory. An A100 has 1.5 to 2.0 TB/s of HBM bandwidth, an order of magnitude less than the roughly 19 TB/s available from the on-chip SRAM local to each streaming multiprocessor.), 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.

> Figure. IO-aware attention. One block of K and V is loaded into on-chip SRAM at a time while blocks of Q stream past it. The score block, the online softmax, and the multiply by V are fused into a single kernel, so the N×N score matrix is never written to HBM; a running max and sum rescale the partial output as each new block arrives. Shaded blocks mark one tile's worth of data.Illustrative numbers

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** (The fraction of the warps an SM could run at once that are actually resident.), 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](https://learn-kernels.com/chapters/foundations/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.

Source

IO-aware tiling, fusion, and the BERT-large/GPT-2/long-range-arena speedups from [FlashAttention](https://arxiv.org/abs/2205.14135). Work-partitioning fixes and the A100 utilization and end-to-end training numbers from [FlashAttention-2](https://arxiv.org/abs/2307.08691). Hopper asynchrony techniques, utilization, and FP8 error figures from [FlashAttention-3](https://arxiv.org/abs/2407.08608). FlashInfer’s kernels, backends, and KV-cache support from the [FlashInfer](https://github.com/flashinfer-ai/flashinfer) README. More on fused and distributed attention kernels in [the reading list](https://learn-kernels.com/chapters/reading#attention).
