PaperlayerView exact revision on arXiv

Source-linked AI summary

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

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

arXiv:2603.05451v1cs.CL

TL;DR

Attention becomes increasingly constrained by shared-memory traffic and non-matmul operations as Blackwell tensor-core throughput outpaces other functional units. FlashAttention-4 co-designs algorithms and kernels to mitigate these bottlenecks, achieving up to 1.3× speedup over cuDNN and 2.7× over Triton on B200 GPUs with BF16.

  • Problem

    Blackwell’s asymmetric scaling shifts attention’s bottleneck from matrix multiplication toward shared-memory traffic and non-matmul operations such as softmax.

  • Method

    FlashAttention-4 co-designs asynchronous, larger-tiled pipelines and software techniques with tensor-memory and 2-CTA mechanisms to reduce non-matmul work, shared-memory traffic, and global atomic adds.

  • Results

    Up to 1.3× speedup over cuDNN and 2.7× over Triton is achieved for BF16 attention on B200 GPUs.

  • Takeaways & Limitations

    FlashAttention-4 reaches up to 1613 TFLOPs/s, or approximately 71% of B200’s theoretical maximum, while supporting extensible attention variants in Python.

  • Takeaways & Limitations

    The roofline analysis is simplified and excludes resources such as floating-point math, register bandwidth, and L2 bandwidth.

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. 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.

1 Introduction

FlashAttention-4 co-designs attention algorithms and GPU kernels for Blackwell’s asymmetric bottlenecks, where non-matmul resources increasingly limit performance. It introduces asynchronous pipelines, softmax optimizations, memory-traffic reductions, and a Python-based implementation, achieving substantial B200 speedups.

  • Motivation: Blackwell doubles tensor-core throughput while shared-memory bandwidth and exponential units scale more slowly, shifting attention bottlenecks to non-matmul resources.For typical attention workloads, shared-memory traffic and exponential operations can exceed MMA compute by 25–60%.
  • Method: FlashAttention-4 redesigns forward and backward pipelines to overlap asynchronous MMA operations, softmax computation, and memory operations on Blackwell.The pipelines exploit larger tile sizes and fully asynchronous MMA operations.
  • Method: Software-emulated exponentials and conditional softmax rescaling reduce the forward pass’s exponential-unit bottleneck and unnecessary rescaling work.Polynomial approximation uses FMA units to increase exponential throughput.
  • Method: Tensor memory and 2-CTA MMA reduce backward-pass shared-memory traffic and halve the number of atomic reductions.The design also supports deterministic execution with minimal performance overhead for reproducible reinforcement-learning training.
  • Framework: 20–30× faster compile times than traditional C++ template-based approaches are achieved with a fully CuTe-DSL implementation embedded in Python.The framework maintains full expressivity while lowering the barrier to developing attention variants.
  • Results: Up to 1.3× speedup over cuDNN and 2.7× over Triton is achieved on B200 GPUs with BF16, reaching approximately 1600 TFLOPS/s at 71% of the theoretical maximum.For large sequence lengths, FlashAttention-4 outperforms alternative attention implementations.

2 Background

The background describes attention computation and the Blackwell execution model, emphasizing hardware changes that shift bottlenecks away from matrix multiplication. Blackwell adds tensor memory, larger asynchronous MMA tiles, and cooperative 2-CTA execution to support new kernel designs.

  • 2.1 Multi-Head Attention: Attention computes row-wise softmax probabilities from query–key scores and applies them to values, with numerical stabilization by subtracting each row maximum.For multi-head attention, each head has separate projections and computation parallelizes across heads and batches.
  • 2.1 Multi-Head Attention: The row-wise softmax gradient is ds = (diag(p) − pp^T)dp for p = softmax(s).This expression describes the gradient transformation used in the backward pass.
  • 2.2 GPU Hardware Characteristics and Execution Model: Blackwell introduces 256 KB of tensor memory per SM for intermediate tensor-core results, enabling direct asynchronous MMA writes without consuming registers.TMEM requires explicit programmer management and supports larger tile sizes by reducing register pressure.
  • 2.2 GPU Hardware Characteristics and Execution Model: Blackwell MMAs process larger 128 × N tiles and write outputs asynchronously to tensor memory, unlike Hopper’s register-based MMA outputs.The architecture supports warp-specialized kernels that separate producer and consumer roles.
  • 2.2 GPU Hardware Characteristics and Execution Model: 2-CTA tensor-core mode lets paired CTAs cooperatively execute one MMA while partitioning operand and accumulator tiles across the pair.Compared with single-CTA MMAs, paired execution supports M = 128 or 256.
  • 2.2 GPU Hardware Characteristics and Execution Model: Blackwell doubles FP16/BF16 tensor-core throughput versus Hopper, while shared-memory bandwidth and exponential throughput remain unchanged or scale more slowly.This imbalance shifts performance pressure toward shared-memory traffic and non-matmul operations such as softmax.

3 Algorithm

FlashAttention-4 uses roofline-guided algorithm and kernel co-design to address Blackwell’s bottlenecks in shared-memory traffic, exponential operations, and overlap. Its forward pipeline overlaps tensor-core computation with softmax while polynomial exponential emulation and revised memory scheduling reduce non-matmul costs.

  • Feeds and Speeds: Roofline analysis identifies shared-memory traffic and exponential operations as important bottlenecks alongside MMA computation for Blackwell attention.For M = N = d = 128, shared memory takes 768 cycles versus 1024 cycles for both MMA compute and exponential operations; larger tiles increase all costs.
  • New pipeline to overlap matmul and softmax: The forward pipeline uses larger tiles and ping-pong scheduling to overlap tensor-core MMAs, softmax computation, and memory operations.Two output tiles are computed per thread block while one tile performs tensor-core operations and the other computes softmax.
  • New pipeline to overlap matmul and softmax: Two softmax warpgroups process complete rows by finding maxima, rescaling, exponentiating, converting precision, and computing row sums.Each Q tile contains 128 query tokens, and each thread processes an entire row.
  • New pipeline to overlap matmul and softmax: Tensor-memory transfer decouples output rescaling into a correction warpgroup, removing that work from the critical path.The design uses tensor-memory partitionings that allocate space for output, S, and P tiles to support pipeline overlap.
  • Emulation of the exponential function: Table 2 evaluates polynomial 2^x emulation against FP64 references using FP32 and BF16 error after rounding.The caption states that BF16 quantization error dominates for polynomial degrees ≥3.

4. Evaluate polynomial to get 2xfrac

FlashAttention-4 addresses Blackwell’s non-MMA bottlenecks by combining asynchronous pipelining, partial exponential emulation, conditional softmax rescaling, and 2-CTA backward execution. These changes target exponential throughput, shared-memory traffic, and atomic reductions while preserving numerical accuracy.

  • Exponential evaluation: Partial polynomial exponential emulation distributes work across FMA and MUFU units, but applies only to 10–25% of softmax entries to limit register pressure.The remaining entries use hardware MUFU.EX2, with the fraction tuned to the ratio of MMA and exponential throughput.
  • Exponential evaluation: Degree-3 polynomial emulation has 8.8 × 10−5 maximum FP32 relative error, while matching hardware within 1 BF16 ULP on 99% of inputs.After BF16 rounding, quantization error of approximately 3.9 × 10−3 dominates polynomial error for degrees ≥3.
  • Softmax rescaling: Conditional softmax rescaling skips updates when the running maximum increases by at most τ, typically τ = log2(256) = 8.0.Final normalization by the true maximum and normalizer preserves correctness while reducing vector rescaling operations.
  • Bottlenecks: Blackwell’s shared-memory traffic and exponential operations can dominate attention execution despite doubled tensor-core throughput.For the typical backward tile M = N = d = 128, shared-memory traffic takes 3328 cycles versus 2560 MMA cycles and 1024 exponential-unit cycles.
  • Pipeline design: FlashAttention-4 overlaps tensor-core, softmax, and memory work through redesigned pipelines using asynchronous MMA operations and larger tiles.The backward pipeline uses tensor memory and schedules dQ and dK MMAs from the previous iteration to overlap softmax latency.
  • Backward pass: The 2-CTA backward mode roughly halves operand-B shared-memory traffic and reduces global dQ atomic reductions by half.The paired CTAs partition the output tile, exchange dS through DSMEM, and each writes only its own dQ slice.

4 Language and Framework

FlashAttention-4 uses CuTe-DSL embedded in Python to retain low-level GPU expressivity while substantially reducing compilation time and improving extensibility for attention variants.

  • Language and framework: FlashAttention-4 is implemented in Python-embedded CuTe-DSL, which lowers Python source to PTX and then SASS without CUDA C++ components.The programming model is isomorphic to CUTLASS C++ and provides direct PTX access for functionality not yet exposed in the APIs.
  • Fast compilation through JIT: 20-30× faster compile times than FlashAttention-3 reduce iteration time for experimentation and debugging during kernel development.Table 4 compares single-kernel compilation for FA3’s C++ templates and FA4’s CuTe-DSL; earlier versions typically precompile hundreds of kernels for attention variants.
  • Flexibility and accessibility: Developers have built FlexAttention and block-sparse attention variants on top of FlashAttention-4 without modifying its core framework.The Python-based framework is intended to lower the barrier to entry for contributors without deep C++ template-metaprogramming expertise.
  • Composable primitives: FlashAttention-4 exposes masking, block-sparse patterns, variable sequence lengths, and scheduling as composable primitives for constructing attention variants.These primitives are designed to be combined independently rather than reimplemented from scratch for each variant.

5 Empirical Evaluation

FlashAttention-4 is evaluated on B200 GPUs across sequence lengths, head dimensions, masking settings, and forward and backward passes, with strong speedups over major implementations.

  • Evaluation scope: FlashAttention-3 does not run on B200, so it is not included as a directly executable baseline in this evaluation.This scope boundary applies specifically to the reported B200 benchmark comparison.
  • Forward pass: Figure 4 compares forward TFLOPS for non-causal and causal attention with head dimension 128 across sequence lengths.The figure reports 1.1-1.3× speedup over cuDNN 9.13.0 and 2.1-2.7× over Triton; newer cuDNN versions incorporate many described techniques and achieve similar performance.
  • Forward pass: For medium and long sequences of 4k tokens or more, FlashAttention-4 consistently outperforms baselines across head dimensions and causal-masking settings.The reported forward-pass gains are larger for causal attention, which the authors attribute to the longest-processing-time-first scheduler.
  • Backward pass: FlashAttention-4 achieves consistent backward-pass speedups across long sequence lengths and causal masking using its 2-CTA backward pass.Figure 6 reports backward-pass TFLOPS for non-causal and causal attention with head dimension 128.
  • Deterministic backward pass: The deterministic backward pass reaches up to 75% of the speed of the nondeterministic backward pass of the 1-CTA implementation.The authors attribute this result to careful swizzling and scheduling.

6 Discussion and Conclusion

FlashAttention-4 co-designs attention algorithms and kernels around Blackwell’s shifted bottlenecks, while its Python-based implementation preserves low-level control and shortens compilation.

  • Discussion and conclusion: Blackwell’s faster tensor cores shift attention bottlenecks toward shared-memory traffic and exponential throughput, motivating algorithmic and kernel co-design.The paper redesigns pipelines around asynchronous MMA, software-emulated exponential operations, and conditional softmax rescaling.
  • Discussion and conclusion: FlashAttention-4 uses tensor memory and 2-CTA MMA mode to reduce shared-memory traffic and atomic reductions in the backward pass.Its implementation also uses larger-tiled matmuls and asynchronous overlap between softmax and matrix multiplication.
  • Discussion and conclusion: Implemented entirely in Python-embedded CuTe-DSL, FlashAttention-4 preserves low-level control while achieving 20-30× faster compilation than C++ template-based kernels.Although optimized for Blackwell GPUs, the authors state that some algorithms may extend to other accelerators as compute outpaces non-matmul units.

A.1 System and libraries

The evaluation measures averaged B200 runtimes using repeated warmup and timing runs, with library versions selected from the March 2025 environment and explicit cuDNN comparisons.

  • System and libraries: Benchmarks warm up with 5 runs, repeat timing 10 times, and report the average timing.The system used is identified as a B100 180GB SXM6 with 1000W in the supplied benchmark-settings passage.
  • System and libraries: The libraries generally use the latest versions available at the time of writing, identified as March 2025.This establishes the reported software environment rather than a timeless version comparison.
  • System and libraries: The paper compares against cuDNN 9.13 and cuDNN 9.19.1.2, after collaborating with the cuDNN team to incorporate some FlashAttention-4 techniques from versions 9.13 and 9.14 onward.The comparison therefore includes both the main-paper cuDNN version and a later version incorporating some described techniques.

A.2 Backward Deterministic non-causal

This section supplements the deterministic backward-kernel evaluation with results for the non-causal setting, alongside causal masking results.

  • The deterministic backward kernel is evaluated without causal masking.
  • Figure 8 includes performance numbers for the non-causal deterministic backward kernel.
  • The non-causal results are presented side by side with results using causal masking.
Loading 2603.05451v1…