[Paper Review] FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning

Tri Dao|arXiv (Cornell University)|Jul 17, 2023

Advanced Neural Network ApplicationsComputer Science140 citations

TL;DR

FlashAttention-2 speeds up attention in Transformers by reworking parallelism and work partitioning, achieving about 2x speedup over FlashAttention and up to 73% of theoretical max FLOPs in forward, 63% in backward, with end-to-end GPT-style training up to 225 TFLOPs/s per A100.

ABSTRACT

Scaling Transformers to longer sequence lengths has been a major problem in the last several years, promising to improve performance in language modeling and high-resolution image understanding, as well as to unlock new applications in code, audio, and video generation. The attention layer is the main bottleneck in scaling to longer sequences, as its runtime and memory increase quadratically in the sequence length. FlashAttention exploits the asymmetric GPU memory hierarchy to bring significant memory saving (linear instead of quadratic) and runtime speedup (2-4$ imes$ compared to optimized baselines), with no approximation. However, FlashAttention is still not nearly as fast as optimized matrix-multiply (GEMM) operations, reaching only 25-40\% of the theoretical maximum FLOPs/s. We observe that the inefficiency is due to suboptimal work partitioning between different thread blocks and warps on the GPU, causing either low-occupancy or unnecessary shared memory reads/writes. We propose FlashAttention-2, with better work partitioning to address these issues. In particular, we (1) tweak the algorithm to reduce the number of non-matmul FLOPs (2) parallelize the attention computation, even for a single head, across different thread blocks to increase occupancy, and (3) within each thread block, distribute the work between warps to reduce communication through shared memory. These yield around 2$ imes$ speedup compared to FlashAttention, reaching 50-73\% of the theoretical maximum FLOPs/s on A100 and getting close to the efficiency of GEMM operations. We empirically validate that when used end-to-end to train GPT-style models, FlashAttention-2 reaches training speed of up to 225 TFLOPs/s per A100 GPU (72\% model FLOPs utilization).

Motivation & Objective

• Motivate scaling Transformers to longer context lengths by reducing attention bottlenecks.
• Improve GPU utilization by rethinking work partitioning across thread blocks and warps.
• Reduce non-matmul FLOPs to keep most time in fast matmul operations.
• Increase parallelism across sequence length, batch, and heads to boost occupancy.
• Validate end-to-end training speedups on GPT-style models.

Proposed method

• Tweak the FlashAttention forward/backward algorithms to reduce non-matmul FLOPs without changing output.
• Parallelize attention across sequence length in addition to batch and heads to increase occupancy.
• Distribute work between warps within a thread block to minimize shared memory traffic.
• Use online softmax tiling to enable block-wise computation with correct final outputs.
• Apply causal masking efficiently by skipping blocks where appropriate and reducing redundant masking work.
• Provide algorithms (Algorithm 1 and Algorithm 2) describing forward and backward passes with block/warp partitioning.

Experimental results

Research questions

• RQ1 Can FlashAttention-2 achieve higher GPU throughput for both forward and backward attention passes compared to FlashAttention?
• RQ2 How much do changes in work partitioning and non-matmul FLOPs reduce runtime and memory traffic?
• RQ3 What end-to-end training speedups are realized when using FlashAttention-2 on GPT-style models across varying context lengths?
• RQ4 How close can FlashAttention-2 get to theoretical max FLOPs/s on modern GPUs (e.g., A100, H100)?

Key findings

• FlashAttention-2 delivers about 2x speedup over FlashAttention in benchmarks.
• Forward pass throughput reaches up to 73% of the theoretical max FLOPs/s on A100; backward pass up to 63%.
• End-to-end training with GPT-style models reaches up to 225 TFLOPs/s per A100 GPU (72% model FLOPs utilization).
• On GPT-3 1.3B/2.7B scale, FlashAttention-2 achieves significant improvements over baseline and FlashAttention across 2k and 8k contexts.
• On H100 GPUs, forward+backward speedups and raw throughput are expected to improve further with new hardware features.