The attention mechanism at the heart of transformer models has a memory problem. As sequence lengths grow, the standard attention algorithm becomes bottlenecked not by compute capacity, but by memory bandwidth. Flash Attention addresses this by rethinking how attention computations interact with GPU memory hierarchy.

The Standard Attention Bottleneck

Standard attention computes a compatibility score between all pairs of tokens in a sequence, materializing an N×N matrix where N is sequence length. For a 2048-token sequence, that’s over 4 million scores stored in GPU high-bandwidth memory (HBM). The algorithm reads this matrix multiple times during forward and backward passes, moving data between HBM and on-chip SRAM repeatedly.

Modern GPUs have abundant compute throughput but limited memory bandwidth. An A100 can perform 312 teraflops of operations per second, but memory bandwidth tops out around 2 TB/s. When an algorithm is memory-bound rather than compute-bound, adding more FLOPS doesn’t help. The GPU sits idle waiting for data transfers.

Standard attention is memory-bound. It performs O(N²) operations but requires O(N²) memory reads and writes. The compute-to-memory-access ratio is roughly 1:1, far below what’s needed to saturate modern GPU compute units.

Tiling and Kernel Fusion

Flash Attention exploits the GPU memory hierarchy through two techniques: tiling and kernel fusion. Instead of materializing the full attention matrix in HBM, it breaks the computation into blocks that fit in fast on-chip SRAM.

SRAM on modern GPUs is small—around 20 MB on an A100—but operates at much higher bandwidth than HBM. By loading blocks of queries, keys, and values into SRAM, computing attention for those blocks, and immediately using the results, Flash Attention reduces HBM traffic dramatically.

Kernel fusion combines multiple operations into a single GPU kernel. Standard attention implementations split the computation across separate kernels: one for matrix multiplication, another for softmax, another for dropout. Each kernel launch writes intermediate results back to HBM and reads them again for the next step. Flash Attention fuses these operations, keeping intermediate values in registers and SRAM.

Recomputation in the Backward Pass

The backward pass during training presents another challenge. Standard implementations save the full attention matrix from the forward pass to use during backpropagation, consuming additional memory. Flash Attention recomputes attention scores during the backward pass instead of storing them.

This trades compute for memory. Recomputation seems inefficient, but when memory bandwidth is the bottleneck, performing extra compute while data is already loaded into SRAM costs little. The overall speedup comes from reduced memory traffic, not reduced operations.

Online Softmax

Computing softmax in tiles requires a careful algorithmic adjustment. Softmax depends on global statistics—the maximum value and sum across the entire sequence—but tiling means processing the sequence in chunks.

Flash Attention uses an online softmax algorithm that maintains running statistics as it processes each block. Instead of computing softmax in two passes (find max, then normalize), it updates normalization factors incrementally. This allows correct softmax computation without materializing the full attention matrix.

Practical Impact

The memory efficiency gains are substantial. Flash Attention reduces memory usage from O(N²) to O(N) in sequence length, enabling longer contexts within fixed memory budgets. A model that maxes out at 1024 tokens with standard attention might handle 4096 tokens with Flash Attention on the same hardware.

Speed improvements follow from better hardware utilization. By keeping the GPU compute units fed with data from fast SRAM rather than slow HBM, Flash Attention achieves 2-4x speedups on typical workloads. The gains increase with sequence length, precisely where standard attention struggles most.

This architectural approach—analyzing where compute and memory bandwidth interact, then restructuring algorithms to match hardware characteristics—represents a broader pattern in systems optimization. The fastest code isn’t always the cleverest algorithm on paper. It’s the algorithm that understands the machine.