FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning

Tri Dao

2023 · arXiv

FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning

Problem

Framing

FlashAttention removes quadratic memory traffic, but still leaves GPU throughput far below GEMM efficiency. FlashAttention-2 closes that gap with better work partitioning, fewer non-matmul operations, and sequence-parallel attention, reaching 50–73% of peak A100 FLOPs and 225 TFLOPs/s per GPU in GPT training.

Currently Used Methods

Foundational

@vaswaniAttentionAllNeed2017 — standard exact attention for Transformers.
- Limitation in context: quadratic memory traffic blocks long-context efficiency.
@daoFlashAttention2022 — tiled exact attention with linear-memory IO complexity.
- Limitation in context: poor warp and block partitioning leaves low occupancy.
"Longformer: The Long-Document Transformer" — sparse local-global attention for long contexts.
- Limitation in context: changes the operator instead of accelerating exact attention.
"Self-attention Does Not Need O(n²) Memory" — memory-efficient recomputation for exact attention.
- Limitation in context: less optimized for GPU kernel throughput.
"GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints" — shares KV heads to reduce cache cost.
- Limitation in context: inference-memory savings do not fix attention kernel efficiency.

Proposed Method

Architecture

The kernel keeps exact attention but changes how GPU work is assigned. It parallelizes a single head across sequence-length blocks and repartitions warps so $Q$ stays shared while $K^T$ and $V$ tiles are split across warps, reducing shared-memory traffic.

Verified diagram: FlashAttention forward tiling with online softmax across key/value blocks, plus a warp-partition view where Q is shared and K^T/V are split across warps.

Verified diagram: intra-block warp partitioning where all warps access Q while K^T and V are sharded across four warps.

Loss / Objective

The method is an exact kernel rewrite.Its core forward update uses online softmax with fewer rescalings:

\mathbf{m}^{(1)} = \operatorname{rowmax}(\mathbf{S}^{(1)}), \qquad \boldsymbol{\ell}^{(1)} = \operatorname{rowsum}\!\left(e^{\mathbf{S}^{(1)}-\mathbf{m}^{(1)}}\right), \qquad \tilde{\mathbf{O}}^{(1)} = e^{\mathbf{S}^{(1)}-\mathbf{m}^{(1)}} \mathbf{V}^{(1)}

\mathbf{m}^{(2)} = \max\!\left(\mathbf{m}^{(1)}, \operatorname{rowmax}(\mathbf{S}^{(2)})\right), \qquad \boldsymbol{\ell}^{(2)} = e^{\mathbf{m}^{(1)}-\mathbf{m}^{(2)}} \boldsymbol{\ell}^{(1)} + \operatorname{rowsum}\!\left(e^{\mathbf{S}^{(2)}-\mathbf{m}^{(2)}}\right)

\mathbf{O}^{(2)} = \operatorname{diag}\!\left((\boldsymbol{\ell}^{(2)})^{-1}\right)\!\left(e^{\mathbf{m}^{(1)}-\mathbf{m}^{(2)}} \tilde{\mathbf{O}}^{(1)} + e^{\mathbf{S}^{(2)}-\mathbf{m}^{(2)}} \mathbf{V}^{(2)}\right)

Sampling Rule / Algorithm

The exact attention target remains unchanged:

\operatorname{Attn}(\mathbf{Q},\mathbf{K},\mathbf{V}) = \operatorname{softmax}\!\left(\frac{\mathbf{Q}\mathbf{K}^\top}{\sqrt{d}}\right)\mathbf{V}

FlashAttention-2 computes this by tiling over sequence blocks and parallelizing those blocks across thread blocks and warps.

Training Procedure

Hardware: 8\times A100 80GB SXM for end-to-end GPT training.
Models: GPT-style 1.3B and 2.7B.
Context lengths: 2k and 8k.
Benchmarks: causal and non-causal attention.
Head dimensions: 64, 128, 256.
Variants supported: MQA and GQA.

Evaluation

Datasets

Synthetic attention benchmarks on A100 80GB SXM4.
GPT-style language-model training, 1.3B and 2.7B parameters.
Context lengths: 512 to 16k in kernel benchmarks.
End-to-end training at 2k and 8k context.

Metrics

Attention throughput: TFLOPs/s.
Device efficiency: percent of theoretical max FLOPs/s.
End-to-end training throughput: TFLOPs/s per GPU.
Relative speedup over PyTorch, FlashAttention, xFormers, and Triton baselines.

Headline results

A100 forward pass: up to 73% theoretical max throughput.
A100 backward pass: up to 63% theoretical max throughput.
Attention benchmarks: 1.7–3.0\times faster than FlashAttention.
GPT training: up to 1.3\times faster than FlashAttention.
GPT training: up to 225 TFLOPs/s per A100, 72% model FLOPs utilization.

Verified results plot: bar chart comparing PyTorch, FlashAttention, xFormers, FlashAttention Triton, and FlashAttention-2 throughput across sequence lengths; FlashAttention-2 is highest in every shown setting.

Ablations

Work partitioning across sequence length raises occupancy when batch size or head count is small.
Warp repartitioning avoids split- $K$ accumulation and cuts shared-memory reads and writes.
Algorithmic rescaling changes reduce non-matmul FLOPs without changing exact outputs.
The kernel extends to head dimension 256, MQA, and GQA.

Method Strengths and Weaknesses

Strengths

Exact attention, not an approximation.
Up to 2\times faster than FlashAttention on kernel benchmarks.
Reaches 225 TFLOPs/s per A100 in GPT training.
Supports causal, non-causal, MQA, GQA, and head dimension 256.

Weaknesses

Gains depend on GPU-specific kernel engineering.
End-to-end speedup is smaller than kernel-level speedup.
Evaluation centers on A100; cross-hardware portability is unclear.
The contribution is systems-heavy, not a new modeling capability.

Suggestions from the authors

Train 16k-context models at the cost of prior 8k-context runs.
Push exact attention closer to GEMM-level efficiency.
Extend long-context training to books, high-resolution images, and video.
Explore new applications unlocked by cheaper long-context Transformers.

FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning

FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning

Problem

Framing

Currently Used Methods

Foundational

Proposed Method

Architecture

Loss / Objective

Sampling Rule / Algorithm

Training Procedure

Evaluation

Datasets

Metrics

Headline results

Ablations

Method Strengths and Weaknesses

Strengths

Weaknesses

Suggestions from the authors

Links

Prior Papers

Further Papers