FlashAttention

The core innovation is tiling the attention computation to fit into the limited, high-bandwidth GPU SRAM. Instead of materializing the full N x N attention matrix in slow HBM, FlashAttention loads blocks of Query (Q), Key (K), and Value (V) matrices from HBM to SRAM, computes attention on the block, and writes only the final output back. This drastically reduces HBM read/write operations, the primary bottleneck in standard attention.

Unlike sparse attention or low-rank approximation methods, FlashAttention computes mathematically identical results to standard attention. It uses a numerically stable online softmax algorithm, known as tiled softmax or online normalizer calculation, to correctly rescale block statistics without needing the full row of attention scores. This guarantees no loss in model quality while achieving significant speedups.

To remain memory-efficient during training, FlashAttention avoids storing the large intermediate attention matrix for the backward pass. Instead, it recomputes the attention scores and softmax statistics during the backward pass by reloading the Q, K, V blocks from HBM. This recomputation strategy trades extra FLOPs for a massive reduction in memory footprint, enabling training with longer sequences.

FlashAttention implements kernel fusion by combining multiple CUDA kernels into a single, highly optimized kernel. Operations like the matrix multiply, softmax scaling, dropout masking, and the final output accumulation are all performed in a single pass over the data while it resides in SRAM. This eliminates the overhead of launching multiple kernels and the associated round-trips to HBM.

Standard attention has a memory complexity of O(N²) with sequence length N, which quickly exhausts GPU HBM for long sequences. FlashAttention reduces the memory footprint to O(N) by never materializing the full attention matrix. This allows training Transformers on sequence lengths of 8k, 16k, or even 64k tokens, which was previously infeasible without model parallelism or approximation.

By shifting the bottleneck from memory bandwidth to compute speed, FlashAttention delivers 2-4x wall-clock speedups for the attention layer. This translates to significant end-to-end training time reductions for models like GPT and BERT. For inference, FlashAttention accelerates the prefill phase and, in its FlashDecoding variant, optimizes the generation phase by parallelizing the attention computation across the sequence length dimension.

The core design principle behind FlashAttention. Traditional attention is memory-bound, not compute-bound. The algorithm minimizes reads/writes between slow High Bandwidth Memory (HBM) and fast Static Random-Access Memory (SRAM) on the GPU.

• Tiling: Splits the Q, K, V matrices into blocks that fit entirely in SRAM.
• Recomputation: Recalculates the attention matrix in the backward pass instead of storing it, trading FLOPs for memory savings.
• Kernel Fusion: Combines multiple operations (matmul, softmax, dropout) into a single CUDA kernel.

Understanding the hardware FlashAttention optimizes for is essential. Modern GPUs have a deep memory hierarchy with vastly different bandwidths.

• HBM (High Bandwidth Memory): Large capacity (80GB on A100) but high latency. Standard attention stores the full N x N attention matrix here.
• SRAM (Static RAM): Small (192KB per streaming multiprocessor on A100) but ~10x faster than HBM. FlashAttention keeps all working data here.
• Bandwidth Gap: The key insight is that attention is bottlenecked by HBM bandwidth, not compute throughput.

FlashAttention is critical for optimizing the Key-Value (KV) cache during autoregressive decoding. The KV cache stores previous keys and values to avoid recomputation, but its size grows linearly with sequence length.

• Memory Pressure: Long sequences cause the KV cache to exceed HBM capacity, making memory bandwidth the primary bottleneck.
• FlashDecoding: An extension of FlashAttention that parallelizes the attention computation over the long KV sequence dimension during inference.
• Works synergistically with Multi-Query Attention (MQA) and Grouped-Query Attention (GQA) to further reduce cache size.