17 July 2023

FlashAttention-2 roughly doubles attention speed on A100 GPUs

Tri Dao's FlashAttention-2 improved GPU work partitioning to deliver roughly 2x the speed of FlashAttention, reaching 50–73% of the theoretical maximum FLOPs/s in the configurations benchmarked in the paper, and up to 225 TFLOPs/s per A100 in the specific GPT-style training configurations reported.

Compute and infrastructure Transformers and attention AI hardware Deep learning Foundational method Independently validated

Background

The attention mechanism at the heart of Transformer models has a costly property: its runtime and memory grow quadratically with sequence length. Longer contexts mean much more compute, which is why scaling to long sequences had been a persistent problem in language modelling, high-resolution image understanding, and applications in code, audio, and video.

The original FlashAttention addressed part of this by exploiting the asymmetric memory hierarchy on GPUs, keeping intermediate results in fast on-chip memory rather than writing them back to slow high-bandwidth memory. That brought memory use down from quadratic to linear and sped things up by 2–4x compared to standard baselines, with no approximation. Even so, it was reaching only 25–40% of the theoretical maximum FLOPs/s the hardware could deliver. The gap came from how work was distributed across the GPU: suboptimal partitioning between thread blocks and warps caused low occupancy and unnecessary reads and writes to shared memory.

What happened

Tri Dao, submitting to arXiv on 17 July 2023, proposed FlashAttention-2 with three specific changes to fix that partitioning problem. First, the algorithm was adjusted to reduce non-matrix-multiply operations, since those run much slower than the dedicated matrix-multiply units on modern GPUs. Second, the attention computation was parallelised across different thread blocks even for a single attention head, increasing how much of the GPU stays occupied. Third, within each thread block, work was redistributed between warps to cut down on communication through shared memory.

Together those changes yielded roughly 2x speedup over FlashAttention. On A100 GPUs the method reached 50–73% of the theoretical maximum FLOPs/s in the configurations benchmarked in the paper, getting close to the efficiency of optimised matrix-multiply, which is the baseline that GPU compute is usually measured against. End-to-end training of GPT-style models reached up to 225 TFLOPs/s per A100 GPU in the specific training configurations reported, corresponding to 72% model FLOPs utilisation.

Why it mattered

The attention layer was the main bottleneck for training on longer sequences, with runtime and memory growing quadratically in sequence length. FlashAttention-2 brought GPU utilisation close to what optimised matrix-multiply operations achieve, reducing the constant-factor cost of longer-context training without any approximation, though quadratic scaling in sequence length persists. That made it possible to reach up to 225 TFLOPs/s per A100 GPU in the specific GPT-style training configurations reported in the paper.

Sources

FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning. arxiv.org. Primary source
FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning. arxiv.org. Primary source

Cite this page

AI Achievements. (2023). FlashAttention-2 roughly doubles attention speed on A100 GPUs. Retrieved 2026-08-29, from https://achievements.ai/milestone/flashattention-2-roughly-doubles-attention-speed-on-a100

@misc{achievements_flashattention_2_roughly_doubles_attention_speed_on_a100,
  title  = {FlashAttention-2 roughly doubles attention speed on A100 GPUs},
  author = {{AI Achievements}},
  year   = {2023},
  url    = {https://achievements.ai/milestone/flashattention-2-roughly-doubles-attention-speed-on-a100}
}