FlashAttention

FlashAttention is explicitly designed around the memory hierarchy of modern GPUs. It treats high-bandwidth memory (HBM) as a slow, large storage and SRAM as a fast, small cache. The algorithm tiles the large attention matrix into blocks, performing the entire softmax operation for a block within SRAM. This minimizes the number of expensive reads/writes to HBM, making the computation memory-bandwidth optimal. The core insight is that recomputation (a compute cost) is cheaper than frequent HBM I/O (a memory cost) for attention operations.

Standard attention algorithms must materialize the intermediate S = QK^T matrix, which scales quadratically (O(N²)) with sequence length (N). For a 64K sequence, this matrix requires ~16GB of memory in FP16. FlashAttention never stores this full matrix to HBM. By fusing the attention scoring, softmax, and weighted sum operations, it only writes the final output. This reduces memory usage from quadratic to linear in sequence length, enabling the training and inference of models with extremely long contexts (e.g., 64K, 128K tokens) on the same hardware.

By minimizing HBM I/O, FlashAttention provides substantial wall-clock speed improvements. Benchmarks show:

• Up to 3x faster training for GPT-style models.
• Up to 10-20x faster inference for long-sequence generation. The speedup is most pronounced for long context lengths where standard attention becomes severely memory-bound. This directly translates to lower cloud compute costs and faster iteration cycles for research and production deployment. The algorithm is exact, so these speedups come with no loss in model accuracy.

FlashAttention's memory efficiency directly enabled a new generation of long-context LLMs. Before its development, training models on sequences longer than 2K tokens was prohibitively expensive. By eliminating the quadratic memory bottleneck, FlashAttention made it feasible to train models on 8K, 32K, and even 128K token sequences. This capability is critical for applications requiring deep document analysis, long-form content generation, and extended multi-turn conversations, forming the technical backbone for models like GPT-4 and Claude.

The core tiling and recomputation principles of FlashAttention have been extended to create more specialized, high-performance variants:

• FlashAttention-2: A more optimized implementation with better work partitioning, achieving near-peak GPU utilization.
• FlashAttention-3: Leverages new hardware features like NVIDIA's FP8 Tensor Cores and asynchronous copy instructions for H100 GPUs.
• FlashDecoding: Optimizes the decode-phase attention for batch size 1, crucial for low-latency inference.
• FlashDecoding++: Further enhances decode-phase performance for very long sequences. These variants ensure the algorithm family continues to push hardware limits.

While FlashAttention optimizes the prefill phase, its principles are complementary to KV cache optimizations for the decode phase. Efficient attention computation reduces the initial cost of populating the KV cache. Furthermore, techniques like PagedAttention (used in vLLM) manage the storage of the KV cache in GPU memory, while FlashAttention optimizes the computation that creates and uses that cache. Together, they form a complete optimization stack for transformer inference, addressing both computational and memory bottlenecks across the entire generation lifecycle.

FlashAttention's primary innovation is its IO-aware design. It splits the large Softmax(QK^T)V operation into smaller blocks or tiles that fit in fast SRAM. Instead of writing the massive attention matrix (size N x N) to slow HBM, it recomputes parts of the attention scores on-the-fly during the backward pass. This kernel fusion avoids the expensive HBM reads/writes that are the true bottleneck, trading extra FLOPs for vastly reduced memory I/O.

• Tiling: Loads queries, keys, and values in blocks from HBM to SRAM.
• Recomputation: Recalculates attention scores in SRAM during backpropagation.
• Result: Achieves exact attention with 2-4x faster training and 10-20x faster inference on long sequences.

The algorithm is explicitly designed for modern GPU memory architecture. It treats SRAM (on-chip, ~20TB/s bandwidth) as a fast cache and HBM (off-chip, ~1TB/s bandwidth) as slow main memory. By keeping the working set of tiles in SRAM, it minimizes the number of costly HBM accesses.

• Bandwidth Reduction: Reduces HBM reads/writes from O(N²) to O(N²d²/M), where M is SRAM size.
• Memory Footprint: Requires only O(N) memory instead of O(N²) to store the attention matrix.
• Practical Impact: Enables training transformers with 16k+ context length on a single GPU, where standard attention would run out of memory.

FlashAttention-2 refines the original algorithm with better parallelism and work partitioning, achieving near-perfect utilization of NVIDIA GPU tensor cores. It rearranges the loop order to parallelize over sequence length and better saturate warps.

• Improved Parallelism: Reduces non-matrix-multiply operations, better aligning with GPU warp schedulers.
• Kernel Variants: Includes FlashDecoding for optimized inference (uneven keys/queries) and support for multi-query and grouped-query attention.
• Hardware Support: Optimized kernels exist for Hopper (FP8), Ampere, and AMD MI300 architectures.

FlashAttention directly enables new research and production capabilities by removing memory constraints.

• Long Context Models: Made training GPT-3 (8k context), Claude (100k context), and other long-context LLMs feasible.
• Higher Batch Sizes: Increases training throughput by allowing larger batches within the same GPU memory.
• Inference Speedup: Critical for reducing latency in autoregressive decoding, where attention over the KV cache is memory-bandwidth bound.
• Multimodal Models: Enables efficient attention over long sequences of image patches in vision transformers.

FlashAttention is part of a broader ecosystem of attention and memory optimizations.

• PagedAttention (vLLM): Manages the KV cache in non-contiguous blocks for efficient memory sharing in inference serving.
• xFormers: A library of optimized transformer building blocks, including memory-efficient attention.
• Causal Masking: FlashAttention includes highly optimized kernels for the causal masks used in language modeling.
• Alibi & Rotary Embeddings: Supports common positional encoding schemes like RoPE without performance degradation.
• Distributed Training: Supports tensor parallelism and sequence parallelism for scaling to extremely long sequences across multiple GPUs.