FlashAttention-2

FlashAttention-2 is an advanced algorithm designed to compute the attention mechanism in large language models (LLMs). The algorithm was developed by Tri Dao and researchers from Stanford University and was introduced in July 2023. Its key objective is to significantly accelerate the training and inference of transformer models by more efficiently utilizing GPU hardware resources, while maintaining full computational identity with the standard attention mechanism, i.e., with no loss of accuracy.

FlashAttention-2 is a logical successor to the FlashAttention algorithm, introduced by the same team in 2022. The new version addresses the issue of incomplete GPU utilization observed in its predecessor and achieves a nearly twofold speed increase compared to the first version.

Background: The Attention Bottleneck in Transformers

The standard self-attention mechanism is a bottleneck when processing long text sequences in transformers. Its computational complexity and memory consumption grow quadratically (O(N²)) with the sequence length (N), imposing severe limitations on the maximum context length and scalability of LLMs.

To address this problem, the FlashAttention algorithm was introduced in 2022. Its key ideas were:

• IO-awareness: The algorithm minimizes expensive read/write operations between the slow GPU High Bandwidth Memory (HBM) and the fast on-chip Static Random-Access Memory (SRAM).
• Tiling: Computations are broken down into small blocks (tiles) that are processed in the fast SRAM, which avoids materializing the full attention matrix in memory.

This allowed for linear memory consumption growth (O(N)) and a 2–4x speedup compared to standard implementations. FlashAttention became widely adopted and facilitated the emergence of models with significantly larger context windows, for example, from 2–4k tokens (GPT-3) to 128k (GPT-4) and beyond. For instance, in the Falcon-40B model, using FlashAttention sped up inference by 3x and overall generation throughput by 5x compared to GPT-3.

Development and Goals of FlashAttention-2

Despite its success, the first version of FlashAttention did not fully utilize the GPU's computational resources. On NVIDIA A100 GPUs, its performance only reached 25–40% of the theoretical maximum (FLOPs/s). The main reason was suboptimal utilization of Streaming Multiprocessors and redundant operations with shared memory.

The goal of FlashAttention-2 was to further accelerate computations through more effective work parallelization and minimization of auxiliary operations. The algorithm was completely rewritten using low-level primitives from the NVIDIA CUTLASS 3.x library to achieve maximum performance.

Technical Architecture and Principles of Operation

FlashAttention-2 introduces three key improvements to enhance parallelism and efficiency:

1. Minimizing Non-Matrix Operations

The algorithm reduces the number of auxiliary floating-point operations that are not matrix multiplications (non-matmul FLOPs). Since GPU Tensor Cores are specifically optimized for matrix operations (GEMM) and perform them up to 16 times faster, this change allows the most powerful GPU units to be utilized for a larger portion of the time.

2. Improved Parallelism

In the original FlashAttention, the work on a single attention head was not parallelized, leading to idle time with long sequences and small batch sizes. FlashAttention-2 introduces inter-block parallelism: computations for a single attention head are now distributed across different GPU Streaming Multiprocessors, significantly increasing their utilization.

3. Optimized Work Partitioning within a Block

At the level of a single compute block, work was re-partitioned among thread groups (warps) to reduce data exchange via shared memory. This reduces the number of redundant read/write operations required for the Softmax normalization.

Performance and Efficiency

Thanks to these architectural improvements, FlashAttention-2 demonstrates a significant increase in performance:

• Twofold Speedup: The algorithm runs approximately 2 times faster than the first version of FlashAttention.
• High GPU Utilization: On NVIDIA A100 GPUs, it achieves 50–73% of the theoretical maximum throughput (TFLOPs), which is close to the efficiency of optimized matrix multiplication (GEMM) operations.
• Record-breaking Computation Speed:
  • On an A100 GPU, it reaches speeds of up to 225 TFLOP/s in an end-to-end training loop for a GPT-style model, corresponding to 72% utilization of the compute units. For comparison, standard attention under the same conditions loaded the GPU at less than 100 TFLOP/s.
  • On an H100 GPU, performance reaches 335 TFLOP/s.

This performance boost allows, for example, training a model with a 16k token context window in the same amount of time previously required for an 8k token window. Importantly, the algorithm remains exact and deterministic, so its application does not affect the model's prediction quality.

Application and Ecosystem Integration

FlashAttention-2 quickly became a standard tool in the LLM ecosystem. It is integrated into many popular frameworks and libraries:

• PyTorch: Native support.
• Hugging Face Transformers: Support is enabled with the parameter `attn_implementation="flash_attention_2"` when loading a model. It is compatible with dozens of architectures (GPT, Llama, Falcon, BERT, etc.).
• TensorRT-LLM, xFormers, and Triton: The algorithm is implemented for these platforms, ensuring its widespread adoption.

This integration allows FlashAttention-2 to be easily combined with other optimization techniques, such as quantization (GPTQ, QLoRA) and parameter-efficient fine-tuning (PEFT).

Comparison with Subsequent Versions

FlashAttention-3

Research in attention optimization is ongoing. In July 2024, Tri Dao introduced FlashAttention-3, which is aimed at leveraging the capabilities of the NVIDIA Hopper GPU architecture (H100/H200). Key innovations include:

• FP8 Support: Uses 8-bit floating-point computations for further acceleration.
• Asynchronous Operations: More effectively utilizes the asynchronous capabilities of the GPU.

FlashAttention-3 provides a 1.5–2x speedup compared to FlashAttention-2 on H100 GPUs, reaching performance of up to 740 TFLOP/s (75% of the theoretical maximum).

Literature

• Dao, T. (2023). FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning. arXiv:2307.08691.
• Dao, T. et al. (2022). FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. arXiv:2205.14135.
• 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.
• Ye, Z. et al. (2025). FlashInfer: Efficient and Customizable Attention Engine for LLM Inference Serving. arXiv:2501.01005.
• Chen, Y. et al. (2023). FlashDecoding++: Faster Large Language Model Inference on GPUs. arXiv:2311.01282.
• Liu, Y. et al. (2024). FastAttention: Extending FlashAttention-2 to NPUs and Low-Resource GPUs. OpenReview: 76NYyOrnfk.
• 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. (2024). FlashMask: Efficient and Rich Mask Extension of FlashAttention. OpenReview: rog0J435OO.
• Abbott, V.; Zardini, G. (2025). FlashAttention on a Napkin: A Diagrammatic Approach to Deep Learning IO-Awareness. arXiv:2412.03317.