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, then examines Flash-Decoding, which is not a fifth schedule but a different parallel decomposition chosen for single-query inference. 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:

O=softmax⁡(QK⊤)V. O=\operatorname{softmax}(QK^\top)V.

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.

VersionPrimary targetBottleneck addressedPrincipal change
v1AmpereQuadratic HBM trafficFused tiling and online softmax
v2AmpereWork partitioning and occupancyQuery-side parallelism and sliced-Q warps
v3HopperSerial movement, GEMM, and softmax phasesTMA, WGMMA, warp specialization, and overlap
v4BlackwellExponential and shared-memory throughput lag tensor coresTMEM, cooperative MMA, and deeper pipelines
An alternating chain: quadratic HBM traffic removed by v1, exposing idle SMs addressed by v2, exposing serialized phases addressed by v3, exposing exponential and shared-memory limits addressed by v4.
Figure 1. Each version removes the cost above it and exposes the one below. Red is whatever dominates the kernel at that moment; green is the scheduling change that removes it. Nothing in the green boxes alters the function being computed.

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

1.1 Before comparing versions, identify the kernel

The version names describe schedules, and a schedule only exists once it is compiled for a particular architecture. On hardware newer than the newest kernel a library ships, the code that runs is not the version whose name appears in the API — it is an older kernel executing through forward compatibility.

This is not a hypothetical. On an SM 10.3 GPU, the flash-attn package’s kernels are compiled for SM80, and what executes is the FlashAttention-2 CUTLASS kernel with the split-KV (Flash-Decoding) variant selected. Meanwhile the vendor math library ships a native kernel for the same architecture, and PyTorch’s own SDPA “flash” backend resolves to the same SM80 kernel through a different binding.

Three naming facts follow, each of which can quietly corrupt a version comparison:

  • The same kernel appears under two names. Forcing SDPA’s flash backend and calling the flash-attn package produce identical numerics — the same maximum error to every digit — because they are the same kernel behind different Python wrappers. A comparison between them measures the wrappers.
  • A newer version number can be slower than an older one. The SM80 kernel is roughly 6× slower than the native SM100 kernel at long sequence lengths, despite the former carrying the FlashAttention name.
  • The kernel a library selects can depend on the head dimension. At dh=32d_h=32 the same family runs 4.7× faster than at dh=256d_h=256, because the tile budget shifts from length to depth.

The practical rule for anything in this article: before comparing two implementations, record the executed kernel name from a profiler and confirm it is the one intended. A benchmark that cannot name its kernel has not established what it measured.

2. FlashAttention-1: The Baseline That Leaves Work Unused

The first implementation establishes the algorithmic foundation described in the previous article: fused tiling, score tiles consumed on chip, and an exact online-softmax recurrence that never materializes the N×NN\times N matrices. What that design achieved, and on which shapes, is reported there and is not repeated here.

For this article, what matters is what version 1 left on the table. Three limits follow directly from its schedule:

  • Parallelism came from batch size and head count alone. When those two do not supply enough independent thread blocks, streaming multiprocessors sit idle.
  • Every warp contributed a partial result to the same output rows, so completing a row required a shared-memory reduction and a barrier.
  • Output normalization ran inside the inner loop, adding divisions and rescaling to the path executed once per key/value tile.

Version 2 addresses the first two directly; the third falls out of the same reorganization.

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.

Under sliced-K four warps each own a key slice and must reduce their partial rows through shared memory; under sliced-Q each warp owns a distinct block of query rows and keeps its output local.
Figure 2. The partitioning change behind most of v2's gain. Splitting the key dimension makes every warp a partial contributor to the same output rows, so a shared-memory reduction and a barrier stand between the warps and their result. Splitting the query dimension gives each warp rows nobody else writes; the shared K/V tile is read-only, so no reduction is needed at all.

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.

Two schedules on three engine rows: in the first, copy, matrix multiply and softmax occupy one time slot each in turn; in the second, all three rows are busy simultaneously on different tiles.
Figure 3. Why warp specialization matters more on Hopper than it did on Ampere. In the serial schedule each engine is idle two thirds of the time, and the kernel runs at the speed of the sum of the phases. Once producer and consumer warps are separated, the three engines work on different tiles at the same time and the kernel runs at the speed of the slowest one.

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.

That separation is testable, and the test is unusually clean: if the versions differ only in schedule, every one of them must produce the same numbers. Measuring five independent implementations — three fused kernels, one conventional reference, and one serving engine — against a float64 reference at a sequence length of 4,096 gives a maximum absolute error of 7.554e-03 for all five, identical to the last digit. They differ in latency by a factor of 20. The function is fixed; only the execution changes.

The corollary is a limitation on what a single machine can demonstrate. Comparing versions v1 through v4 requires kernels for four GPU generations, and a host with one GPU generation provides at most the kernels that run on it — plus whatever older kernels execute through compatibility. A version-by-version benchmark therefore cannot be produced on arbitrary hardware, and a table of throughput numbers across versions is usually a table of different GPUs, not different schedules. The comparison in this article is between schedules as published, each measured on the hardware it was written for; the reproduction of the underlying invariant is what transfers.

Horizontal bars of reported kernel throughput: 230 TFLOP/s on A100, 740 on H100 in FP16, 1200 on H100 in FP8, and 1600 on B200, each annotated with the fraction of that GPU's peak.
Figure 4. Reported best-case kernel throughput, each version on the hardware it was written for. The bars are not a like-for-like comparison — the GPU changes with every row — and that is the point: absolute throughput rises roughly sevenfold while the fraction of peak stays near three quarters. The schedule is what keeps that fraction constant as the hardware moves.

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 Θ(N2d)\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.