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
Aug 19, 2026

FlashAttention for Scalable Vector Architectures

Inference with transformer models on CPUs is increasingly important, especially for Small Language Models (SLMs), where vector architectures are emerging as a promising execution substrate. The attention module is a major bottleneck due to high memory bandwidth requirements; FlashAttention mitigates this by fusing operations to improve data locality and reduce intermediate memory traffic. In this paper, we present FlashAttention-V, a blocked FlashAttention for scalable vector architectures that adapts efficiently from short to very long vectors by exploiting parallelism across attention heads, inter-head packing to enable efficient utilization of vector lengths beyond the head dimension, and improving vector register utilization and memory access locality. We integrate FlashAttention-V into ggml within llama.cpp and evaluate it on TinyLlama, Llama 3.2, Qwen2.5, and Pythia-410M using gem5 and a Banana Pi BPI-F3. On the Banana Pi BPI-F3, we confirm that loop reordering and loop unrolling across attention heads are effective optimization principles, scaling performance gains with larger models and most pronounced with short contexts and during decoding. Simulation-based analysis shows that FlashAttention-V achieves 22x-42x speedup over scalar FlashAttention at 512-bit VL in prefill, with an additional 2x-2.5x gain scaling to 64 lanes and 4096-bit VL. During decode, FlashAttention-V achieves 8x-11x speedup using 512-bit vector lengths over scalar FlashAttention, with performance showing diminishing sensitivity to vector width and lane count due to single-token, memory-bound execution. We further identify structural bottlenecks in Q8_0 quantized linear layers that limit arithmetic amortization under long-vector execution, consistent across RVV and Arm SVE, indicating that current quantization formats pose a fundamental challenge to long-vector scalability.

Sonia Rani Gupta, Nikela Papadopoulou, Miquel Pericàs

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

Efficient Transformer InferenceEfficient AttentionTransformer AttentionLLM Inference Acceleration
AskPDF
Sep 30, 2026

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

While most attention logits can be computed in low precision without degrading numerical stability, current attention kernels fail to exploit this phenomenon. We introduce a novel hardware-algorithm co-design in the form of mixed-precision FlashAttention. Our method accumulates key-query products and evaluates their exponentials in 8-bit formats, then adaptively identifies sensitive sub-blocks and recomputes them in 16-bit formats. We propose the specifications for a dedicated accelerator capable of executing this pipeline efficiently. Simulated experiments with Qwen3 and Gemma 3 show that rerouting a selective minority of sub-blocks to high precision is sufficient to recover the baseline model performance.

Stanislav Budzinskiy, Marian Gloser, Tolunay Yilmaz +5

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

AI Accelerator InferenceLLM Inference AccelerationEfficient AttentionAttention Mechanisms
AskPDF
Sep 3, 2026

Hardware-Aware FP4 FlashAttention-4

Blackwell's 4-bit floating-point (FP4) tensor cores do not automatically make attention faster because softmax conversion and on-chip dependencies dominate once its matrix products shrink. We address this with Direct-P for noncausal inference and a causal path that passes the forward quantization directly into backward. Direct-P maps scores directly to FP4 probabilities and reaches up to 2.13×\times the bfloat16 (BF16) forward throughput on an NVIDIA GB200. The causal path reconstructs probabilities from saved quantized queries and keys and uses 8-bit floating-point (FP8) gradient operands, accelerating a complete single-GPU 8-billion-parameter update by up to 1.14×\times. Matched distributed training retains FP8 probabilities and values; every tested MXFP4 probability/value training trajectory diverges.

Robert Hu

GRAPHCORE RESEARCH

Mixed-Precision QuantizationGPU Kernel OptimizationTransformer AttentionEfficient AttentionAttention QuantizationAttention Mechanisms
AskPDF