cs.CLMar 5, 2026

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

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

Organizations: Princeton University · Together AI · Meta · Colfax Research · NVIDIA · Georgia Tech

Read PDF
Ask AI
Cite

Copy citation

DOINot available for this paper

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.

Figures & tables

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.
Resource1283128^{3}256×1282256\times 128^{2}
MMA compute10242048
Shared memory7681536
Exponential unit10242048
Table 1: Roofline analysis (cycles) for the attention forward pass. For both tile sizes, MMA compute and exponential unit are the primary bottlenecks.
 FP32 vs FP64BF16 vs FP64
MethodMax rel errMean rel errMax rel errMean rel err
Ideal (FP64 →\to BF16)——3.89×10−33.89\times 10^{-3}1.41×10−31.41\times 10^{-3}
Hardware MUFU.EX21.41×10−71.41\times 10^{-7}3.04×10−83.04\times 10^{-8}3.89×10−33.89\times 10^{-3}1.41×10−31.41\times 10^{-3}
Degree 38.77×10−58.77\times 10^{-5}5.43×10−55.43\times 10^{-5}3.90×10−33.90\times 10^{-3}1.41×10−31.41\times 10^{-3}
Degree 43.05×10−63.05\times 10^{-6}1.84×10−61.84\times 10^{-6}3.89×10−33.89\times 10^{-3}1.41×10−31.41\times 10^{-3}
Degree 51.44×10−71.44\times 10^{-7}5.48×10−85.48\times 10^{-8}3.89×10−33.89\times 10^{-3}1.41×10−31.41\times 10^{-3}
Table 2: Accuracy of 2x2^{x} polynomial emulation on [0,1)[0,1) , measured against FP64 reference on 4M random inputs. FP32 columns measure the raw polynomial output; BF16 columns measure after rounding to BF16. The BF16 quantization error dominates for all degrees ≥3\geq 3 .
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.
ResourceCycles ( N=d=128N=d=128 )Cycles ( N=d=128N=d=128 )
 1-CTA ( M=128M=128 )2-CTA ( M=256M=256 )
MMA compute25602560
Shared memory (MMA operands)20481536
Shared memory ( dS\mathbf{dS} write)256256
Shared memory ( dS\mathbf{dS} DSMEM )0384
Shared memory ( dQ\mathbf{dQ} write + read)1024512
Table 3: Roofline analysis for the attention backward pass with M=N=d=128M=N=d=128 . Shared memory traffic is the bottleneck, exceeding MMA compute time by approximately 30%. In the 2-CTA setting with M=256M=256 and N=d=128N=d=128 (with an exception to the dQ\mathbf{dQ} mma, with M=N=128M=N=128 and d=256d=256 ), shared memory traffic exceeds MMA compute time by approximately 5%.
Figure 3: In the 2-CTA backward dQdQ step, the CTA pair uses DSMEM to exchange half of the dSdS tile so each CTA forms an (M2×2N)(\frac{M}{2}\times 2N) operand and can run a CTA-pair UMMA with a doubled reduction.
MethodForward passBackward pass
FlashAttention-355s45s
FlashAttention-42.5s1.4s
Speedup22 ×\times32 ×\times
Table 4: Compile time for a single kernel: FA3 (C++ templates) and FA4 (CuTe-DSL). Typically FA2 and FA3 require precompiling hundreds of kernels for different attention variants.
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)
Figure 6: Backward pass TFLOPS on B200 (FP16/BF16) with head dimension 128. Left: non-causal attention. Right: causal attention.
Figure 7: Ablations for Deterministic Backward pass on B200 (FP16/BF16) with head dimension 128. Causal attention – SPT, LPT with reverse mblock order, LPT, and naive with no batch/head swizzle.
Appendix figures & tables1 asset

Supplementary material from the paper’s appendix.

Appendix

Figure 8: Ablations for Deterministic Backward pass on B200 with head dimension 128. Non-causal attention with batch/head swizzle versus naive. Right: causal attention – SPT, LPT with reverse mblock order, LPT, and naive with no batch/head swizzle.

Explore similar work

CardsList

  1. FlashAttention for Scalable Vector Architectures

    Aug 19, 2026Sonia Rani Gupta, Nikela Papadopoulou, Miquel PericàsEfficient Transformer InferenceEfficient Attention

    Chalmers University of Technology Gothenburg, Sweden · University of Glasgow Glasgow, UK

  2. LampAttention: Look-Ahead Mixed-Precision FlashAttention for Dedicated Accelerators

    Sep 30, 2026Stanislav Budzinskiy, Marian Gloser, Tolunay Yilmaz +5AI Accelerator InferenceLLM Inference Acceleration

    Faculty of Mathematics University of Vienna, Austria · Huawei Heisenberg Research Center, Munich, Germany · Huawei Technologies Co. Ltd

  3. Hardware-Aware FP4 FlashAttention-4

    Sep 3, 2026Robert HuMixed-Precision QuantizationGPU Kernel Optimization

    GRAPHCORE RESEARCH