Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

FlashAttention-4: Algorithm and Kernel Pipelining Co-Design for Asymmetric Hardware Scaling

Ted Zadouri, Markus Hoehnerbach, Jay Shah, Timmy Liu, Vijay Thakkar, Tri Dao

TL;DR

The proposed method, FlashAttention-4, achieves up to 1.3$\times$ speedup over cuDNN 9.13 and 2.7$\times$ over Triton on B200 GPUs with BF16, reaching up to 1613 TFLOPs/s (71% utilization).

Abstract

Attention, as a core layer of the ubiquitous Transformer architecture, is the bottleneck for large language models and long-context applications. While FlashAttention-3 optimized attention for Hopper GPUs through asynchronous execution and warp specialization, it primarily targets the H100 architecture. The AI industry has rapidly transitioned to deploying Blackwell-based systems such as the B200 and GB200, which exhibit fundamentally different performance characteristics due to asymmetric hardware scaling: tensor core throughput doubles while other functional units (shared memory bandwidth, exponential units) scale more slowly or remain unchanged. We develop several techniques to address these shifting bottlenecks on Blackwell GPUs: (1) redesigned pipelines that exploit fully asynchronous MMA operations and larger tile sizes, (2) software-emulated exponential and conditional softmax rescaling that reduces non-matmul operations, and (3) leveraging tensor memory and the 2-CTA MMA mode to reduce shared memory traffic and atomic adds in the backward pass. We demonstrate that our method, FlashAttention-4, achieves up to 1.3$\times$ speedup over cuDNN 9.13 and 2.7$\times$ over Triton on B200 GPUs with BF16, reaching up to 1613 TFLOPs/s (71% utilization). Beyond algorithmic innovations, we implement FlashAttention-4 entirely in CuTe-DSL embedded in Python, achieving 20-30$\times$ faster compile times compared to traditional C++ template-based approaches while maintaining full expressivity.

FlashAttention-4: Algorithm and Kernel Pipelining Co-Design for Asymmetric Hardware Scaling

TL;DR

The proposed method, FlashAttention-4, achieves up to 1.3 speedup over cuDNN 9.13 and 2.7 over Triton on B200 GPUs with BF16, reaching up to 1613 TFLOPs/s (71% utilization).

Abstract

Attention, as a core layer of the ubiquitous Transformer architecture, is the bottleneck for large language models and long-context applications. While FlashAttention-3 optimized attention for Hopper GPUs through asynchronous execution and warp specialization, it primarily targets the H100 architecture. The AI industry has rapidly transitioned to deploying Blackwell-based systems such as the B200 and GB200, which exhibit fundamentally different performance characteristics due to asymmetric hardware scaling: tensor core throughput doubles while other functional units (shared memory bandwidth, exponential units) scale more slowly or remain unchanged. We develop several techniques to address these shifting bottlenecks on Blackwell GPUs: (1) redesigned pipelines that exploit fully asynchronous MMA operations and larger tile sizes, (2) software-emulated exponential and conditional softmax rescaling that reduces non-matmul operations, and (3) leveraging tensor memory and the 2-CTA MMA mode to reduce shared memory traffic and atomic adds in the backward pass. We demonstrate that our method, FlashAttention-4, achieves up to 1.3 speedup over cuDNN 9.13 and 2.7 over Triton on B200 GPUs with BF16, reaching up to 1613 TFLOPs/s (71% utilization). Beyond algorithmic innovations, we implement FlashAttention-4 entirely in CuTe-DSL embedded in Python, achieving 20-30 faster compile times compared to traditional C++ template-based approaches while maintaining full expressivity.
Paper Structure (31 sections, 14 equations, 8 figures, 4 tables)

This paper contains 31 sections, 14 equations, 8 figures, 4 tables.

Table of Contents

  1. Introduction
  2. Background
  3. Multi-Head Attention
  4. GPU Hardware Characteristics and Execution Model
  5. Memory hierarchy:
  6. Thread hierarchy:
  7. Tensor cores and increased asynchrony:
  8. 2-CTA tensor core:
  9. Shifting bottlenecks:
  10. Algorithm
  11. Attention forward pass
  12. Feeds and Speeds
  13. New pipeline to overlap matmul and softmax
  14. Emulation of the exponential function
  15. Skipping online softmax rescaling
  16. ...and 16 more sections

Figures (8)

  • Figure 1: FlashAttention-4 forward pipeline. The superscript $^H$ denotes the matrices corresponding to the "high" Q tile, and superscript $^L$ denotes matrices corresponding to the "low" Q tile. Each Q tile corresponds to 128 query tokens.
  • Figure 2: FlashAttention-4 backward computation graph (5 MMA operations + 2 elementwise operations), showing the 1-CTA MMA mode software pipeline order across the prologue, main loop, and tail.
  • Figure 3: In the 2-CTA backward $dQ$ step, the CTA pair uses DSMEM to exchange half of the $dS$ tile so each CTA forms an $(\frac{M}{2} \times 2N)$ operand and can run a CTA-pair UMMA with a doubled reduction.
  • Figure 4: Forward pass TFLOPS on B200 (FP16/BF16) with head dimension 128. Left: non-causal attention. Right: causal attention. FA4 achieves 1.1-1.3$\times$ speedup over cuDNN 9.13.0 and 2.1-2.7$\times$ over Triton across sequence lengths. Since the initial release of our implementation, newer versions of cuDNN have incorporated many of the techniques described in this paper, yielding similar performance to FA4.
  • Figure 5: Forward pass TFLOPS comparison between cuDNN and FA4 on B200 (FP16/BF16) with head dimension (192, 128) for causal attention (typically used in DeepSeek V3 architecture)
  • ...and 3 more figures