FlashAttention

FlashAttention is a revolutionary algorithm for computing the attention mechanism, designed to significantly accelerate the training and inference of large language models (LLMs) while maintaining full computational accuracy. The algorithm was first introduced in 2022 by a team of researchers from Stanford University led by Tri Dao.

The key idea behind FlashAttention is to reorganize computations with an awareness of the GPU memory hierarchy, which minimizes the number of accesses to slow memory and eliminates the main bottleneck of the standard attention mechanism.

The Problem with Standard Attention

The standard self-attention mechanism in Transformers is calculated using the formula: $\text{Attention}(Q,K,V)=\text{softmax}\left(\frac{Q{K}^T}{\sqrt{d_k}}\right)V$ where Q, K, and V are the query, key, and value matrices.

The main problem with this approach is its quadratic complexity in time and memory (O(N²)) with respect to the sequence length N. A naive implementation requires computing and storing the full N×N attention matrix S in GPU memory, which leads to two critical problems:

1. High memory consumption: Storing the N×N matrix becomes infeasible when working with long contexts.
2. Input/Output (IO) Operations: The primary bottleneck is not the number of arithmetic operations, but the constant accesses to slow GPU memory.

GPU Memory Hierarchy

To understand the problem, it is important to distinguish between two types of memory in a GPU (using the NVIDIA A100 as an example):

• SRAM (Static RAM): Fast on-chip memory of small capacity (~20 MB) with enormous bandwidth (up to 19 TB/s).
• HBM (High Bandwidth Memory): Slower, large-capacity memory (40–80 GB) with much lower bandwidth (around 1.5 TB/s).

This asymmetry makes the standard attention algorithm memory-bound, as it constantly reads and writes large matrices from the slow HBM, which is the main source of latency.

Key Innovations of FlashAttention

FlashAttention is an IO-aware algorithm that solves the problem by minimizing accesses to HBM. This is achieved through three main techniques.

Tiling and Block Processing

Instead of processing the entire matrix at once, FlashAttention divides the input matrices Q, K, and V into small blocks (tiles) that fit into the fast SRAM. The algorithm sequentially loads these blocks, performs all attention computations for them, and updates the final result without storing the full attention matrix in the slow HBM.

Online Softmax Computation

A key technical breakthrough was the "online" computation of Softmax. The standard Softmax requires knowledge of all elements in the input vector for normalization. FlashAttention uses a modified algorithm that allows Softmax to be computed in parts. It maintains two intermediate values (the current maximum and the sum of exponents), which are updated as new blocks are processed, allowing for an exact result without accessing the entire matrix at once.

Fusing Operations into a Single CUDA Kernel

All attention operations (the QKᵀ matrix multiplication, masking, Softmax, and multiplication by V) are combined into a single fused CUDA kernel. This drastically reduces the number of read/write operations to HBM: instead of multiple passes over the entire matrix, the algorithm loads a block into SRAM once, performs all computations, and writes only the final result.

Theoretical and Practical Efficiency

Complexity and Optimality

FlashAttention reduces memory consumption from O(N²) to O(N), enabling linear scaling. It has been proven that the algorithm's IO complexity is theoretically optimal for computing attention in a two-level memory hierarchy, meaning it is impossible to perform exact attention faster without hardware changes.

Empirical Results

The first version of FlashAttention demonstrated significant improvements:

• Speedup:
  • BERT-large (sequence length 512): 15% training speedup.
  • GPT-2 (sequence length 1K): 3x speedup.
  • Long-Range Arena tasks (1K-4K): 2.4x speedup.
• Memory Savings: Up to 20x memory savings compared to exact baseline implementations.
• Improved Model Quality: By enabling work with longer contexts, FlashAttention not only avoids quality loss but actually improves model quality. For example, GPT-2's perplexity improved by 0.7 points, and accuracy on long-document classification tasks increased by 6.4 points.

Evolution and Further Developments

The success of FlashAttention initiated a whole series of hardware-aware algorithms.

FlashAttention-2 (2023)

The second version was aimed at more fully utilizing GPU resources. In the original FlashAttention, efficiency on the NVIDIA A100 was only 25–40% of the maximum. FlashAttention-2 introduced improvements in computation parallelization, which allowed for:

• Achieving a 2x speedup compared to the first version.
• Increasing GPU utilization to 50–73% of the theoretical maximum.
• Expanding support for attention heads of size 256, as well as for Multi-Query Attention (MQA) architectures.

FlashAttention-3 (2024)

The third version was specifically optimized for the NVIDIA Hopper (H100) GPU architecture. It utilizes new hardware features such as Tensor Core asynchrony and FP8 support, which made it possible to:

• Achieve another 1.5–2x speedup compared to FlashAttention-2.
• Reach performance of up to 740 TFLOPS on FP16 and close to 1.2 PFLOPS on FP8.

Specialized Solutions

The ideas behind FlashAttention have been extended in other projects:

• FlashInfer (2025): A customizable attention engine specifically optimized for LLM inference tasks. It focuses on efficient handling of the KV cache in streaming generation mode.
• FlashMLA (2024): An implementation of attention with context cache compression (latent attention), which saves memory on very long sequences with minimal information loss.

Impact on the Industry and Ecosystem

FlashAttention became a fundamental breakthrough and quickly turned into the industry standard for efficient LLM training and inference. It has been integrated into key libraries such as PyTorch and Hugging Face and is used in most major language models (LLaMA, MPT, Falcon, Claude, etc.).

FlashAttention and its subsequent versions played a crucial role in expanding the context windows of language models: from 2–4k tokens (GPT-3) to 128k tokens (GPT-4) and even to millions of tokens in experimental models. The algorithm eliminated one of the main obstacles to scaling Transformers, opening up new possibilities for AI applications, from analyzing long documents to multimodal understanding.

Links

• Official FlashAttention repository on GitHub

Literature

• Dao, T. et al. (2022). FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. arXiv:2205.14135.
• Dao, T. (2023). FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning. arXiv:2307.08691.
• Shah, J. et al. (2024). FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-Precision. arXiv:2407.08608.
• Kwon, W. et al. (2023). Efficient Memory Management for Large Language Model Serving with PagedAttention. arXiv:2309.06180.
• Hong, K. et al. (2023). FlashDecoding++: Faster Large Language Model Inference on GPUs. arXiv:2311.01282.
• Ye, Z. et al. (2025). FlashInfer: Efficient and Customizable Attention Engine for LLM Inference Serving. arXiv:2501.01005.
• Dege, P. et al. (2025). FlashMLA-ETAP: Efficient Transpose Attention Pipeline for Accelerating MLA Inference on NVIDIA H20 GPUs. arXiv:2506.01969.
• Wang, G. et al. (2025). FlashMask: Efficient and Rich Mask Extension of FlashAttention. OpenReview wUtXB43Chi.
• Dao, T. et al. (2022). FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness (OpenReview version). OpenReview H4DqfPSibmx.
• Gholami, A. et al. (2024). FlashAttention on a Napkin: A Diagrammatic Approach to Deep Learning IO-Awareness. OpenReview pF2ukh7HxA.