PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Tri Dao

arXiv:2307.08691v1cs.LG

TL;DR

Attention’s quadratic runtime and memory challenge long-context modeling, while FlashAttention remains less efficient than optimized matrix multiplication. FlashAttention-2 improves parallelism and work partitioning and reduces non-matrix-multiplication operations, achieving around 2× speedup over FlashAttention across benchmark settings.

  • Problem

    Long-context modeling faces quadratic attention costs, and FlashAttention remains less efficient than optimized matrix multiplication.

    Source: §1 · pp. 1–2

  • Method

    FlashAttention-2 improves parallelism and work partitioning while reducing non-matrix-multiplication operations without changing the output.

    Source: §1 · p. 2

  • Results

    Around 2× speedup over FlashAttention is achieved across benchmarks varying causal masking and head dimensions.

    Source: §1 · p. 2

  • Takeaways & Limitations

    16k-context models can be trained for the same price as previously training 8k-context models.

    Source: §5 · p. 11

    FlashAttention-2 uses manually tuned block sizes for each head dimension, with automatic tuning left to future work.

    Source: §3.3 · p. 9

Abstract

from arXiv · show

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$\times$ 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$\times$ 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).

1 Introduction

FlashAttention-2 reduces non-matmul operations and improves parallelism and work partitioning to address FlashAttention’s GPU inefficiencies. It achieves around 2× speedup over FlashAttention and GPT-style training speeds up to 225 TFLOPs/s per A100 GPU.

1 Introduction

  • 2-4× wall-clock speedup and up to 10-20× memory saving were achieved by FlashAttention [5] through tiling and recomputation, without approximation.Attention’s quadratic runtime and memory requirements constrain longer-context Transformers [18], motivating applications such as long-document querying and story writing. FlashAttention reduces memory usage to linear in sequence length and has seen wide adoption in Transformer training and inference.Source: §1 · pp. 1–2
  • 30-50% forward and 25-35% backward theoretical maximum throughput expose FlashAttention’s efficiency gap against optimized GEMM, which reaches up to 80-90%.The backward figure is reported on A100. Profiling attributes inefficiency to suboptimal thread-block and warp partitioning, causing low occupancy or unnecessary shared-memory reads and writes.Source: §1 · p. 2
  • Up to 16× higher matmul throughput than non-matmul throughput motivates FlashAttention-2’s reduction of non-matmul FLOPs without changing outputs.Specialized matrix-multiply units make non-matmul operations disproportionately expensive despite their small share of total FLOPs.Source: §1 · p. 2
  • FlashAttention-2 parallelizes forward and backward computation along sequence length, alongside batch and head dimensions, improving occupancy for long sequences and often-small batches.This expands parallelism beyond the batch and head dimensions.Source: §1 · p. 2
  • Within each thread block, FlashAttention-2 partitions work between warps to reduce communication and shared-memory reads and writes.Source: §1 · p. 2
  • Around 2× faster than FlashAttention, FlashAttention-2 reaches up to 73% forward and 63% backward theoretical maximum throughput.Benchmarks cover different head dimensions, with and without causal masking. End-to-end GPT-style training reaches up to 225 TFLOPs/s per A100 GPU.Source: §1 · p. 2

2 Background

Standard attention incurs quadratic memory costs by materializing intermediate matrices in GPU high-bandwidth memory. FlashAttention [5] reduces memory traffic without approximation through tiled computation and backward-pass recomputation.

2.1 Hardware characteristics

  • GPU execution combines specialized low-precision matrix-multiply units with an HBM–SRAM hierarchy and thread blocks whose warps communicate through shared memory.Thread blocks run on streaming multiprocessors; each warp contains 32 threads that can communicate through fast shuffle instructions or cooperate on matrix multiplication. Kernels load inputs from HBM into registers and SRAM, compute, and write outputs back to HBM.Source: §2.1 · p. 2

2.2 Standard Attention Implementation

  • O(N^2) memory is required by standard attention because it materializes score and probability matrices in HBM, making memory traffic a runtime bottleneck.For inputs Q, K, V and output O of shape N × d, computation proceeds as S = QK^T, row-wise P = softmax(S), then O = PV. Heads and batches run in parallel, and P is saved for backpropagation.Source: §2.2 · p. 3

2.3 FlashAttention

  • FlashAttention [5] targets accelerator memory reads and writes while preserving attention outputs without approximation.Source: §2.3 · p. 3

2.3.1 Forward pass

  • 2-4× wall-clock speedup over optimized baseline attention implementations comes from FlashAttention’s tiled forward pass, which avoids writing intermediate matrices to HBM.Input blocks move from HBM to SRAM; online softmax [11] [13] rescales block outputs to recover exact attention.Source: §2.3.1 · p. 3

2.3.2 Backward pass

  • 10-20× memory savings and 2-4× backward-pass wall-clock speedup come from recomputing attention matrices in SRAM rather than storing them.Memory requirements become linear rather than quadratic in sequence length, with savings depending on sequence length. Backward tiling requires keeping more values in SRAM for five matrix multiplications, versus two in the forward pass, complicating implementation despite needing no softmax rescaling.Source: §2.3.2 · p. 5

3 FlashAttention-2: Algorithm, Parallelism, and Work Partitioning

FlashAttention-2 reduces non-matmul operations, adds sequence-length parallelism across thread blocks, and repartitions work between warps to reduce shared-memory access. Together, these improvements yield 2-3× speedup.

3 FlashAttention-2: Algorithm, Parallelism, and Work Partitioning

  • 2-3× speedup comes from reducing non-matmul FLOPs, increasing thread-block parallelism, and reducing shared-memory communication between warps.These changes combine algorithmic tweaks with improved GPU work partitioning.Source: §3 · p. 5

3.1 Algorithm

  • 16× greater cost per non-matmul FLOP motivates reducing non-matmul operations on A100.A100's theoretical throughput is 312 TFLOPs/s for FP16/BF16 matmul versus 19.5 TFLOPs/s for non-matmul FP32 operations.Source: §3.1 · p. 5

3.1.1 Forward pass

  • The forward pass maintains unscaled output accumulators until final normalization and stores only logsumexp rather than separate maxima and exponential sums.These online-softmax changes reduce non-matmul FLOPs.Source: §3.1.1 · p. 5
  • O(N^2d) FLOPs and O(N) additional memory suffice for exact attention without approximation.Skipping fully causally masked blocks gives around 1.7-1.8× speedup compared to attention without the causal mask. With square blocks, only one block per row requires explicit causal masking.Source: §3.1.1 · pp. 6–7

3.1.2 Backward pass

  • The backward pass largely preserves FlashAttention, replacing separate row-wise maxima and exponential sums with row-wise logsumexp.MQA [15] and GQA [1] share key/value heads across query heads to reduce inference KV-cache size. Index manipulation avoids duplicating key/value heads, while backward computation sums dK and dV across the implicitly duplicated heads.Source: §3.1.2 · p. 7

3.2 Parallelism

  • Sequence-length parallelism improves GPU occupancy for long sequences when batch sizes or head counts are small.Forward workers process row blocks without communicating; backward workers process column blocks and use atomic adds to accumulate dQ.Source: §3.2 · p. 8, Figure 2 · p. 8

3.3 Work Partitioning Between Warps

  • Splitting Q across warps instead of K and V eliminates forward inter-warp communication and reduces shared-memory reads and writes.FlashAttention's split-K scheme requires shared-memory writes, synchronization, and intermediate-result reduction. The backward pass also avoids split-K, although gradient dependencies still require some synchronization.Source: §3.3 · p. 9
  • Block sizes are manually tuned per head dimension; larger blocks reduce shared-memory traffic but risk register spilling or exceeding shared-memory capacity.Typical block sizes are {64, 128} × {64, 128}, chosen according to head dimension and device shared memory. Automating these four choices through auto-tuning remains future work.Source: §3.3 · p. 9

4 Empirical Validation

FlashAttention-2 accelerates both attention computation and end-to-end GPT-style training, reaching up to 230 TFLOPs/s for attention and 225 TFLOPs/s per A100 GPU during training.

4 Empirical Validation

  • 1.7-3.0× speedups over FlashAttention, 1.3-2.5× over FlashAttention in Triton, and 3-10× over standard attention are reported for FlashAttention-2.Attention throughput reaches up to 230 TFLOPs/s, or 73% of the theoretical maximum on A100 GPUs.Source: §4 · pp. 9–10

4.2 End-to-end Performance

  • 2.8× faster training than a baseline without FlashAttention and up to 1.3× faster training than FlashAttention are achieved by FlashAttention-2.Training 1.3B/2.7B GPT-style models at 2k/8k sequence lengths on eight A100 GPUs reaches up to 225 TFLOPs/s per GPU (72% model FLOPs utilization). Reported FLOPs follow Megatron-LM [16], without halving the attention term for causal masking.Source: §4 · p. 10, §4.2 · p. 11

4.1 Benchmarking Attention

  • 1.3-1.5× forward-pass and around 2× backward-pass speedups over FlashAttention in Triton accompany FlashAttention-2’s roughly 2× speedup over FlashAttention and xformers’ cutlass implementation.Benchmarks use an A100 80GB SXM4, sequence lengths 512–16k with 16k total tokens, hidden dimension 2048, head dimensions 64/128, and causal/noncausal attention. FlashAttention-2 is up to 10× faster than standard PyTorch attention.Source: Figure 4 · p. 10, §4.1 · p. 10, Figure 5 · p. 11
  • 335 TFLOPs/s is reached on H100 using the unchanged implementation, without special instructions exploiting TMA or fourth-generation Tensor Cores.A further 1.5-2× speedup from new instructions is an expectation left to future work, not a measured result.Source: §4.1 · p. 11

5 Discussion and Future Directions

FlashAttention-2 is 2× faster than FlashAttention, enabling 16k-context training at the previous cost of 8k-context training. Future work targets broader hardware and datatype support, longer contexts, and easier programming of its optimizations.

5 Discussion and Future Directions

  • 2× faster than FlashAttention, FlashAttention-2 enables training with 16k context for the same price as previously training with 8k context.The authors envision applications to understanding long books and reports, high-resolution images, audio, and video.Source: §5 · pp. 11–12
  • 225 TFLOPs/s/GPU (72% model FLOPs utilization) is the peak FlashAttention-2 training speed reported for GPT-style models on 8×A100 GPUs.The reported comparison uses a baseline running without FlashAttention.Source: Table 1 · p. 12
  • FlashAttention-2 is also expected to accelerate training, finetuning, and inference of existing models.These benefits extend beyond the proposed longer-context applications.Source: §5 · p. 12
  • Planned hardware and datatype expansion includes H100 GPUs, AMD GPUs, and FP8, with H100 optimization identified as the immediate next step.The H100 work would exploit TMA, fourth-generation Tensor Cores, and FP8.Source: §5 · p. 12
  • Combining FlashAttention-2’s low-level optimizations with local, dilated, or block-sparse attention could enable training with much longer context.The authors also propose collaborating with compiler researchers to make these optimization techniques easily programmable.Source: §5 · p. 12

Related papers

Loading 2307.08691v1…