The Evolution of FlashAttention: From Ampere to Blackwell

Abstract

The original FlashAttention algorithm removes a fundamental source of GPU memory traffic: it computes attention score tiles on chip and consumes them without materializing the full score or probability matrices in high-bandwidth memory. Once this cost is removed, different bottlenecks become visible. FlashAttention-2 improves work partitioning on Ampere; FlashAttention-3 overlaps data movement, matrix multiplication, and softmax on Hopper; and FlashAttention-4 adapts the pipeline to Blackwell, where tensor-core throughput has grown faster than exponential and shared-memory throughput. This article compares these implementations as a sequence of bottleneck shifts. It assumes the tiled algorithm and online-softmax recurrence developed in FlashAttention: Exact Attention Without Materialization.

1. The Invariant and the Moving Bottleneck

All numbered FlashAttention versions compute the same dense attention function:

Their shared invariant is that score and probability tiles remain on chip. A tile is produced, incorporated into the online-softmax state, and discarded. The versions differ principally in how that dataflow is assigned to thread blocks, warps, matrix units, memory engines, and on-chip storage.

This progression illustrates a common systems effect: removing one dominant cost does not end optimization. It reveals the next cost in the execution pipeline.

2. FlashAttention-1: Remove Quadratic HBM Traffic

The first implementation establishes the algorithmic foundation. It fuses score computation, masking, softmax, and value accumulation so the $N\times N$ score and probability matrices are never stored in HBM. An online-softmax recurrence permits each key/value tile to update an exact running result.

This changes the relevant optimization target from isolated matrix multiplications to the complete attention operator. The NeurIPS 2022 paper reports up to 7.6x kernel-level speedup over a PyTorch implementation, together with end-to-end improvements for BERT, GPT-2, and Long Range Arena workloads [1]. The implementation nevertheless reaches only a fraction of available matrix throughput because its work partitioning and synchronization leave hardware underused.

3. FlashAttention-2: Improve Parallel Work Partitioning

FlashAttention-2 retains the same non-materializing algorithm but reorganizes work at two levels.

First, it parallelizes over query blocks in addition to batch and attention heads. This matters when batch size and head count do not provide enough independent thread blocks to occupy all streaming multiprocessors. Independent query tiles supply another parallel dimension without requiring communication between thread blocks.

Second, it changes warp partitioning from sliced-K to sliced-Q. Under sliced-K, warps compute partial contributions for the same output rows and must exchange and reduce those contributions through shared memory. Under sliced-Q, each warp owns complete query rows while sharing the key/value tile. The output remains local to the owning warp, reducing synchronization and shared-memory traffic.

The implementation also delays output normalization until all key/value blocks have been processed. Carrying an unnormalized accumulator removes divisions and rescaling operations from the hot inner loop.

The paper reports approximately 2x speedup over FlashAttention-1, up to 230 TFLOPs/s on A100, and up to 73% of theoretical peak throughput in the forward pass. Reported end-to-end training reaches up to 225 TFLOPs/s per A100, corresponding to 72% model FLOPs utilization in the measured configuration [2].

4. FlashAttention-3: Overlap Hopper’s Asynchronous Engines

Hopper exposes hardware that changes how the kernel can be scheduled:

• The Tensor Memory Accelerator (TMA) moves multidimensional tiles between HBM and shared memory asynchronously.
• Warpgroup matrix multiply-accumulate (WGMMA) executes asynchronous tensor-core operations across a warpgroup.

FlashAttention-3 uses warp specialization to separate producer and consumer roles. Producer warps issue TMA transfers for future key/value tiles. Consumer warps execute WGMMA operations and softmax updates. Shared-memory buffers and barriers connect the stages.

A ping-pong schedule overlaps different resources: while one warpgroup performs matrix multiplication for one tile, another can process softmax for a previous tile. This is important because exponential throughput does not scale with tensor-core throughput. Executing the phases serially would leave one resource idle while the other works.

On H100, the paper reports up to 740 TFLOPs/s in FP16, approximately 75% of theoretical peak, and close to 1.2 PFLOPs/s in FP8. Relative to FlashAttention-2, it reports 1.5–2.0x faster forward execution and 1.5–1.75x faster backward execution [3].

5. FlashAttention-4: Adapt to Asymmetric Blackwell Scaling

On Blackwell, matrix throughput grows faster than several supporting resources. Exponential units and shared-memory bandwidth can therefore become limiting even though they represent a small fraction of the arithmetic operation count.

FlashAttention-4 uses Tensor Memory (TMEM), a dedicated on-chip accumulator store connected to the matrix units. Keeping accumulators in TMEM reduces pressure on general-purpose registers and permits a deeper pipeline. Cooperative two-CTA matrix operations distribute larger tiles across paired thread blocks.

The implementation also redistributes work around the newly exposed bottlenecks. It overlaps multiple query tiles, avoids unnecessary online-softmax rescaling when the running maximum is unchanged enough to permit the cheaper path, and can evaluate part of the exponential workload using otherwise available FMA resources. In the backward pass, layouts are selected to reduce shared-memory transposes and operand traffic.

For B200, the paper reports approximately 1.6 PFLOPs/s BF16, about 71% of peak throughput, as well as 1.3x speedup over cuDNN 9.13 and 2.7x speedup over a Triton implementation in its reported configurations [4]. These results depend strongly on sequence length, head dimension, masking, and precision.

6. Flash-Decoding: A Different Parallel Shape

Autoregressive decoding presents a different workload. Query length is one, while the KV-cache length may contain tens of thousands of tokens. Query-side parallelism therefore supplies little work.

Flash-Decoding partitions the key/value sequence instead. Multiple thread blocks independently compute partial maxima, normalization sums, and weighted outputs over different KV partitions. A final reduction merges those states using the same online-softmax identity as FlashAttention. The attention function is unchanged; the parallel decomposition changes to match the decode-time shape [5].

Flash-Decoding is best understood as a related branch rather than FlashAttention-3 or FlashAttention-4. The numbered versions primarily track architecture-specific schedules for attention, whereas Flash-Decoding addresses insufficient parallelism in single-query inference.

7. Discussion

The history of FlashAttention separates algorithm from schedule. Non-materialization and online softmax define the stable algorithmic core. Loop ordering, warp ownership, asynchronous copies, accumulator placement, and pipeline depth are schedules chosen for a particular hardware balance.

The optimization sequence is therefore:

1. Eliminate unnecessary HBM materialization.
2. Expose sufficient independent work to occupy the GPU.
3. Overlap specialized engines so no phase serializes the pipeline.
4. Rebalance the schedule when matrix throughput outgrows supporting resources.

This framing is more durable than treating the versions as independent algorithms. A future GPU may require another schedule, but it can preserve the same non-materializing computation.

8. Limitations

• Cross-version throughput is not directly comparable unless hardware, precision, shapes, and masking are controlled.
• Reported figures are best-case measurements from the corresponding papers, not universal end-to-end speedups.
• Architecture-specific kernels require continued tuning and may not transfer efficiently to other GPU families.
• FlashAttention still performs $\Theta(N^2d)$ arithmetic for dense attention; scheduling improvements do not remove quadratic compute.

References

1. Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. NeurIPS 2022.
2. Tri Dao. FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning. 2023.
3. Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, and Tri Dao. FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision. 2024.
4. A. Zadouri et al. FlashAttention-4: Algorithm and Kernel Pipelining Co-Design for Asymmetric Hardware Scaling. 2026.
5. Tri Dao, Daniel Haziza, Francisco Massa, and Grigory Sizov. Flash-Decoding for Long-Context Inference. 2023.