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

Abstract: Attention, as a core layer of the ubiquitous Transformer architecture, is the bottleneck for LLMs 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.

• The paper details a hardware-aware co-design that pipelines asynchronous MMA operations and emulates exponential functions to address GPU scaling bottlenecks.
• It integrates conditional softmax rescaling and TMEM-based tile scheduling to reduce memory traffic, warp divergence, and atomic contention.
• Empirical results demonstrate up to 1.3× to 2.7× performance gains over prior methods, underscoring the value of algorithm-hardware co-design for modern AI accelerators.

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

Motivation and Problem Formulation

The proliferation of Transformer-based models has generated substantial demand for efficient GPU-accelerated attention mechanisms, particularly as hardware architectures exhibit increasingly asymmetric scaling patterns. On NVIDIA Blackwell GPUs (B200, GB200), tensor core throughput has doubled relative to Hopper H100, while other resources such as shared memory bandwidth and exponential unit throughput scale more slowly or remain unchanged. This asymmetry creates new bottlenecks in attention workloads, fundamentally altering performance regimes and necessitating algorithm-hardware co-design.

FlashAttention-4 directly addresses these bottlenecks by redesigning the attention pipeline, exploiting Blackwell GPU features, and mitigating limiting factors outside of pure matrix multiplication (MMA). Key innovations include asynchronous MMA operation pipelining, software-emulated exponential computation, conditional softmax rescaling, and reductions in shared memory traffic and global atomic adds. All kernel development is realized in CuTe-DSL embedded in Python, which accelerates compilation (20–30×) and fosters greater extensibility.

Hardware-Aware Algorithm Design

Asynchronous MMA and Tile Scheduling

The introduction of tensor memory (TMEM) in Blackwell enables fully asynchronous MMA instructions. Where Hopper held MMA accumulators in registers (dramatically constraining tile sizes and register allocation), Blackwell's tensor cores write to TMEM directly, supporting larger tiles and more aggressive pipelining. FlashAttention-4 employs overlapping pipelines—while one tile executes tensor core operations, the other computes softmax—maximizing resource utilization.

Software-Emulated Exponential Functions

Softmax computation is a critical non-MMA bottleneck, limited by the MUFU's modest throughput (16 ops/clock/SM vs. 8192 ops/clock/SM for MMA). FlashAttention-4 addresses this by distributing exponential computations across both MUFU and floating-point FMA units via polynomial approximation. Range reduction (Cody-Waite) and bit manipulation yield efficient $2^x$ evaluation. Empirical tests demonstrate that a degree-3 polynomial matches MUFU hardware to within 1 BF16 ULP on 99% of inputs; BF16 quantization dominates the error profile for all polynomial degrees ≥3.

Conditional Softmax Rescaling

FlashAttention online softmax involves recursive renormalization steps for numerical stability, but rescaling is only necessary upon encountering new maxima. FlashAttention-4 introduces conditional rescaling, skipping unnecessary renormalizations and tolerating “slack” within thresholds, with a final normalization step ensuring exact outputs. This reduces vector multiplications, warp divergence, and register pressure without compromising correctness.

Backward Pass Pipeline and Memory Optimization

Pipeline Overlap and Resource Partitioning

Backward computation is dominated by five MMA operations (recomputing $S$, gradients of $QK$ and $PV$, and their respective partial derivatives). FlashAttention-4 exploits TMEM and aggressive tile partitioning to overlap MMA and non-MMA work, minimizing serialization and hiding softmax latency. The figure below illustrates the backward computational graph:

Figure 1: 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.

TMEM partitioning is carefully managed—the memory can only fit four tiles of $128 \times 128$, leading to sharing between $S$, $P$ and $dP$, $dS$, $dQ$. The pipeline is ordered for overlap and resource conservation.

2-CTA Mode: Traffic and Atomic Reduction

Blackwell’s 2-CTA MMA mode enables CTA pairs to jointly execute MMAs and share TMEM. In the backward $dQ$ step, DSMEM is used to exchange half of the $dS$ tile across CTAs, allowing each to process $(\frac{M}{2} \times 2N)$ operands with doubled reductions, halving the atomic operations. This alleviates shared memory contention and atomic bottlenecks.

Figure 2: 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 runs a CTA-pair UMMA with a doubled reduction.

Deterministic Backward Execution

To ensure reproducible gradient computation (essential for RL and debugging), FlashAttention-4 implements deterministic reduction using semaphore locks and careful CTA order scheduling. Swizzling over heads/batches and shortest-processing-time-first (SPT) scheduling is used to minimize stalls and achieve high deterministic throughput.

Scheduling and Framework Implementation

Longest-Processing-Time-First Scheduling

Work tiles for attention (especially with causal masking or variable sequence lengths) exhibit load imbalance. By sorting and processing tiles in LPT order (with batch as the outer dimension and head swizzling to minimize L2 cache thrashing), FlashAttention-4 achieves improved performance—empirical benchmarks show 4–8% FLOPS gains for MHA and up to 14% for MQA.

CuTe-DSL: Python-Based Kernel Assembly

The entire kernel suite is written in CuTe-DSL, providing CUTLASS-equivalent expressivity in Python and yielding 20–30× faster compile times than C++ templates. Primitive/flexible abstractions support rapid prototyping (e.g., block-sparse variants, FlexAttention), orthogonal features (masking, varlen, scheduling), and composable optimizations, lowering the entry barrier for GPU attention research.

Empirical Results

FlashAttention-4 achieves up to 1.3× speedup over cuDNN 9.13 and 2.7× over Triton for BF16 workloads on B200, reaching 1613 TFLOPs/s (71% theoretical peak). The pipeline and scheduling improvements result in robust performance across sequence lengths, head dimensions, and masking types. Deterministic backward passes retain up to 75% the speed of nondeterministic variants.

Implications and Future Directions

FlashAttention-4 exemplifies the necessity of algorithm-hardware co-design in response to rapidly evolving accelerators. As future GPU generations further increase MMA throughput relative to memory and non-matmul units, attention algorithms must integrate pipelined scheduling, operator emulation, and conditional computation. The Pythonic kernel framework foreshadows democratization of GPU programming and faster iteration cycles for algorithmic innovation.

Beyond Blackwell, analogous pipelining/memory optimization concepts may transfer to TPUs, IPUs, ASICs, and future neural accelerators. The trend toward hardware-aware algorithm adaptation will accelerate as device heterogeneity rises, and will influence all stages of AI model training and inference, especially in the context of long-sequence modeling and multimodal systems.

Conclusion

FlashAttention-4 provides robust, hardware-scaled attention kernels by co-designing algorithmic and kernel strategies with direct reference to hardware bottlenecks and constraints. The method leverages asynchronous computation, memory-aware scheduling, operator emulation, and conditional processing, yielding significant speedups and utilization on Blackwell GPUs. The modular, Python-based framework empowers rapid research and application deployment. These contributions will inform future attention kernel development as hardware asymmetry intensifies and deep learning workloads continue to scale (2603.05451).