PaperlayerView exact revision on arXiv

Source-linked AI summary

FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision

Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri Dao

arXiv:2407.08608v2cs.LGcs.AI

TL;DR

Attention remains a bottleneck for long-context Transformers, and FlashAttention-2 does not explicitly exploit newer GPU asynchrony and low-precision features. FlashAttention-3 redesigns the computation with warp specialization, overlapped GEMM-softmax execution, and FP8-aware quantization. On H100, it reports 1.5-2.0× FP16 speedups, close to 1.2 PFLOPs/s with FP8, and 2.6× lower FP8 numerical error than per-tensor quantization.

  • Problem

    Attention is a primary computational bottleneck for long-context Transformer applications, while FlashAttention-2 does not explicitly use hardware asynchrony and low-precision computation.

  • Method

    FlashAttention-3 combines producer-consumer warp specialization, asynchronous overlap of block-wise GEMMs and softmax, and FP8 block quantization with incoherent processing.

  • Results

    1.5-2.0× FP16 forward-pass speedup over FlashAttention-2 is reported on H100, with up to 740 TFLOPs/s, close to 1.2 PFLOPs/s for FP8, and 2.6× lower FP8 numerical error than standard per-tensor quantization.

  • Takeaways & Limitations

    The reported results show that exploiting Hopper’s asynchrony and low-precision hardware can improve attention efficiency while preserving or improving numerical accuracy.

  • Takeaways & Limitations

    The 2-stage pipeline requires extra registers for intermediate results, creating a trade-off with larger block sizes that must be resolved through profiling.

Abstract

from arXiv · show

Attention, as a core layer of the ubiquitous Transformer architecture, is the bottleneck for large language models and long-context applications. FlashAttention elaborated an approach to speed up attention on GPUs through minimizing memory reads/writes. However, it has yet to take advantage of new capabilities present in recent hardware, with FlashAttention-2 achieving only 35% utilization on the H100 GPU. We develop three main techniques to speed up attention on Hopper GPUs: exploiting asynchrony of the Tensor Cores and TMA to (1) overlap overall computation and data movement via warp-specialization and (2) interleave block-wise matmul and softmax operations, and (3) block quantization and incoherent processing that leverages hardware support for FP8 low-precision. We demonstrate that our method, FlashAttention-3, achieves speedup on H100 GPUs by 1.5-2.0$\times$ with FP16 reaching up to 740 TFLOPs/s (75% utilization), and with FP8 reaching close to 1.2 PFLOPs/s. We validate that FP8 FlashAttention-3 achieves 2.6$\times$ lower numerical error than a baseline FP8 attention.

1 Introduction

Attention is a computational bottleneck for long-context Transformers, while FlashAttention-2 leaves newer GPU asynchrony and low-precision capabilities underused. FlashAttention-3 redesigns attention around these features and reports substantial H100 speedups with improved FP8 accuracy.

  • Motivation: Quadratic self-attention scaling makes attention a primary computational bottleneck for long-context Transformer applications.Longer contexts support document, code, image, audio, video, interaction-history, and agent-workflow applications.
  • Motivation: FlashAttention-2 uses a simplified synchronous model without explicitly exploiting hardware asynchrony or low-precision computation.The redesign must overlap dependent matmul and softmax operations while controlling FP8 quantization error around outlier features.
  • Contributions: FlashAttention-3 introduces producer-consumer warp specialization to overlap data movement and Tensor Core computation.Separate producer and consumer warps extend the ability to hide memory and instruction-issue latencies.
  • Contributions: FlashAttention-3 overlaps block-wise GEMMs with softmax by exploiting asynchronous WGMMA and circumventing sequential dependencies.In the 2-stage pipeline, softmax processes one score block while WGMMA computes the next block asynchronously.
  • Contributions: FlashAttention-3 targets FP8 Tensor Cores with block quantization and incoherent processing to nearly double measured throughput while mitigating accuracy loss.The method also bridges WGMMA layout requirements for FP32 accumulators and FP8 operands.
  • Results: 1.5-2.0× forward-pass speedup over FlashAttention-2 is reported for FP16, reaching 740 TFLOPs/s, while FP8 reaches close to 1.2 PFLOPs/s.The H100 evaluation also reports 1.5-1.75× backward-pass speedup.
  • Results: 2.6× lower numerical error is reported for FP8 FlashAttention-3 than standard attention with per-tensor quantization in cases with outlier features.FP16 FlashAttention-3 matches FlashAttention-2’s numerical error and improves over standard attention by retaining intermediate results such as softmax rescaling in FP32.

2 Background: Multi-Head Attention and GPU Characteristics

The attention equations operate on query, key, and value sequences, while GPU execution depends on hierarchies of memory and threads. Hopper adds asynchronous TMA and WGMMA execution plus FP8 Tensor Core support that FlashAttention-3 can exploit.

  • Multi-Head Attention: For one attention head, Q, K, and V have shape R^N×d, where N is sequence length and d is head dimension.The attention output is formed from these inputs using scaled dot products, row-wise softmax, and multiplication by V.
  • Multi-Head Attention: The attention computation uses α = 1/√d as a typical scaling factor and subtracts each row maximum to stabilize exponentiation.Multi-head attention parallelizes this computation across heads and batches.
  • Multi-Head Attention: The backward pass computes dQ, dK, and dV from dO using the chain rule and row-wise softmax derivatives.The softmax derivative is expressed as (diag(p) − pp^T)dp for p = softmax(s).
  • GPU Characteristics: Hopper’s memory hierarchy spans global memory, L2 cache, shared memory, and registers, with capacity inversely related to bandwidth.Shared memory is addressable by threads within a CTA, while registers are private to individual threads.
  • GPU Characteristics: Hopper organizes execution from threads through warps, warpgroups, threadblocks, clusters, and grids.A warp contains 32 threads, and threads in the same CTA are co-scheduled on one SM.
  • GPU Characteristics: Hopper provides asynchronous TMA memory copies and asynchronous WGMMA Tensor Core operations that can source inputs directly from shared memory.Warp-specialized kernels assign separate producer and consumer roles to data movement and computation.
  • GPU Characteristics: FP8 WGMMA provides 2× the per-SM throughput of FP16 or BF16, but accepts only k-major shared-memory input layouts.These layout restrictions require modifications to FP8 attention algorithms.
  • Standard and FlashAttention: FlashAttention avoids materializing intermediate S and P matrices in HBM by using local softmax reduction and fusing attention into one kernel.Standard attention is defined here as materializing S and P in HBM.

3 FlashAttention-3: Algorithm

FlashAttention-3 redesigns the attention forward pass around Hopper’s asynchronous data movement and Tensor Core execution. It combines warp-specialized buffering, pipelined GEMM-softmax overlap, and FP8-specific layout and accuracy techniques.

  • 3.1 Producer-Consumer asynchrony through warp-specialization and pingpong scheduling: FlashAttention-3 integrates warp-specialization with a circular shared-memory buffer, separating producer and consumer warpgroups for asynchronous data movement and computation.TMA loads populate shared memory while consumer warpgroups execute WGMMA-based GEMMs.
  • 3.1 Producer-Consumer asynchrony through warp-specialization and pingpong scheduling: 989 TFLOPS of FP16 matmul versus 3.9 TFLOPS of special functions motivates overlapping softmax with GEMMs on H100.Pingpong scheduling alternates softmax and GEMM work between two warpgroups.
  • 3.1 Producer-Consumer asynchrony through warp-specialization and pingpong scheduling: 570 to 620-640 TFLOPS: pingpong scheduling improves FP16 forward performance for head dimension 128 and sequence length 8192.The implementation’s schedule is less clean than the idealized depiction but generally improves performance.
  • 3.2 Intra-warpgroup overlapping GEMMs and softmax: The two-stage pipeline breaks sequential dependencies across iterations, allowing softmax on one score block to overlap with asynchronous WGMMA for the next block.Within one warpgroup, additional register buffers enable interleaving GEMMs and softmax operations.
  • 3.2 Intra-warpgroup overlapping GEMMs and softmax: 3-stage pipelining can further overlap the second WGMMA with softmax but requires more registers, complicating the tile-size and pipeline-depth trade-off.The added register demand makes balancing pipeline depth and tile size more difficult.
  • 3.3 Low-precision with FP8: 2.6×: block quantization and incoherent processing reduce FP8 FlashAttention-3’s numerical error.Incoherent processing applies an orthogonal transformation before FP8 quantization to spread outlier values without changing QKᵀ.

4 Empirical Validation

FlashAttention-3 is empirically evaluated on H100 GPUs for speed, ablations, FP8 performance, and numerical accuracy across attention settings.

  • Benchmarking attention: Up to 2.0× faster than FlashAttention-2, FlashAttention-3 reaches 740 TFLOPs/s, or 75% of H100’s theoretical maximum.The benchmark compares standard attention, FlashAttention-2, Triton, cuDNN, and FlashAttention-3 across sequence lengths.
  • Benchmarking attention: 1.5-2.0× forward and 1.5-1.75× backward speedups over FlashAttention-2 are observed across FP16 settings.FlashAttention-3 is also up to 3-16× faster than standard attention and surpasses cuDNN for medium and long sequences.
  • Accuracy of FP8 attention: The FP8 experiments measure forward runtime on H100 for comparable settings and report head dimension 256 results in Figure 7.Full FP8 results are provided in Appendix C.2.
  • Ablation study: From 570 to 661 TFLOPs, the ablation shows that warp-specialization and 2-stage WGMMA-softmax pipelining contribute to speedup.The ablation fixes batch, sequence length, head count, and head dimension while removing or including the two improvements.
  • Accuracy of FP8 attention: FP8 FlashAttention-3 is 2.6× more accurate than the FP8 baseline with per-tensor quantization under outlier-feature conditions.In FP16, both FlashAttention-2 and FlashAttention-3 achieve 1.7× lower RMSE than standard attention because intermediate softmax results remain in FP32.

5 Dicussion, Limitations, Conclusion

The paper concludes that Hopper’s asynchronous execution and low-precision features improve attention efficiency and accuracy, while identifying inference and training questions for future work.

  • Conclusion: Asynchrony and low-precision techniques substantially improve attention efficiency and accuracy in FlashAttention-3.The conclusion reports 1.5-2.0× speedup over FlashAttention-2 and a 2.6× reduction in FP8 numerical error versus standard per-tensor quantization.
  • Limitations: Future work includes optimizing LLM inference, integrating a persistent kernel into the FP8 kernel, and studying low-precision attention in large-scale training.These are stated limitations and open directions of the current work.
  • Scope: The methods are developed for Hopper GPUs, although the authors expect them to apply to other accelerators with sufficiently robust asynchronous and low-precision capabilities.The expected transfer is presented as a scope expectation rather than an evaluated result.
  • Implications: FlashAttention-3 can serve as an improved primitive for distributed attention methods that currently use FlashAttention or FlashAttention-2.The related-work discussion specifically connects the improvement to Ring attention and related multi-GPU approaches.
  • Related work: Existing quantization work largely targets KV-cache reduction for inference, while quantization during training remains challenging because stable training typically requires higher precision.This positions FlashAttention-3’s FP8 accuracy work within a broader low-precision attention context.

B.1 Asynchrony Through Warp Specialization for the Backward Pass

The backward pass extends warp specialization with a dedicated dQ-writer role to manage asynchronous computation and accumulation across thread blocks.

  • Warp specialization: A separate dQ-writer warp handles accumulation of local dQ results into global dQ, avoiding memory-contention stalls for other warps.Many thread blocks write to the same dQ location, motivating the dedicated writer role.
  • Preprocessing: A preprocessing kernel computes D as rowsum(dO ◦ O), writes it to HBM, and partitions it into blocks for the backward pass.The resulting D blocks are loaded alongside Q and dO during the pipeline.
  • Pipeline roles: The backward algorithm uses producer warps to load K and V, consumer warpgroups to compute gradients, and a circular shared-memory pipeline.The algorithm divides inputs and gradients into blocks and coordinates their movement through staged buffering.
  • Gradient accumulation: Local dQ values are written to shared memory and atomically added to global dQ using a semaphore in the dQ-writer warp.This separates local gradient production from contended global accumulation.

B.2 2-Stage Pipelining SASS Analysis

SASS analysis confirms that the 2-stage pipeline overlaps softmax-related operations with the first WGMMA while leaving the second WGMMA unoverlapped.

  • Pipeline overlap: Softmax is reordered to the beginning before the first WGMMA, while exponentiation, row sums, output rescaling, and conversions are also interleaved.These reordered operations create opportunities for asynchronous overlap.
  • Pipeline overlap: The first WGMMA is interleaved with softmax and FP32-to-FP16 conversion, indicating parallel execution of GEMM and non-WGMMA operations.This directly verifies the intended overlap in generated SASS code.
  • Pipeline boundaries: The second WGMMA is not overlapped with other instructions, as expected from the pipeline design.Its constituent HGMMAs are issued as a packed group rather than interleaved with other operations.
  • Validation: Overall, the SASS inspection shows that the 2-stage pipelining idea works as expected.The analysis validates the compiler-generated instruction schedule against the intended execution pattern.

B.3 3-Stage Pipelining Algorithm

The 3-stage pipeline overlaps work from successive iterations to increase concurrency, but it performs worse than the 2-stage pipeline because of compiler scheduling and register-pressure costs.

  • B.3 3-Stage Pipelining Algorithm: The 3-stage algorithm pipelines WGMMA operations and softmax across successive iterations using producer-consumer coordination.It processes WGMMA work from iterations j, j+1, and j+2 while softmax and data loading proceed asynchronously.
  • B.3 3-Stage Pipelining Algorithm: The 3-stage algorithm performs worse than the 2-stage pipeline.The supplied passage introduces compiler-overlap and register-pressure explanations for this result.
  • B.3 3-Stage Pipelining Algorithm: The consumer warpgroup loads Q_i, K_j, and V_j blocks, computes score matrices with WGMMA, updates softmax statistics, and accumulates output blocks.The epilogue rescales O_i, computes L_i, and writes both outputs to HBM.
  • B.3 3-Stage Pipelining Algorithm: The compiler overlaps softmax with only the first WGMMA, leaving the second WGMMA unoverlapped for unclear instruction-reordering reasons.This limits the intended overlap between softmax and both WGMMA operations.
  • B.3 3-Stage Pipelining Algorithm: The 3-stage pipeline requires more registers than the 2-stage version, including an extra P̃_i and scale_o, forcing a smaller block size.The additional storage includes a B_r × B_c input-data buffer and a B_r-sized float buffer.

C.1 System and libraries

The experiments benchmark FlashAttention on an H100 80GB SXM5 using recent software and controlled timing conditions.

  • C.1 System and libraries: The benchmarks run on an H100 80GB SXM5 with a 700W power specification.The authors generally use the latest library versions available at the time of writing.
  • C.1 System and libraries: The reported software environment includes Triton nightly 3.0.0.post20240424212437.
  • C.1 System and libraries: Benchmark variability is reduced by fixing the GPU clock at 1830MHz and averaging timings over 100 repetitions.The fixed clock is the one used to calculate the 989 TFLOPS FP16 theoretical maximum throughput.

C.2 FP8 Attention Full Results

The FP8 attention evaluation uses multiple sequence lengths on the H100, with longer sequences constrained to avoid wave quantization; results are presented for H100 FP8 forward speed.

  • C.2 FP8 Attention Full Results: For sequence lengths at least 4k, lengths are made divisible by 132, matching the number of SMs in the H100 SXM5.This choice is intended to avoid wave quantization.
  • C.2 FP8 Attention Full Results: Figure 9 reports attention forward speed in FP8 on an H100 GPU.
Loading 2407.08608v2…