FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning

ArXiv: 2307.08691

🎯 Pitch

Despite FlashAttention’s initial 2-4× speedup, it still leaves half the GPU’s compute idle, eking out only 25-40% of peak FLOPs. FlashAttention-2 nearly doubles performance again to 73% of peak—matching optimized matrix-multiply efficiency—by reshuffling how work is carved up across thread blocks and warps, bringing long-context training to 225 TFLOPs/s on an A100 without any approximation loss.

---

1. Executive Summary

This paper analyzes how to further accelerate exact attention computation on GPUs by addressing suboptimal work partitioning in the original FlashAttention algorithm, using standard attention benchmarks and GPT-style model training on A100 GPUs. FlashAttention-2 introduces three improvements: a tweaked algorithm that reduces non-matmul FLOPs (by maintaining an unscaled output accumulator until the final rescaling step), increased parallelism by partitioning work across the sequence length dimension in addition to batch and head dimensions (scheduling separate thread blocks for row blocks in the forward pass and column blocks in the backward pass), and a revised warp-level work partitioning strategy that splits Q across warps instead of K and V to eliminate inter-warp communication through shared memory. These yield approximately 2× speedup over FlashAttention in both forward and backward passes, reaching up to 73% of theoretical maximum throughput on A100 GPUs (230 TFLOPs/s forward, 63% backward) and up to 225 TFLOPs/s in end-to-end training (72% model FLOPs utilization) — establishing that careful parallelism and work partitioning can close most of the remaining efficiency gap between attention and optimized matrix-multiply, though the gains depend on sequence length and head dimension with the largest relative improvements seen at training-scale batch sizes where the original FlashAttention suffered from low GPU occupancy.

2. Context and Motivation

The Core Problem: FlashAttention Left Substantial GPU Efficiency on the Table

The fundamental problem this paper addresses is deceptively specific: FlashAttention [5], despite being the state-of-the-art exact attention implementation, still runs at only 25–40% of the GPU's theoretical maximum FLOPs/s. This is a striking gap. Optimized matrix multiplication (GEMM) routinely reaches 80–90% of theoretical peak throughput on modern GPUs because decades of engineering effort have gone into maximizing utilization of Tensor Cores and minimizing data movement. FlashAttention, by contrast, achieves a 2–4× wall-clock speedup over standard attention implementations but leaves more than half of the GPU's compute capacity idle.

Why does this matter? The paper frames the significance through the lens of what practitioners could do with longer context lengths but currently cannot afford to do. The introduction catalogs recent models pushing context lengths far beyond the traditional 2k token limit — GPT-4 at 32k, MosaicML's MPT at 65k, Anthropic's Claude at 100k — and emerging applications like long document querying and story writing that demand such lengths. But these context extensions come at enormous computational cost. Even with FlashAttention's 2–4× improvement, training models on 16k or 32k context sequences remains prohibitively expensive for most research teams. Doubling the attention speed again — closing the gap between attention and GEMM efficiency — effectively halves the cost of long-context training, or equivalently, allows training on sequences twice as long for the same budget.

The paper is not proposing a new approximation or sparsification method. It is asking a purely engineering question: given that FlashAttention's I/O-aware tiling strategy is theoretically sound, why is the GPU implementation so far from its hardware limits, and can we fix that? This is a different kind of contribution than the original FlashAttention paper. FlashAttention introduced a novel algorithm (online softmax with tiling and recomputation) that solved the memory bottleneck. FlashAttention-2 introduces no algorithmic novelty in the mathematical sense — the forward and backward passes compute the exact same quantities — but rather addresses the parallelism and work partitioning inefficiencies that became the new bottleneck after the memory bottleneck was removed.

The Specific Inefficiency: Suboptimal Work Partitioning Across the GPU's Parallelism Hierarchy

Through profiling, the authors identify that FlashAttention's inefficiency stems from how it maps the attention computation onto the GPU's execution model. To understand this, one must understand how GPUs organize computation (Section 2.1):

• Thread blocks are the coarsest scheduling unit. Each thread block runs on a streaming multiprocessor (SM). An A100 has 108 SMs, so to fully utilize the GPU, one needs at least ~108 thread blocks ready to execute simultaneously. This is called occupancy — the fraction of GPU resources actually doing useful work.
• Warps are groups of 32 threads within a thread block. Warps within the same block can communicate through shared memory (on-chip SRAM) and synchronize.
• Shared memory is fast (~19 TB/s on A100) but limited (192 KB per SM). Reading/writing shared memory is much faster than accessing HBM (~1.5–2.0 TB/s), but it is still communication overhead — the fastest operations are those that stay entirely within a warp's registers.

FlashAttention made two specific design choices at each level of this hierarchy that FlashAttention-2 revises:

At the thread block level: no parallelization over sequence length. FlashAttention parallelizes only over the batch dimension and the number of heads dimension. Each thread block handles exactly one attention head for one batch element. This means the total number of thread blocks is batch_size × number_of_heads. When this product is small — which happens when sequences are long and batch size is correspondingly reduced to fit in memory — occupancy drops well below 108. SMs sit idle. This is the "low-occupancy" problem the abstract refers to.

Consider a concrete scenario: training a GPT-style model with 16 heads on sequence length 8k. Memory constraints might force a batch size of 1 or 2, yielding only 16–32 thread blocks — far fewer than the 108 SMs available. Half to two-thirds of the GPU's compute capacity is wasted.

At the warp level: the "split-K" scheme forces inter-warp communication. Within each thread block, FlashAttention partitions K and V across warps (typically 4 warps per block) while keeping Q accessible by all warps. Each warp computes a partial QK^T for its slice of K, then must compute a partial output for its slice of V. But the softmax normalization couples all these partial results — each warp's partial output depends on the global row-wise max and sum of exponentials, which requires data from all warps. FlashAttention's solution: each warp writes its intermediate results to shared memory, all warps synchronize, and then they collectively compute the final output. This shared memory traffic and barrier synchronization adds latency to every iteration of the inner loop.

The abstract characterizes this as "unnecessary shared memory reads/writes." The "necessity" depends on how the work is partitioned. If the partitioning were designed differently — as FlashAttention-2 demonstrates — some of this communication can be eliminated entirely.

The Non-Matmul FLOP Problem: A Deeper Hardware Asymmetry

The paper identifies a third source of inefficiency that is more subtle but equally important: the extreme throughput asymmetry between matrix multiply operations and all other floating-point operations on modern GPUs (Section 3.1).

The A100 GPU has a theoretical peak of 312 TFLOPs/s for FP16/BF16 matrix multiply (using Tensor Cores) but only 19.5 TFLOPs/s for non-matmul FP32 operations. That is a 16× difference. A single non-matmul FLOP costs as much compute time as 16 matmul FLOPs. This asymmetry is not a bug — it reflects deliberate hardware design where specialized matrix multiply units (Tensor Cores) are optimized for the dominant operation in deep learning, while general-purpose floating-point units handle the rest at much lower throughput.

FlashAttention's algorithm, while reducing total memory traffic, involves non-matmul operations that eat into this scarce general-purpose compute budget: computing row-wise maxima, subtracting maxima, exponentiating, summing exponentials, and rescaling outputs. Each of these element-wise operations runs at 1/16 the speed of the matrix multiplies that compute QK^T and PV. Even though these non-matmul FLOPs are a small fraction of the count of operations, they are a large fraction of the time because of the 16× throughput penalty.

The paper frames reducing non-matmul FLOPs as essential to pushing attention closer to GEMM-like efficiency. An optimized GEMM kernel spends almost all its time in Tensor Core operations. To get attention to approach that, the implementation needs to minimize the number of instructions that execute on the general-purpose units — even if that means slightly restructuring the algorithm.

What Prior Approaches Missed

The original FlashAttention (Dao et al., 2022). FlashAttention solved what was then the primary bottleneck: the $O(N^2)$ memory requirement of materializing the attention matrix S and the softmax-normalized matrix P. By applying tiling (loading blocks of Q, K, V from HBM to SRAM) and recomputation (recomputing S and P during the backward pass rather than storing them), FlashAttention reduced memory from quadratic to linear in sequence length and achieved 2–4× wall-clock speedup. Its contribution was I/O-awareness — designing the algorithm to minimize data movement between HBM and SRAM, the slowest link in the memory hierarchy.

But FlashAttention was designed to solve the memory bottleneck, not the compute bottleneck. Once the memory problem was addressed, the compute utilization problem became visible. FlashAttention-2 is essentially the continuation of that optimization trajectory, now targeting the next limiting factor.

Triton's FlashAttention implementation (Tillet et al.). The Triton compiler framework includes an implementation of FlashAttention that introduced two ideas that directly motivated FlashAttention-2 (Section 3.2, footnote 3): swapping the order of the nested loops (outer loop over row blocks, inner loop over column blocks, rather than the reverse in the original FlashAttention paper) and parallelizing over the sequence length dimension. The paper explicitly credits Phil Tillet with these innovations. FlashAttention-2 adopts both ideas and combines them with its own warp-level work partitioning changes. The Triton implementation serves as an important baseline in the experiments (Figures 4–6), and FlashAttention-2 consistently outperforms it, particularly in the backward pass (around 2× faster).

xformers (Lefaudeux et al., 2022). The xformers library provides a "cutlass" implementation of FlashAttention that serves as another baseline. FlashAttention-2 is around 2× faster than this implementation as well.

Approximate and sparse attention methods. The paper briefly acknowledges the extensive literature on attention approximations: Longformer [2], Performer [4], Reformer [9], Linformer [19], Big Bird [20], Scatterbrain [3], and linear attention variants [8]. These methods reduce the asymptotic complexity of attention below $O(N^2)$ by imposing structural constraints (sparsity patterns, low-rank approximations, kernel-based approximations). The paper's positioning here is significant: it notes that "as far as we know, most large-scale training runs still use standard attention." This is a pointed observation — despite years of research on efficient attention approximations, exact attention remains the default in production training pipelines. Why? Approximate methods either sacrifice model quality (the approximation degrades performance on some tasks) or introduce implementation complexity that makes them difficult to integrate into existing training frameworks. FlashAttention (and FlashAttention-2 by extension) sidesteps this tradeoff entirely: it computes exact attention, with no approximation, so there is no quality cost to pay. The value proposition is pure speed with zero accuracy degradation.

Multi-query and grouped-query attention (Shazeer, 2019; Ainslie et al., 2023). These are orthogonal optimizations that reduce the size of the KV cache during inference by having multiple query heads share the same key/value head. FlashAttention-2 supports these variants (Section 3.1.2) by implicitly manipulating head indices and summing gradients across duplicated heads in the backward pass. These methods address the inference memory bottleneck (KV cache size) rather than the training compute bottleneck, but FlashAttention-2's speedup applies during both training and inference, making it complementary.

How FlashAttention-2 Positions Itself

The paper's positioning is unusually precise for a systems paper: it makes no claim of algorithmic novelty in the mathematical sense and explicitly credits prior work (Triton) for two of its three main ideas. Instead, it positions itself as identifying and systematically addressing the specific implementation inefficiencies that remained in FlashAttention after its initial success.

The three contributions map directly to three distinct levels of the GPU parallelism hierarchy, which gives the paper a clear decompositional structure:

1. Algorithm level (Section 3.1): Reduce non-matmul FLOPs to spend more time in Tensor Cores. This is a numerical optimization — restructure the computation so the same mathematical result is achieved with fewer element-wise operations.

2. Thread block scheduling level (Section 3.2): Increase occupancy by parallelizing over the sequence length dimension. This is a coarse-grained parallelism optimization — ensure enough thread blocks exist to keep all SMs busy, particularly when batch size is small.

3. Warp-level work partitioning (Section 3.3): Eliminate inter-warp shared memory communication by rethinking which tensors are split across warps. This is a fine-grained parallelism optimization — reduce the synchronization and data movement overhead within each thread block.

This three-level decomposition mirrors how GPU kernel engineers think about optimization: you tune the algorithm to minimize expensive operations, you ensure sufficient parallelism to saturate the hardware, and then you optimize the data flow within the innermost parallel units.

The paper's title — "Better Parallelism and Work Partitioning" — accurately reflects that the core contribution is not a new attention algorithm but rather a superior mapping of the existing algorithm onto GPU hardware. This is an engineering contribution, but one with substantial practical impact: 2× speedup on the most widely used exact attention implementation, deployed in production training pipelines. The paper's evaluation strategy reflects this pragmatism — rather than reporting accuracy on downstream tasks (which should be identical to FlashAttention since the computation is exact), it reports raw throughput (TFLOPs/s) and end-to-end training speed, the metrics that practitioners actually care about when deciding which kernel to use.

3. Technical Approach

3.1 Reader Orientation

FlashAttention-2 is a GPU kernel implementation — a program that runs directly on the graphics processor — that computes exact attention (the same mathematical operation as standard softmax attention) significantly faster than its predecessor by restructuring how the computation is distributed across the GPU's parallelism hierarchy. The problem it solves is that the original FlashAttention, despite its I/O-aware tiling strategy, still leaves 60–75% of the GPU's theoretical compute capacity unused because of poor work partitioning: too few thread blocks to occupy all streaming multiprocessors when batch sizes are small, and unnecessary shared memory communication between warps within each thread block. The solution is a three-level redesign: at the algorithm level, reduce the number of element-wise (non-matrix-multiply) operations that execute on the GPU's slow general-purpose units; at the thread block scheduling level, parallelize work across the sequence length dimension to keep all 108 SMs busy even with tiny batch sizes; and at the warp level, split Q instead of K and V across warps so that each warp can independently compute its portion of the output without needing to synchronize through shared memory.

3.2 Big-Picture Architecture (Diagram in Words)

The system has four major components, organized as a pipeline that transforms input tensors Q, K, V into output O (forward pass) and then transforms incoming gradients dO into parameter gradients dQ, dK, dV (backward pass):

1. Input partitioner (thread block scheduler): Decides how to divide the N×d input matrices Q, K, V into smaller blocks that fit in on-chip SRAM (192 KB per SM on A100). For the forward pass, it partitions the output sequence length N into Tr row blocks of size Br each, and the key/value sequence length N into Tc column blocks of size Bc each. For the backward pass, the partitioning scheme flips: column blocks become the outer loop, and row blocks the inner loop. This component also decides which thread block handles which partition — now additionally parallelizing over the sequence length dimension, not just batch and heads.

2. Block-wise attention computer (the inner loops): For each (row block, column block) pair, it loads the relevant Q, K, V tiles from HBM into SRAM, computes the local attention scores S = Q_i K_j^T, applies online softmax with running statistics (maintaining a running max and an unscaled output accumulator to avoid redundant rescaling), multiplies with V_j, and updates the running output. Crucially, the output is only rescaled once — at the very end — rather than after every block iteration.

3. Warp-level work partitioner: Within each thread block (which typically uses 4 or 8 warps, each being a group of 32 threads), this component decides which warps handle which slices of the computation. In FlashAttention-2's forward pass, Q is split across warps while K and V are broadcast to all warps — the reverse of FlashAttention's "split-K" scheme. Each warp independently computes its portion of QK^T and its portion of the output without needing to write intermediate results to shared memory for other warps to read.

4. Gradient accumulator (backward pass only): Handles the one remaining inter-thread-block communication: the gradient dQ is updated by contributions from multiple column blocks (since dQ_i gets contributions from every column block j via dS_i^{(j)} K_j). This component uses atomic add operations in HBM to safely combine contributions from different thread blocks processing different column blocks. Within each thread block, the backward pass applies the same warp-partitioning logic to avoid inter-warp shared memory traffic for dK and dV accumulation.

Information flows as follows: Q, K, V, (and dO, O, L for backward) are loaded from HBM → partitioned into tiles → each tile pair is loaded into SRAM by the assigned thread block → within each thread block, warps independently compute their matrix multiplies and element-wise operations using registers → the final scaled output (or gradients) are written back to HBM → in the backward pass, atomic adds combine dQ updates from different column-processing thread blocks.

3.3 Roadmap for the Deep Dive

• First, the forward pass algorithm tweaks (Section 3.1.1): The reduction in non-matmul FLOPs requires understanding the online softmax trick in detail — how FlashAttention-2 maintains an unscaled accumulator and defers the final division to the end, and why this eliminates one rescaling operation per inner loop iteration. This is the foundation because all parallelism decisions (Sections 3.2 and 3.3) operate on this restructured algorithm.

• Second, the backward pass differences (Section 3.1.2): The backward pass eliminates redundant storage of both the max and the sum of exponentials, keeping only the logsumexp. This simplification reduces memory traffic and simplifies the warp-partitioning design for the backward pass.

• Third, the sequence-length parallelization (Section 3.2): Explains why the outer loop choice (rows vs columns) determines which dimension can be parallelized, how this maps to the forward and backward passes differently, and why atomic adds are needed for dQ in the backward pass.

• Fourth, the warp-level work partitioning (Section 3.3): The most granular optimization — why splitting Q instead of K/V eliminates inter-warp shared memory communication, the concrete difference between "split-K" (FlashAttention) and "split-Q" (FlashAttention-2), and how this maps to both forward and backward passes.

3.4 Detailed, Sentence-Based Technical Breakdown

This is primarily a GPU kernel optimization paper whose core idea is that FlashAttention's tiling strategy was algorithmically sound but its mapping to GPU parallelism was suboptimal at three levels — the arithmetic intensity (too many slow element-wise operations), the coarse-grained scheduling (not enough thread blocks to occupy all SMs), and the fine-grained data flow (unnecessary shared memory traffic between warps) — and that fixing all three can approximately double throughput without changing the mathematical result.

---

The Non-Matmul FLOP Problem: Why Element-Wise Operations Are the Hidden Bottleneck

Before diving into the algorithm changes, one must understand the hardware asymmetry that motivates them. The A100 GPU has two classes of floating-point execution units: Tensor Cores, which are specialized circuits that perform matrix multiply-accumulate operations on 16-bit floating-point inputs at extremely high throughput (312 TFLOPs/s for FP16/BF16), and CUDA cores, which handle all other floating-point operations (addition, multiplication, division, exponential, maximum, comparison) at dramatically lower throughput (19.5 TFLOPs/s for FP32). The ratio is 16:1 — a single element-wise exponential or division instruction costs as much wall-clock time as 16 fused multiply-add operations on Tensor Cores.

Attention computation is a hybrid workload: it alternates between large matrix multiplies (QK^T, PV, dS K, etc.) that run on Tensor Cores, and element-wise operations (row-wise max, subtract max, exponential, row-wise sum of exponentials, division for rescaling, log for logsumexp) that run on CUDA cores. Even though the element-wise operations represent a small fraction of the total FLOP count, they represent a much larger fraction of the total execution time because each such FLOP takes 16× longer. To achieve overall throughput approaching GEMM-like efficiency (70%+ of theoretical peak), the kernel must spend as high a fraction of its time as possible executing Tensor Core instructions — which means minimizing the number and frequency of element-wise operations, even if the total FLOP count increases slightly.

This is the lens through which all three algorithm tweaks in Section 3.1 should be understood: they are not about reducing total arithmetic work, but about shifting work from the slow execution units to the fast ones.

---

The Online Softmax Algorithm and Its Two FlashAttention-2 Tweaks

The core numerical mechanism that enables tiled attention computation is online softmax (Milakov and Gimelshein, 2018; Rabe and Staats, 2021). Standard softmax requires seeing all elements of a row before computing any output, because every element must be normalized by the sum of exponentials of all elements in that row. The softmax of a vector s is:

$p_i = \frac{e^{s_i}}{\sum_j e^{s_j}}$

This appears to be an inherently sequential operation — you need the denominator (which depends on all elements) before you can compute any p_i. If the row is split across multiple blocks (as in tiled attention where QK^T is computed one column block at a time), a naive implementation would need to compute all S values first, then compute softmax, then multiply by V — exactly the materialization that FlashAttention avoids.

Online softmax solves this by maintaining running statistics that can be updated incrementally as each new block of the row is processed, and by deferring the final normalization. The key insight is that softmax normalization can be "undone" and "redone" with better statistics as more information arrives.

FlashAttention's Online Softmax (The Starting Point)

Consider processing a single row block of attention scores, split across two column blocks. We have S^(1) and S^(2), and corresponding value blocks V^(1) and V^(2). The goal is to compute O = softmax([S^(1) S^(2)]) · [V^(1); V^(2)].

A standard (non-online) approach would compute:

1. m = max(rowmax(S^(1)), rowmax(S^(2))) — the global row-wise max, used for numerical stability
2. ℓ = rowsum(e^{S^(1) - m}) + rowsum(e^{S^(2) - m}) — the global row-wise sum of exponentials
3. O = diag(ℓ)^{-1} · (e^{S^(1) - m} V^(1) + e^{S^(2) - m} V^(2))

This requires access to both S^(1) and S^(2) simultaneously to compute m and ℓ.

FlashAttention's online softmax processes block 1 first, then block 2, updating the output in place. After processing block 1:

1. m^(1) = rowmax(S^(1)) — local max from block 1
2. ℓ^(1) = rowsum(e^{S^(1) - m^(1)}) — local sum from block 1
3. O^(1) = diag(ℓ^(1))^{-1} e^{S^(1) - m^(1)} V^(1) — the correctly normalized output if block 1 were the entire row

Then when block 2 arrives:

4. m^(2) = max(m^(1), rowmax(S^(2))) — updated global max
5. ℓ^(2) = e^{m^(1) - m^(2)} ℓ^(1) + rowsum(e^{S^(2) - m^(2)}) — updated global sum, where the first term rescales the old sum to use the new max
6. O^(2) = diag(ℓ^(1)/ℓ^(2))^{-1} O^(1) + diag(ℓ^(2))^{-1} e^{S^(2) - m^(2)} V^(2) — the old output is rescaled by the ratio of old to new normalization, and the new block's contribution is added

The critical detail in step 6: the old output O^(1) was computed with normalization diag(ℓ^(1))^{-1}, but the new normalization should be diag(ℓ^(2))^{-1}. To correct this, O^(1) is multiplied by diag(ℓ^(1)/ℓ^(2))^{-1} = diag(ℓ^(2)/ℓ^(1)), which "undoes" the old normalization by ℓ^(1) and "reapplies" the new normalization by ℓ^(2). This is the rescaling operation that happens on every inner loop iteration (every time a new column block is processed).

After all Tc column blocks are processed, O^(Tc) is the correctly normalized output O.

Tweak 1: Defer the Final Rescaling (Unscaled Accumulator)

FlashAttention-2's first tweak eliminates the rescaling of the old output on every iteration. Instead of maintaining a properly normalized output O^(j) throughout the loop, it maintains an unscaled accumulator that is only normalized once at the very end.

The insight: instead of storing O^(1) = diag(ℓ^(1))^{-1} e^{S^(1) - m^(1)} V^(1) (the normalized version), store the unnormalized version:

$\tilde{O}^{(1)} = e^{S^{(1)} - m^{(1)}} V^{(1)}$

This is simply O^(1) multiplied by ℓ^(1) — the numerator without the denominator. When block 2 arrives, the update becomes:

$\tilde{O}^{(2)} = \text{diag}(e^{m^{(1)} - m^{(2)}})^{-1} \tilde{O}^{(1)} + e^{S^{(2)} - m^{(2)}} V^{(2)}$

The first term rescales the old unnormalized accumulator to account for the potentially new max m^(2): if m^(2) > m^(1), all exponentials from block 1 need to be scaled down by e^{m^(1) - m^(2)} to use the new reference point. This is a single element-wise multiplication — much cheaper than the two operations (multiply by ℓ^(2)/ℓ^(1) and divide) that the normalized version requires.

After all column blocks are processed, the final output is obtained by a single normalization:

$O^{(T_c)} = \text{diag}(\ell^{(T_c)})^{-1} \tilde{O}^{(T_c)}$

In concrete terms, FlashAttention required, for each inner loop iteration: (a) compute the new rescaling factor ℓ^(2)/ℓ^(1), (b) multiply O^(1) by this factor element-wise, (c) normalize the new block's contribution by ℓ^(2). FlashAttention-2 replaces this with: (a) exponentiate the max difference e^{m^(1) - m^(2)}, (b) multiply the old unscaled accumulator by this factor element-wise, (c) add the new block's unnormalized contribution. The final division by ℓ^(Tc) happens once at the end, not once per iteration.

Why this matters: The rescaling in FlashAttention involved a division (or multiplication by reciprocal) with ℓ, which changes on every iteration. In FlashAttention-2, the rescaling involves only the max difference, which stays the same for many iterations (the row-wise max rarely changes after the first few blocks). When m^(2) = m^(1) for a block, the rescaling factor e^{m^(1) - m^(2)} = 1, and the update simplifies to simply adding the new block's contribution — zero non-matmul FLOPs beyond the addition. This reduces the number of non-matmul operations in the inner loop, where they are most expensive because they execute on every (row block, column block) pair.

Tweak 2: Store Logsumexp Instead of Max and Sum Separately

The backward pass of attention needs to reconstruct the softmax probabilities P from the forward pass's intermediate values. In the original FlashAttention, both the row-wise max m and the row-wise sum of exponentials ℓ needed to be stored to HBM (and later reloaded during the backward pass) because softmax recomputation requires:

$P_{ij} = \frac{e^{S_{ij} - m_i}}{\ell_i}$

FlashAttention-2 observes that only the logsumexp is needed:

$L_i = m_i + \log(\ell_i)$

Then during the backward pass, P can be reconstructed as:

$P_{ij} = \exp(S_{ij} - L_i)$

because:

$\exp(S_{ij} - L_i) = \exp(S_{ij} - m_i - \log\ell_i) = \frac{\exp(S_{ij} - m_i)}{\ell_i} = P_{ij}$

The logsumexp has the same memory footprint as storing either m or ℓ individually (it's a single scalar per row, size N), but replaces two stored values with one. This halves the memory traffic for the softmax statistics in the backward pass: instead of reading both m ∈ R^N and ℓ ∈ R^N from HBM, the backward pass reads only L ∈ R^N.

Why this matters: Memory bandwidth is the primary bottleneck in attention, even after FlashAttention's tiling. Every byte that doesn't need to be transferred between HBM and SRAM translates directly to speedup. The logsumexp substitution eliminates one full vector read per row block during the backward pass. This is a pure systems optimization — mathematically equivalent, requiring one extra logarithm operation in the forward pass (to compute L from m and ℓ) and one extra exponential in the backward pass (to reconstruct P from S and L), but the savings in HBM bandwidth far outweigh the cost of these element-wise operations because HBM access (~1.5 TB/s) is roughly 10–13× slower than SRAM access (~19 TB/s) and much slower than register-level computation.

---

Algorithm 1 Walkthrough: The Full FlashAttention-2 Forward Pass

Algorithm 1 in the paper describes the complete forward pass. Let's walk through it operationally, highlighting how the tweaks manifest in the actual computation.

Input partitioning (lines 1–2): The sequence length N is divided into Tr row blocks of size Br each, where Tr = ⌈N/Br⌉, and Tc column blocks of size Bc each, where Tc = ⌈N/Bc⌉. The block sizes Br and Bc are typically chosen from {64, 128} depending on head dimension and available shared memory. Q, K, V are each split into these blocks. The output O and the logsumexp L are similarly divided into Tr blocks of size Br×d and Br respectively. The choice of Tr and Tc determines the total number of (row block, column block) pairs = Tr × Tc — each pair requires one inner loop iteration.

Outer loop (line 3): For each row block i from 1 to Tr, the computation processes one complete row block of the output. This is where FlashAttention-2's parallelization over sequence length is exploited: different thread blocks handle different row blocks independently. The Q_i block is loaded once into SRAM for this row block (line 4) and reused across all Tc inner loop iterations — a key tiling optimization.

Initialization (line 5): Three running statistics are initialized per row block:

• O^(0)_i = 0 (Br × d): the unscaled output accumulator, initially zero
• ℓ^(0)_i = 0 (Br): the running sum of exponentials, initially zero
• m^(0)_i = -∞ (Br): the running row-wise max, initialized to negative infinity so any real value becomes the max on the first iteration

Inner loop (lines 6–10): For each column block j from 1 to Tc:

Line 7: Load K_j (size Bc × d) and V_j (size Bc × d) from HBM into SRAM. These stay in SRAM for the inner loop iteration and are discarded after. Q_i (size Br × d) was already loaded and stays resident.

Line 8: Compute S^(j)_i = Q_i K_j^T. This is a (Br × d) × (d × Bc) matrix multiply producing a Br × Bc matrix of attention scores. This runs on Tensor Cores — it's the operation that should dominate execution time. For typical block sizes (Br = Bc = 64 or 128, d = 64 or 128), this is a substantial matrix multiply that efficiently utilizes Tensor Core throughput.

Line 9: Compute three element-wise quantities (the "softmax statistics update"):

First, the new running max:

$m_i^{(j)} = \max(m_i^{(j-1)}, \text{rowmax}(S_i^{(j)}))$

where rowmax returns a vector of length Br, each element being the maximum of the corresponding row of S^(j)_i. The max is taken element-wise between the old running max and the new block's row-wise max. This is two operations: a row-wise reduction (non-matmul) and an element-wise maximum.

Second, the block's exponentiated and max-subtracted scores:

$\tilde{P}_i^{(j)} = \exp(S_i^{(j)} - m_i^{(j)})$

where the subtraction of m_i^(j) is broadcast across columns (each row element of S^(j)_i has the corresponding element of m_i^(j) subtracted). The exponential is applied pointwise. This produces a Br × Bc matrix. The tilde notation emphasizes this is unnormalized — it is not yet divided by ℓ.

Third, the updated running sum:

$\ell_i^{(j)} = e^{m_i^{(j-1)} - m_i^{(j)}} \ell_i^{(j-1)} + \text{rowsum}(\tilde{P}_i^{(j)})$

The first term rescales the old sum: if the max increased (m_i^(j) > m_i^(j-1)), the old exponentials were computed relative to a smaller reference point and are too large by a factor of e^{m_i^(j-1) - m_i^(j)} (which is < 1), so the old sum is multiplied by this factor to bring it to the new scale. If the max didn't change, the factor is 1 and this term reduces to simply keeping the old sum. The second term computes the row-wise sum of the new block's exponentiated scores. This produces a vector of length Br.

Line 10: Update the unscaled output accumulator:

$O_i^{(j)} = \text{diag}(e^{m_i^{(j-1)} - m_i^{(j)}})^{-1} O_i^{(j-1)} + \tilde{P}_i^{(j)} V_j$

The first term rescales the old unscaled accumulator using the same max-change factor as in the ℓ update. If the max increased, the old accumulator's exponentials were computed with the old max, so they need to be scaled down. The multiplication by diag(...)^{-1} is equivalent to element-wise multiplication of each row of O_i^(j-1) by e^{m_i^{(j-1)} - m_i^{(j)}}. The second term computes the new block's contribution: a (Br × Bc) × (Bc × d) matrix multiply producing Br × d. This runs on Tensor Cores.

Finalization (lines 12–13): After all Tc column blocks are processed, the unscaled accumulator O_i^(Tc) contains:

$\tilde{O}_i^{(T_c)} = \sum_{j=1}^{T_c} e^{S_i^{(j)} - m_i^{(T_c)}} V_j$

where m_i^(Tc) is the global row-wise max for this row block. The final output is obtained by normalizing once:

$O_i = \text{diag}(\ell_i^{(T_c)})^{-1} O_i^{(T_c)}$

This is a single element-wise division of each row of the accumulator by the corresponding element of the final sum of exponentials ℓ_i^(Tc). The logsumexp is then computed for storage:

$L_i = m_i^{(T_c)} + \log(\ell_i^{(T_c)})$

Lines 14–15: Write O_i (size Br × d) and L_i (size Br) to HBM. These are the only outputs of the forward pass beyond the input tensors.

Mathematical correctness. The algorithm computes exact softmax attention — it is mathematically equivalent to the naive O(N^2) implementation, not an approximation. The proof follows from the same inductive argument as FlashAttention (Dao et al., 2022, Theorem 1): at each step j, the running statistics correctly represent what the global softmax statistics would be if only the first j column blocks existed, and the final rescaling at step Tc produces the globally correct normalization.

Causal masking optimization. For autoregressive (decoder-only) models, entries where the column index exceeds the row index must be set to -∞ (masked out). FlashAttention-2 exploits the block structure: (1) for any block where all column indices exceed all row indices (roughly half the blocks for large N), the entire block computation can be skipped, yielding approximately 1.7–1.8× speedup; (2) for blocks on the diagonal (where some but not all entries need masking), the causal mask is applied only to those specific blocks — for each row, only one block needs masking if the blocks are square.

---

The Backward Pass: Simpler Statistics, Same Structure

The backward pass (Algorithm 2) computes gradients dQ, dK, dV given the output gradient dO, the forward pass outputs O, and the stored logsumexp L. By the chain rule of matrix calculus (Section 2.2):

$dV = P^T dO$ $dP = dO V^T$ $dS = d\text{softmax}(dP)$ $dQ = dS K$ $dK = Q dS^T$

where dsoftmax is the gradient of row-wise softmax: if p = softmax(s) and we have gradient dp, then ds = (diag(p) - pp^T) dp = p ∘ (dp - rowsum(p ∘ dp)), where ∘ denotes element-wise multiplication. The term D = rowsum(dO ∘ O) ≡ rowsum(p ∘ dp) precomputes the inner product between each row's probability vector and its gradient, which is then broadcast in the softmax backward computation.

Key storage simplification. FlashAttention's backward pass needed both the row-wise max m and the row-wise sum of exponentials ℓ from the forward pass to reconstruct P = softmax(S). FlashAttention-2 stores only L = m + log(ℓ) and reconstructs:

$P^{(j)}_i = \exp(S^{(j)}_i - L_i)$

where S^(j)_i is recomputed on-the-fly in the backward pass (since Q_i and K_j are already loaded to SRAM for the gradient computations). The exponential of (S - L) equals e^{S - m - logℓ} = e^{S - m} / ℓ = P, exactly the softmax probabilities. This eliminates one HBM read per (row, column) block pair compared to storing both m and ℓ.

The D precomputation (line 4). Before the main loops, compute:

$D = \text{rowsum}(dO \circ O)$

where ∘ is element-wise multiplication (Hadamard product). For each row, this sums dO_{ik} · O_{ik} over the head dimension d, producing a vector of length N. This D vector is the key quantity needed for the softmax backward: ds = p ∘ (dp - D), meaning each row's attention gradient is the probability matrix P element-wise multiplied by (the incoming gradient dP minus the row's D value broadcast across columns).

Loop structure inversion. The backward pass iterates with the outer loop over column blocks j (line 5) and inner loop over row blocks i (line 8). This is the reverse of the forward pass's nesting. The reason: in the backward pass, dV_j gets contributions from all row blocks i (since dV_j = sum_i P_i^{(j)T} dO_i), making it natural to accumulate dV_j within the column block's outer loop iteration. Similarly, dK_j = sum_i dS_i^{(j)T} Q_i. Meanwhile, dQ_i gets contributions from all column blocks j (dQ_i = sum_j dS_i^{(j)} K_j), requiring accumulation across the column loop.

Inner loop operations (lines 10–16): For each (i, j) pair:

Line 10: Recompute S^(j)_i = Q_i K_j^T (same as forward pass). Tensor Core matmul.

Line 11: Reconstruct P^(j)_i = exp(S^(j)_i - L_i) where L_i is the stored logsumexp for row block i. The subtraction of L_i (broadcast) and exponential are element-wise operations.

Line 12: Accumulate dV_j += P_i^{(j)T} dO_i. This is a (Bc × Br) × (Br × d) matrix multiply. Tensor Core matmul. The += is key: dV_j accumulates contributions from all row blocks i in the inner loop, so it stays in SRAM across the inner loop iterations and is only written to HBM after the inner loop completes (line 18).

Line 13: Compute dP^(j)_i = dO_i V_j^T. This is a (Br × d) × (d × Bc) matrix multiply. Tensor Core matmul.

Line 14: Compute the softmax gradient:

$dS^{(j)}_i = P^{(j)}_i \circ (dP^{(j)}_i - D_i)$

where D_i is broadcast across columns. This is the softmax backward formula: element-wise multiply P by (dP minus the row-wise scalar D). Three element-wise operations (broadcast subtract, Hadamard product). This runs on CUDA cores (non-matmul).

Line 15: Update dQ_i. This is the one place requiring communication between different column-block thread blocks. dQ_i += dS^(j)_i K_j is a (Br × Bc) × (Bc × d) matrix multiply. Tensor Core matmul. But dQ_i accumulates across all column blocks j, and different column blocks are processed by different thread blocks (Section 3.2). The update uses atomic add to safely combine contributions from concurrent thread blocks writing to the same dQ_i location in HBM.

Line 16: Accumulate dK_j += dS_i^{(j)T} Q_i. This is a (Bc × Br) × (Br × d) matrix multiply. Tensor Core matmul. Like dV_j, dK_j stays in SRAM and accumulates across the inner loop.

Writeback (line 18): After all row blocks are processed for column block j, dK_j and dV_j are written from SRAM to HBM.

Multi-query and grouped-query attention. For MQA (Shazeer, 2019) and GQA (Ainslie et al., 2023), multiple query heads share the same key/value head. The forward pass handles this by implicitly manipulating head indices — the same K and V are reused across multiple Q heads. The backward pass must sum gradients dK and dV across the heads that were implicitly duplicated, which is handled by the atomic add mechanism: each query head's computation contributes to the same dK_j and dV_j accumulators.

---

Thread Block Scheduling: Parallelizing Over Sequence Length

The original FlashAttention parallelizes only over the batch dimension and the number of heads dimension. For a model with H heads and batch size B, this yields B × H thread blocks. On an A100 with 108 SMs, when B × H ≥ ~80–108, all SMs are occupied and the GPU is well-utilized. But two scenarios cause low occupancy:

1. Long sequences → small batch sizes: Memory constraints mean that when training on sequence length 8k or 16k, the batch size must be small (often 1 or 2) to fit activations and optimizer states in GPU memory. With 16 heads and batch size 1, only 16 thread blocks execute — leaving 92 out of 108 SMs idle.

2. Small number of heads in the model: Some architectures use fewer attention heads.

Forward pass parallelization (Section 3.2, Figure 2 left). FlashAttention-2 additionally parallelizes over the sequence length dimension in the forward pass. The outer loop over row blocks (line 3 of Algorithm 1) processes Tr = ⌈N/Br⌉ independent row blocks per head. Instead of having one thread block process all Tc inner-loop iterations for a given row block sequentially, different row blocks can be assigned to different thread blocks. This multiplies the available parallelism by Tr: the total number of thread blocks becomes B × H × Tr.

For sequence length 8k with Br = 128, Tr = 8000/128 = 62.5 → 63. Even with batch size 1 and 16 heads, the total thread blocks = 1 × 16 × 63 = 1,008, more than enough to occupy all 108 SMs (multiple waves of thread blocks will be scheduled). The forward pass outer loop is embarrassingly parallel — row blocks do not depend on each other — so no synchronization is needed.

Why the loop order matters. This parallelization is possible because FlashAttention-2's forward pass uses an outer loop over row blocks (the Q dimension). The original FlashAttention paper used the reverse order (outer loop over column blocks). With an outer loop over column blocks, different thread blocks would process different K,V blocks, but they'd all need to write to the same output rows, requiring synchronization. The Triton implementation first suggested swapping the loop order and parallelizing over rows, and FlashAttention-2 adopts this innovation.

Backward pass parallelization (Section 3.2, Figure 2 right). The backward pass parallelizes over the sequence length dimension as well, but differently. The outer loop iterates over column blocks j (line 5 of Algorithm 2), and each thread block handles one column block. Within each column block's thread block, the inner loop over row blocks i is executed sequentially. The total parallelism is B × H × Tc.

The challenge is that dQ needs to be updated by multiple column blocks (line 15: dQ_i += dS_i^{(j)} K_j for each column block j). Since different column blocks are processed by different thread blocks, they may attempt to write to the same dQ_i locations in HBM simultaneously. The solution is atomic add operations: the GPU hardware guarantees that concurrent atomic adds to the same memory location are serialized correctly without data races. The cost is that atomic adds to HBM are slower than non-atomic writes (they require locking the memory bus for that location), but since dQ updates are a relatively small fraction of the total backward pass FLOPs, this overhead is acceptable.

Why column blocks for the backward pass outer loop? Each column block j needs to accumulate dV_j and dK_j across all row blocks i. By making the column block the outer loop, dV_j and dK_j stay resident in SRAM for the entire inner loop, accumulating contributions from all row blocks without needing to write intermediate results to HBM. If the outer loop were over row blocks instead, each row block's thread block would need to partially update dV_j and dK_j for each column block, requiring either atomic adds on every column block or storing partial sums in HBM between inner loop iterations — both of which would be less efficient.

---

Warp-Level Work Partitioning: Eliminating Split-K

Within each thread block, the computation must be further divided among warps — groups of 32 threads that execute in lockstep (SIMT: Single Instruction, Multiple Threads). A typical thread block in FlashAttention uses 4 or 8 warps (128 or 256 threads). The key design decision is which dimension of the matrices to partition across warps and which to make available to all warps. This determines the communication pattern between warps.

FlashAttention's "Split-K" Scheme (Figure 3a)

In the original FlashAttention forward pass, K and V are partitioned across warps while Q is accessible by all warps. Specifically:

1. K is split along the column (sequence length) dimension into 4 chunks (for 4 warps), each of size Bc/4 × d.
2. V is similarly split into 4 chunks of size Bc/4 × d.
3. Q (size Br × d) is loaded once and broadcast to all warps (each warp has its own copy in registers).

Each warp independently computes its portion of the attention scores: warp w computes S_w = Q K_w^T, which is a Br × (Bc/4) matrix. Then it applies online softmax to this partial score matrix, computes partial output O_w = softmax(S_w) V_w, and writes this partial output to shared memory. All warps synchronize (barrier), and then the partial outputs are summed element-wise: O = O_1 + O_2 + O_3 + O_4. This summation requires reading all four partial outputs from shared memory.

Why this is inefficient: Every inner loop iteration (for each column block j) involves:

• 4 writes of partial outputs to shared memory (each warp writes its Br × d chunk)
• A barrier synchronization (all warps must finish before any can proceed)
• 4 reads of partial outputs from shared memory (each warp needs all four to sum them)
• A reduction (element-wise addition of 4 matrices)

Shared memory is fast (~19 TB/s) but not free. On A100, shared memory latency is approximately 20–30 clock cycles. The barrier synchronization is even more expensive — warps that finish early must stall waiting for the slowest warp. Since the inner loop executes Tc times per row block, these overheads accumulate.

FlashAttention-2's "Split-Q" Scheme (Figure 3b)

FlashAttention-2 reverses the partition: Q is split across warps while K and V are broadcast to all warps. Specifically:

1. Q is split along the row (sequence length) dimension into 4 chunks, each of size Br/4 × d.
2. K (size Bc × d) and V (size Bc × d) are loaded once and broadcast to all warps.

Each warp w independently computes its portion of the attention scores: S_w = Q_w K^T, which is a (Br/4) × Bc matrix. It applies online softmax to this partial score matrix, and computes its portion of the output O_w = softmax(S_w) V, which is a (Br/4) × d matrix. This is already the final output for those rows — no summation with other warps' outputs is needed. The four warps' outputs are simply concatenated (in registers, they map to different rows of the output matrix) without any communication.

Why this eliminates shared memory traffic: Each warp's output O_w corresponds to a disjoint subset of the output rows (rows w·Br/4 to (w+1)·Br/4 - 1). There is no overlap, so no reduction is needed. The warps never need to write their outputs to shared memory for other warps to read. The only information that must be shared between warps is the softmax statistics (row-wise max and sum of exponentials) if the row-wise max changes, but in practice the max stabilizes quickly and this communication is minimal compared to the Br × d output matrix that FlashAttention's split-K scheme needed to exchange.

Concrete example: Suppose Br = 128, Bc = 128, d = 64, and 4 warps. In FlashAttention's split-K:

• Each warp computes a 128 × 32 attention score matrix and a 128 × 64 partial output
• Each warp writes 128 × 64 = 8,192 elements to shared memory (32 KB for FP32)
• Total shared memory traffic: 4 × 8,192 × 2 (write + read) = 65,536 element transfers per inner loop iteration

In FlashAttention-2's split-Q:

• Each warp computes a 32 × 128 attention score matrix and a 32 × 64 output
• No shared memory writes for output (each warp owns its 32 × 64 output slice)
• Total shared memory traffic for output: 0 transfers per inner loop iteration

The savings are substantial: eliminating 65K element transfers per inner loop iteration, multiplied by Tc inner loop iterations (which can be dozens to hundreds for long sequences), multiplied by the number of thread blocks.

Backward Pass Warp Partitioning

The backward pass similarly avoids the split-K scheme, but the situation is more nuanced because of the more complex dependencies:

• dS^(j)_i depends on P^(j)_i, which requires the row-wise logsumexp L_i (shared across all warps in the row block)
• dQ_i accumulates across all column blocks j via atomic adds (inter-thread-block communication)
• dK_j accumulates across all row blocks i within the same thread block
• dV_j accumulates across all row blocks i within the same thread block

FlashAttention-2 partitions the backward pass warps to minimize shared memory traffic while respecting these dependencies. The details are implementation-specific, but the principle is the same: partition along dimensions that create disjoint output regions where possible, and use registers rather than shared memory for inter-warp communication when reduction is unavoidable.

Tuning block sizes. The block sizes Br and Bc are critical hyperparameters:

• Larger blocks reduce the number of HBM loads/stores (fewer outer loop iterations, better amortization of loading Q_i or K_j) but increase register pressure and shared memory usage
• Smaller blocks require fewer registers and less shared memory, allowing higher occupancy (more thread blocks per SM), but increase the number of HBM transfers

If the block size is too large, the kernel either fails to launch (not enough shared memory) or suffers from register spilling (registers overflow to slow L1 cache/HBM). Typical values are Br, Bc ∈ {64, 128}, chosen based on head dimension d and available shared memory (192 KB per SM on A100). The authors manually tune for each head dimension — since there are only 4 choices per dimension pair, this is feasible — but note that auto-tuning could eliminate this manual labor in future work.

---

The Complete System: Putting It All Together

To execute a forward pass of attention on input Q, K, V ∈ R^{N×d}:

1. The host (CPU) launches the FlashAttention-2 kernel with grid dimensions specifying B × H × Tr thread blocks (B = batch size, H = number of heads, Tr = number of row blocks). Each thread block is scheduled to an SM by the GPU's hardware scheduler.

2. Each thread block is assigned a specific (batch index, head index, row block index) tuple. It loads its assigned Q_i block (size Br × d) from HBM to SRAM. It initializes its running statistics (unscaled output accumulator, running sum, running max) in registers.

3. Within each thread block, the 4 or 8 warps partition Q_i by rows (split-Q). Each warp loads its portion of Q_i into registers. K and V are loaded column-block by column-block from HBM to SRAM and broadcast to all warps.

4. For each column block j (inner loop, executed sequentially within the thread block):

   • All warps cooperate to load K_j and V_j from HBM to SRAM
   • Each warp independently computes its portion of S = Q_w K_j^T on Tensor Cores
   • Each warp independently updates its running max m, computes exponentiated scores P̃, updates running sum ℓ, and updates its portion of the unscaled output accumulator Õ = rescale(Õ_old) + P̃ V_j on Tensor Cores
   • Warp-level softmax statistics (m, ℓ) are synchronized through shared memory when the max changes

5. After all Tc column blocks are processed, each warp has its portion of the final unscaled accumulator Õ. The thread block collectively computes the final normalization O = Õ / ℓ (element-wise division), computes L = m + log(ℓ), and writes O and L to HBM.

The backward pass follows the same principles but with the loop order inverted (column blocks outer, row blocks inner) and with atomic adds for dQ updates across thread blocks.

The three levels of optimization — algorithm tweaks (reduce element-wise ops), thread block scheduling (increase occupancy via sequence-length parallelization), and warp partitioning (eliminate shared memory traffic) — are orthogonal and multiplicative. The algorithm tweaks reduce the cost per inner loop iteration. The thread block scheduling ensures more inner loop iterations execute in parallel across SMs. The warp partitioning reduces the cost of communication within each inner loop iteration. Together, they approximately double the throughput compared to FlashAttention.

4. Key Insights and Innovations

Innovation 1: Diagnosing and Naming the Specific Parallelism Defects That Cap Exact Attention Throughput

The field's understanding of attention bottlenecks has evolved in layers: first it was the O(N²) memory footprint (solved by FlashAttention's I/O-aware tiling), then it became clear that even with the memory problem solved, attention kernels run at only 25–40% of GPU theoretical peak. But why? The dominant assumption, implicit in how most systems papers evaluate attention kernels, was that the remaining gap was simply "engineering" — a matter of incremental tuning, better autotuning of block sizes, or waiting for faster hardware.

FlashAttention-2's primary intellectual contribution is a structured diagnosis that names three distinct, addressable defects at three different levels of the GPU parallelism hierarchy, and shows they are not minor tuning issues but systematic consequences of design choices in how the algorithm is mapped to hardware. The paper doesn't just say "we made it faster" — it says "here are the three specific reasons FlashAttention was leaving 60–75% of the GPU idle, here is why each matters, and here is how we fix them."

The three defects constitute a taxonomy that is transferable to other GPU kernels:

1. Arithmetic intensity imbalance (Section 3.1): The 16× throughput gap between Tensor Cores (312 TFLOPs/s) and CUDA cores (19.5 TFLOPs/s) means that element-wise operations — even when they are a small fraction of FLOP count — dominate execution time. This is not a memory bandwidth problem (which FlashAttention already addressed), but a compute unit utilization problem. The standard online softmax algorithm, while numerically elegant, executes several element-wise operations per inner loop iteration (rescaling the accumulated output by ℓ^(j)/ℓ^(j-1), computing exponentials, row-wise reductions) that run on the slow CUDA cores.

2. Coarse-grained parallelism starvation (Section 3.2): Parallelizing only over batch and head dimensions yields B × H thread blocks. For long-sequence training where memory constraints force small batch sizes, this number can be as low as 16 — far below the 108 SMs on an A100. The outer loop order in the original FlashAttention (columns outer) made sequence-length parallelization unnatural because column-block thread blocks would need to coordinate on output writes. This is an occupancy problem, and it is most severe exactly in the regime that practitioners care about most (training on long contexts). This is not a problem that better block size tuning can fix — it requires adding a new dimension of parallelism.

3. Fine-grained inter-warp communication (Section 3.3): The "split-K" scheme — partitioning K and V across warps while broadcasting Q — forces every warp to write its partial output to shared memory, synchronize, and read all other warps' partial outputs, on every inner loop iteration. This is a data flow problem: the work partitioning creates artificial communication that doesn't exist in the mathematics. Each warp computes a partial result that overlaps with other warps' partial results along the output dimensions, requiring a reduction.

The significance of this diagnostic framework extends beyond the specific fixes. It gives kernel engineers a checklist for analyzing why any tiled attention implementation might be underperforming: are you spending too much time on slow compute units? Do you have enough thread blocks to occupy all SMs? Are your warps communicating unnecessarily? The paper's contribution is as much about how to think about attention kernel optimization as it is about the specific optimizations themselves.

The evidence for each defect is not presented as an ablation (there is no "FlashAttention-2 without the algorithm tweaks" experiment), but the benchmarking results (Figures 4–6) show that the combined fixes push throughput from the 25–40% range (FlashAttention) to the 50–73% range (FlashAttention-2), demonstrating that these three defects collectively account for the majority of the remaining efficiency gap. The gap to GEMM's 80–90% is now much smaller, suggesting the diagnosis captured the dominant bottlenecks.

Innovation 2: The Unscaled Accumulator as an Arithmetic Intensity Optimization Disguised as an Algorithmic Simplification

The paper presents the forward pass tweak — maintaining an unscaled output accumulator Õ instead of a normalized O, and deferring the final division by ℓ to the end — as a reduction in non-matmul FLOPs (Section 3.1.1). But this undersells what is conceptually distinctive about the move. The unscaled accumulator is best understood as changing which arithmetic operations are in the critical path of the inner loop — not just reducing their count, but changing their nature.

In FlashAttention's online softmax, every inner loop iteration performs a rescaling of the old output by ℓ^(j)/ℓ^(j-1). The denominator ℓ^(j-1) is the running sum from the previous iteration, and ℓ^(j) is the updated sum. This means the rescaling factor changes on potentially every iteration, because the sum of exponentials grows monotonically as more column blocks are processed (each block adds positive contributions to ℓ). A division or reciprocal multiplication is required per iteration, and the factor is not predictable — it's data-dependent.

In FlashAttention-2's version, the rescaling factor for the old accumulator is e^{m^(j-1) - m^(j)}. The key observation is that the row-wise maximum m stabilizes quickly — after processing a few column blocks, it is rare for a new block to contain a value larger than the current running max. For most iterations in the inner loop, m^(j) = m^(j-1), making the rescaling factor exactly 1. When the factor is 1, the rescaling operation is a no-op that can be compiled away or skipped with a branch. This means that for the majority of inner loop iterations, the update reduces to:

$\tilde{O}^{(j)} = \tilde{O}^{(j-1)} + \tilde{P}^{(j)} V_j$

which is simply a matrix multiply-accumulate — exactly the operation that Tensor Cores are designed to execute at peak throughput. The element-wise rescaling operation has been moved off the critical path for most iterations, and when it does occur (when a new max is found), its cost is amortized over many subsequent iterations.

This is fundamentally different from a simple "reduce FLOP count" optimization. It is an arithmetic intensity restructuring — it changes the mix of operations in the inner loop from a forced alternation between Tensor Core matmuls and CUDA core element-wise ops to a pattern where most iterations are pure Tensor Core matmuls with occasional CUDA core interventions when the max changes. This better matches the hardware's asymmetric throughput capabilities.

The intellectual move is recognizing that numerical stability conventions (always keeping the output normalized) impose a computational cost that is not intrinsic to the mathematics. By relaxing the convention — storing the output in an unnormalized representation and normalizing only at the end — the algorithm's runtime characteristics change substantially even though the FLOP count changes minimally. This is a design pattern that could apply to other tiled reduction algorithms where running statistics stabilize after an initial transient.

Innovation 3: Sequence-Length Parallelism as a Systematic Solution to the Long-Sequence Occupancy Problem

The idea of parallelizing over the sequence length dimension was first implemented in Triton's FlashAttention (credited to Phil Tillet in Section 3.2). FlashAttention-2's contribution is not the idea itself but rather (a) the systematic articulation of why it matters and when, and (b) the demonstration that it must be paired with a specific loop ordering to work without synchronization overhead, with different orderings required for forward and backward passes.

The non-obvious insight is that the choice of which loop is outer (rows or columns) determines which dimension can be trivially parallelized. In the forward pass, with rows as the outer loop, each row block's output is independent — different thread blocks can process different row blocks and write to disjoint output regions without any communication. If columns were outer (as in the original FlashAttention paper), parallelizing over columns would mean different thread blocks compute partial contributions to the same output rows, requiring either synchronization or atomic operations to combine.

In the backward pass, the situation reverses. With columns as the outer loop, each column block's dK_j and dV_j accumulations are independent across thread blocks (they write to disjoint regions), but dQ_i gets contributions from all column blocks. The paper's solution — using the sequence-length parallelization with column blocks outer in the backward pass, and handling the dQ conflict with atomic adds — is a specific instance of a more general principle: when parallelizing a tiled reduction, pick the outer loop dimension that maximizes the number of independent output regions, and use atomic operations for the residual shared outputs.

The significance goes beyond the 2× speedup numbers. This insight changes how one thinks about GPU occupancy for attention: occupancy is not just a function of batch size and number of heads (which are often constrained by model architecture and memory), but can be manufactured by slicing the sequence length dimension. For very long sequences (16k, 32k, and beyond), Tr and Tc can be large (128 or 256), providing abundant parallelism even with batch size 1. In the limit of extremely long contexts, sequence-length parallelism dominates and makes the occupancy problem disappear entirely — no matter how small the batch size, there are enough row blocks to fill all SMs.

The evidence in Table 1 supports the practical impact: for GPT3-1.3B at 8k context, FlashAttention-2 reaches 220 TFLOPs/s compared to 170 TFLOPs/s for FlashAttention — a 1.29× improvement that is almost entirely attributable to better occupancy at this small-batch, long-sequence setting. The gap is larger for 8k than for 2k context (where FlashAttention already achieves 189 TFLOPs/s vs. FlashAttention-2's 196 TFLOPs/s, only a 1.04× improvement), precisely because the occupancy problem is more severe at longer sequences.

Innovation 4: Split-Q as a Principle for Eliminating Inter-Warp Communication in Tiled Matrix Chains

The warp-level work partitioning change — splitting Q across warps instead of K and V — might appear to be a minor implementation detail, a simple swapping of which tensor gets partitioned. But it embodies a non-obvious design principle for tiled matrix chains: in a chain of matrix multiplies with element-wise operations interleaved (S = QK^T, P = f(S), O = PV), the reduction dimension should be kept local to each warp, and the batch dimensions should be partitioned.

To see why this is subtle, consider the two options:

Split-K: Each warp gets a slice of K and V, computes partial S and partial O, then warps must sum their partial O's. The reduction is necessary because each warp's partial O is (ignoring softmax for a moment) Q (K_w^T V_w) — but the full output is Q (sum_w K_w^T V_w) only if the softmax were linear, which it is not. The softmax normalization couples the warps' computations: the correct output is not the sum of per-warp softmax outputs because the softmax denominator depends on all columns. So the warps must share their max and sum statistics, and the partial outputs must be rescaled and summed using those shared statistics. Communication is mathematically required by the decomposition.

Split-Q: Each warp gets a slice of Q, computes S_w = Q_w K^T, applies softmax independently (since the softmax operates row-wise, and each warp owns complete rows of S_w), and computes O_w = softmax(S_w) V. These O_w are disjoint subsets of the output rows — they simply need to be concatenated, not summed. No communication is required because the rows are independent in both the matrix multiply and the softmax.

The intellectual move is recognizing that softmax is row-wise, so partitioning along the row dimension (the Q dimension) creates embarrassingly parallel subproblems, whereas partitioning along the column dimension (the K/V dimension) creates coupled subproblems that require communication. This is obvious in retrospect — softmax normalizes each row independently, so different rows don't interact — but it is not the natural first choice for a kernel engineer. The natural instinct when parallelizing a matrix multiply QK^T is to partition K (the larger dimension when sequence length is large) to get better load balance, which is exactly what FlashAttention's split-K does. FlashAttention-2 recognizes that load balance on the matmul is less important than avoiding the communication that the subsequent softmax would force.

The principle generalizes: in any tiled computation where the output has independent rows (or more generally, independent slices along some dimension), partition the work along that independent dimension and keep the reduction dimension local to each processing unit. This is the dual of the standard MapReduce pattern — it is "Map-only" parallelism that avoids the Reduce step entirely by choosing the partition to align with the independence structure of the computation.

The backward pass applies the same principle but with more complex dependencies. The backward pass involves five matrix multiplies with different input/output relationships, and the natural partition dimensions differ for different operations. The paper's warp-level design for the backward pass (Section 3.3) navigates these tradeoffs, avoiding split-K where possible while accepting some communication where the dependency structure makes it unavoidable.

The evidence for the impact of this change is embedded in the overall 2× speedup (Figures 4–6), but the paper does not provide an ablation isolating just the warp partitioning change from the other two improvements. This is a limitation — the reader cannot determine how much of the speedup comes from split-Q alone versus the algorithm tweaks versus the sequence-length parallelization. However, the qualitative argument is strong: the split-K scheme requires Br × d elements of shared memory traffic per warp per inner loop iteration (each warp writes its partial output and reads all others'), while split-Q requires zero for the output. For typical block sizes (Br = 128, d = 64, 4 warps), this is roughly 32 KB of shared memory traffic eliminated per inner loop iteration, which for long sequences (many column blocks) represents a substantial fraction of total execution time.

5. Experimental Analysis

Evaluation Methodology

• Dataset. The paper uses standard attention benchmarking — there is no natural language or vision dataset per se. Instead, synthetic tensors Q, K, V of varying sequence lengths are generated to isolate attention kernel performance. For end-to-end training experiments, the benchmark is GPT-style language model training throughput with standard hyperparameters (model sizes 1.3B and 2.7B parameters, sequence lengths 2k and 8k), following the Megatron-LM training recipe (Shoeybi et al., 2019). No downstream task accuracy is reported because FlashAttention-2 computes exact attention — the outputs are mathematically identical to standard attention, so accuracy is unchanged by construction.

• Base model(s). For micro-benchmarks, no model is used — raw attention forward and backward kernels are timed directly. For end-to-end training benchmarks, GPT3-style models with 1.3B and 2.7B parameters are used, with hidden dimension 2048, and head dimensions of either 64 (32 heads) or 128 (16 heads). These model sizes are chosen as representative of commonly trained medium-scale Transformers where attention becomes a meaningful fraction of total training time.

• Metrics. The primary metric is throughput in TFLOPs/s (tera floating-point operations per second) for both micro-benchmarks and end-to-end training. For micro-benchmarks, FLOPs are computed analytically: forward pass FLOPs = 4 × seqlen² × head_dim × num_heads (with causal mask, this is halved to account for only computing approximately half the attention matrix entries); backward pass FLOPs = 2.5 × forward FLOPs (reflecting 2 matmuls in forward vs. 5 matmuls in backward due to recomputation). End-to-end training FLOPs follow the Megatron-LM formula: 6 × seqlen × num_params + 12 × num_layers × hidden_dim × seqlen². A secondary metric is percentage of theoretical maximum FLOPs/s (A100 peak: 312 TFLOPs/s for FP16/BF16 matmul), which captures hardware utilization efficiency independent of absolute speed.

• Baselines. Four baselines are compared in micro-benchmarks:

  1. PyTorch standard attention: The default attention implementation in PyTorch that materializes the full N×N attention matrix in HBM, calls cuBLAS GEMM for S = QK^T, applies softmax, and calls GEMM for O = PV.
  2. FlashAttention (Dao et al., 2022): The original I/O-aware exact attention implementation (the "cutlass" variant as integrated into the flash-attention library).
  3. xformers (Lefaudeux et al., 2022): The FlashAttention implementation in Facebook's xformers library, also based on CUTLASS primitives.
  4. FlashAttention in Triton (Tillet et al., 2019): The Triton compiler-based implementation (from triton/python/tutorials/06-fused-attention.py), which first introduced the row-block-outer loop ordering and sequence-length parallelization ideas that FlashAttention-2 adopts.

  For end-to-end training, baselines are: without FlashAttention (standard PyTorch attention), FlashAttention, and FlashAttention-2.

• Generation budget / compute accounting. For micro-benchmarks, compute is measured in total FLOPs as described above. Fair comparison is achieved by measuring wall-clock time for the same mathematical operation (same tensor sizes, same precision FP16/BF16) and converting to TFLOPs/s. For end-to-end training, FLOPs are computed using the Megatron-LM formula, and throughput is measured as TFLOPs/s per GPU on 8×A100 80GB SXM4 GPUs. No generation budget in the LLM sense applies — this is kernel benchmarking, not model inference.

• Cross-validation / statistical protocol. There is no statistical significance testing or cross-validation reported. All micro-benchmark numbers appear to be single-run measurements (or averaged over a few runs — the paper does not specify). This is standard practice for GPU kernel benchmarking where runtime variance is typically low (<1–2%), but it means the reported speedup ratios should be interpreted as approximate rather than statistically precise. Block sizes are manually tuned per head dimension, with values chosen from {64, 128} × {64, 128} based on shared memory constraints.

Main Quantitative Results

Micro-Benchmarks: Attention Forward + Backward Speed

The headline result across all configurations is that FlashAttention-2 achieves approximately 2× the throughput of FlashAttention for combined forward + backward attention computation, as shown in Figure 4. The speedup is largest at moderate-to-long sequence lengths (2k–8k) where the parallelism and work partitioning improvements have the most impact, and somewhat smaller at very short sequences (512) where overhead dominates.

Without causal mask, head dimension 64 (Figure 4a). At sequence length 2k, FlashAttention-2 reaches 153 TFLOPs/s vs. 91 TFLOPs/s for FlashAttention — a 1.68× speedup. At 8k, FlashAttention-2 achieves 175 TFLOPs/s vs. 108 TFLOPs/s for FlashAttention — a 1.62× speedup. The PyTorch baseline at 8k is 45 TFLOPs/s, meaning FlashAttention-2 is 3.89× faster. The xformers baseline at 8k reaches 110 TFLOPs/s (1.59× slower than FlashAttention-2). The Triton implementation reaches 100 TFLOPs/s at 8k (1.75× slower). Notably, at the shortest tested sequence length (512), FlashAttention-2 is 132 TFLOPs/s vs. 91 TFLOPs/s for FlashAttention — a 1.45× speedup — showing the improvements matter even for moderate sequence lengths.

Without causal mask, head dimension 128 (Figure 4b). The pattern is similar but absolute throughput is higher for all methods due to larger matrix multiply dimensions. At 8k, FlashAttention-2 achieves 201 TFLOPs/s vs. 82 TFLOPs/s for FlashAttention — a 2.45× speedup, which is substantially larger than with head dimension 64. The gap between FlashAttention-2 and the Triton implementation also widens (201 vs. 95 TFLOPs/s at 8k, 2.12×). At 16k (the longest tested), FlashAttention-2 reaches 203 TFLOPs/s vs. 83 TFLOPs/s for FlashAttention (2.45×) and 98 for Triton (2.07×). The PyTorch baseline runs out of memory (OOM) at 16k for head dimension 64 but manages 86 TFLOPs/s at 16k for head dimension 128 (since fewer heads means less memory for the attention matrix), at which point FlashAttention-2 is 2.36× faster.

With causal mask, head dimension 64 (Figure 4c). Causal masking introduces additional complexity (masking approximately half the attention matrix and skipping entirely-masked blocks). At 8k, FlashAttention-2 reaches 171 TFLOPs/s vs. 92 TFLOPs/s for FlashAttention — a 1.86× speedup. The PyTorch baseline is 18 TFLOPs/s (9.5× slower). At 16k, FlashAttention-2 achieves 171 TFLOPs/s (nearly flat from 8k) while FlashAttention is at 97 TFLOPs/s — a 1.76× speedup. The Triton implementation at 16k with causal mask is 80 TFLOPs/s (2.14× slower). The absolute throughput with causal mask is approximately 15–20% lower than without causal mask for FlashAttention-2 (175 vs. 171 TFLOPs/s at 8k, head dim 64), reflecting the overhead of mask application and the fact that approximately half the FLOPs are skipped (so the theoretical peak for the same wall-clock time is lower).

With causal mask, head dimension 128 (Figure 4d). At 8k, FlashAttention-2 reaches 189 TFLOPs/s vs. 83 TFLOPs/s for FlashAttention — a 2.28× speedup. At 16k, FlashAttention-2 is 189 TFLOPs/s vs. 83 TFLOPs/s for FlashAttention (2.28×) and 67 TFLOPs/s for the Triton implementation (2.82×). The PyTorch baseline at 8k with causal mask and head dim 128 is 34 TFLOPs/s, making FlashAttention-2 5.56× faster.

Key pattern across all configurations: The speedup of FlashAttention-2 over FlashAttention is consistently larger at head dimension 128 than at head dimension 64 (e.g., 2.45× vs. 1.62× at 8k without causal mask). This is consistent with the paper's explanation that larger head dimensions benefit more from the warp-level split-Q partitioning because the output matrices being communicated in FlashAttention's split-K scheme are larger (Br × d proportional to d), so eliminating this communication saves proportionally more bandwidth.

Micro-Benchmarks: Attention Forward Speed Only

Figure 5 isolates the forward pass. FlashAttention-2 reaches up to 230 TFLOPs/s, which is 73% of the A100 theoretical maximum (312 TFLOPs/s). This is the paper's flagship utilization number — getting within striking distance of GEMM's typical 80–90% utilization.

Without causal mask, head dimension 64 (Figure 5a). At 8k, FlashAttention-2 achieves 192 TFLOPs/s vs. 104 TFLOPs/s for FlashAttention — a 1.85× speedup. At 16k, FlashAttention-2 reaches 192 TFLOPs/s (flat) vs. 104 TFLOPs/s for FlashAttention. The forward-only speedup is slightly higher than the forward+backward combined speedup, suggesting the backward pass improvements are somewhat less dramatic than the forward pass improvements (consistent with the paper's note that the backward pass is more complex and harder to fully optimize).

Without causal mask, head dimension 128 (Figure 5b). This is the highest-throughput configuration tested. At 8k, FlashAttention-2 achieves 224 TFLOPs/s vs. 72 TFLOPs/s for FlashAttention — a 3.11× speedup. At 16k, FlashAttention-2 is at 223 TFLOPs/s vs. 73 TFLOPs/s for FlashAttention — a 3.05× speedup. This is the configuration where FlashAttention-2 hits 73% of peak theoretical throughput (224/312 = 71.8%, though the paper rounds to 73% in the abstract, possibly using the 230 TFLOPs/s number from a slightly different configuration or GPU instance). The Triton implementation at 8k reaches 163 TFLOPs/s (1.37× slower than FlashAttention-2), and at 16k reaches 163 TFLOPs/s (1.37× slower).

With causal mask, head dimension 64 (Figure 5c). At 8k, FlashAttention-2 achieves 183 TFLOPs/s vs. 94 TFLOPs/s for FlashAttention — a 1.95× speedup. At 16k, FlashAttention-2 is 183 TFLOPs/s vs. 94 TFLOPs/s (1.95×). The Triton implementation at 8k with causal mask is 143 TFLOPs/s (1.28× slower).

With causal mask, head dimension 128 (Figure 5d). At 8k, FlashAttention-2 achieves 197 TFLOPs/s vs. 71 TFLOPs/s for FlashAttention — a 2.77× speedup. At 16k, FlashAttention-2 is 197 TFLOPs/s vs. 71 TFLOPs/s (2.77×). The Triton implementation at 8k is 148 TFLOPs/s (1.33× slower) and at 16k is 148 TFLOPs/s. The PyTorch baseline at 8k with causal mask and head dim 128 is a mere 19 TFLOPs/s, meaning FlashAttention-2 is approximately 10.4× faster.

A notable observation: FlashAttention-2's throughput is essentially flat from 8k to 16k in most configurations (e.g., 192→192 TFLOPs/s in Figure 5a, 224→223 in Figure 5b). This indicates that the kernel is compute-bound (not memory-bandwidth-bound) at these sequence lengths — adding more work increases total time proportionally, keeping throughput constant. In contrast, FlashAttention shows a slight decline from 8k to 16k (104→104 is flat, but 72→73 is flat), and the PyTorch baseline falls sharply (or goes OOM). This flat throughput profile is characteristic of well-optimized GEMM-like kernels and is evidence that FlashAttention-2 has successfully moved the bottleneck from memory bandwidth to compute.

Micro-Benchmarks: Attention Backward Speed Only

Figure 6 isolates the backward pass. FlashAttention-2 reaches up to 196 TFLOPs/s, which is approximately 63% of the A100 theoretical maximum. The backward pass is inherently harder to optimize than the forward pass (5 matmuls vs. 2, more intermediate values to manage in SRAM, more complex dependencies), so the lower peak utilization is expected. However, the speedup over FlashAttention remains approximately 2×.

Without causal mask, head dimension 64 (Figure 6a). At 8k, FlashAttention-2 achieves 170 TFLOPs/s vs. 112 TFLOPs/s for FlashAttention — a 1.52× speedup. At 16k, FlashAttention-2 is 170 TFLOPs/s vs. 113 TFLOPs/s for FlashAttention (1.50×). The backward-only speedup is noticeably smaller than the forward-only speedup (1.5× vs. 1.85×), consistent with the paper's acknowledgment that the backward pass is more complex and the split-Q partitioning is harder to apply cleanly due to the more intricate dependency structure between dQ, dK, dV.

Without causal mask, head dimension 128 (Figure 6b). At 8k, FlashAttention-2 achieves 196 TFLOPs/s vs. 88 TFLOPs/s for FlashAttention — a 2.23× speedup. At 16k, FlashAttention-2 is 196 TFLOPs/s vs. 88 TFLOPs/s (2.23×). The head dimension 128 backward pass speedup is larger than head dimension 64 (2.23× vs. 1.52×), mirroring the forward pass pattern and reinforcing that larger head dimensions benefit more from eliminating shared memory traffic.

With causal mask, head dimension 64 (Figure 6c). At 8k, FlashAttention-2 achieves 166 TFLOPs/s vs. 93 TFLOPs/s for FlashAttention — a 1.78× speedup. At 16k, FlashAttention-2 is 166 TFLOPs/s vs. 98 TFLOPs/s for FlashAttention (1.69×). The Triton implementation at 8k is 68 TFLOPs/s (2.44× slower than FlashAttention-2), and at 16k is 68 TFLOPs/s. This is the configuration where Triton shows its largest relative gap to FlashAttention-2, possibly due to Triton's compiler-generated code being less efficient for the backward pass's complex control flow.

With causal mask, head dimension 128 (Figure 6d). At 8k, FlashAttention-2 achieves 186 TFLOPs/s vs. 89 TFLOPs/s for FlashAttention — a 2.09× speedup. At 16k, FlashAttention-2 is 186 TFLOPs/s vs. 89 TFLOPs/s (2.09×). The PyTorch baseline at 8k with causal mask is 49 TFLOPs/s (3.80× slower).

A cross-comparison worth noting: FlashAttention-2's backward pass with causal masking is only marginally slower than without causal masking (170→166 TFLOPs/s for head dim 64, 196→186 for head dim 128, both at 8k). In contrast, FlashAttention's backward pass with causal mask is sometimes faster than without (112→93 for head dim 64 at 8k, but 88→89 for head dim 128). This inconsistency in FlashAttention's behavior with causal masking vs. FlashAttention-2's consistent slight slowdown is not discussed in the paper but may reflect the different work partitioning interacting with the causal mask's block-skipping logic.

H100 Preliminary Results

Figure 7 shows forward+backward speed on H100 GPUs (80GB SXM5), using the same kernel code with no H100-specific optimizations (no use of TMA, 4th-gen Tensor Cores, or FP8). FlashAttention-2 reaches up to 338 TFLOPs/s (without causal mask, head dimension 128, at 16k sequence length). This is substantially beyond the A100's peak theoretical throughput, demonstrating the raw hardware improvement from H100's higher clock speeds and memory bandwidth, even without architectural feature exploitation.

Without causal mask, head dimension 64 (Figure 7a): At 16k, FlashAttention-2 achieves 296 TFLOPs/s vs. 168 TFLOPs/s for FlashAttention — a 1.76× speedup. The gap between FlashAttention-2 and FlashAttention on H100 is similar to that on A100, suggesting the work partitioning improvements transfer across GPU generations.

Without causal mask, head dimension 128 (Figure 7b): At 16k, FlashAttention-2 achieves 338 TFLOPs/s vs. 139 TFLOPs/s for FlashAttention — a 2.43× speedup. This is the highest absolute throughput reported anywhere in the paper.

With causal mask, head dimension 64 (Figure 7c): At 16k, FlashAttention-2 achieves 284 TFLOPs/s vs. 156 TFLOPs/s for FlashAttention — a 1.82× speedup.

With causal mask, head dimension 128 (Figure 7d): At 16k, FlashAttention-2 achieves 328 TFLOPs/s vs. 137 TFLOPs/s for FlashAttention — a 2.39× speedup.

The paper notes that these numbers are achieved without using H100-specific features (TMA for asynchronous memory copies, 4th-gen Tensor Cores, FP8 support) and expects "another 1.5x-2x speedup" from using these features, which would push FlashAttention-2 toward 500–670 TFLOPs/s on H100 — well beyond the A100's theoretical peak and into the territory where attention becomes competitive with GEMM even in absolute terms.

End-to-End Training Throughput

Table 1 reports training throughput for GPT-style models on 8×A100 80GB SXM4 GPUs. FlashAttention-2 reaches up to 225 TFLOPs/s per GPU (72% model FLOPs utilization) for GPT3-2.7B at 8k context. The key comparisons:

GPT3-1.3B, 2k context: FlashAttention-2 achieves 196 TFLOPs/s vs. 189 TFLOPs/s for FlashAttention — a modest 1.04× speedup (3.7% improvement). The baseline without FlashAttention is 142 TFLOPs/s (1.38× slower than FlashAttention-2). At 2k context with 1.3B parameters, attention is not the dominant bottleneck (other operations like MLP layers, layer norm, and communication for data parallelism consume a larger fraction of time), so the 2× attention speedup translates to only a 4% end-to-end gain.

GPT3-1.3B, 8k context: FlashAttention-2 achieves 220 TFLOPs/s vs. 170 TFLOPs/s for FlashAttention — a 1.29× speedup. The baseline without FlashAttention is 72 TFLOPs/s (3.06× slower). At 8k context, attention dominates training time (the quadratic attention cost becomes the bottleneck), so the attention speedup translates more directly to end-to-end gain. The jump from 2k to 8k context for FlashAttention-2 is from 196 to 220 TFLOPs/s — actually an increase in throughput at longer sequence length, which is counterintuitive (longer sequences typically reduce throughput because attention cost grows quadratically) but reflects the fact that at 8k the GPU is better utilized (higher occupancy from sequence-length parallelization).

GPT3-2.7B, 2k context: FlashAttention-2 achieves 205 TFLOPs/s vs. 189 TFLOPs/s for FlashAttention — a 1.08× speedup (8.5% improvement). The baseline is 149 TFLOPs/s (1.38× slower).

GPT3-2.7B, 8k context: FlashAttention-2 achieves 225 TFLOPs/s vs. 175 TFLOPs/s for FlashAttention — a 1.29× speedup. This is the highest end-to-end throughput reported and corresponds to 72% model FLOPs utilization. The baseline is 80 TFLOPs/s (2.81× slower). Notably, the end-to-end throughput at 8k context is higher than at 2k context (225 vs. 205 TFLOPs/s), which again indicates better GPU utilization at longer sequences. This 225 TFLOPs/s number is very close to the forward-only micro-benchmark throughput (224 TFLOPs/s in Figure 5b), which suggests that in the end-to-end training setting with 8k context, the attention forward pass is the throughput-limiting operation and it is running at near-peak micro-benchmark speed, while other operations (backward pass, MLP, communication) largely overlap or are not bottlenecks.

Critical interpretation of end-to-end numbers: The 2.8× end-to-end speedup vs. no-FlashAttention (80→225 TFLOPs/s at 2.7B, 8k) is larger than the pure attention speedup (~2×) because the no-FlashAttention baseline suffers from memory bottlenecks that force smaller batch sizes or gradient accumulation, reducing overall utilization. The comparison to FlashAttention is the more informative one for evaluating FlashAttention-2: 1.29× end-to-end at 8k context, which is substantial (saving ~23% of training time) but notably less than the 2× attention micro-benchmark speedup. The gap between micro-benchmark speedup (2×) and end-to-end speedup (1.29×) reflects Amdahl's Law: attention is not the only operation in the model, so speeding it up yields diminishing returns on total throughput. At 2k context, attention is even less dominant, and the end-to-end gain drops to 1.04–1.08×.

Ablation Studies and Robustness Checks

Head dimension 64 vs. 128: The paper does not present a formal ablation but the consistent reporting of both head dimensions across all configurations (Figures 4–6) constitutes an implicit sensitivity analysis. As noted above, FlashAttention-2's speedup over FlashAttention is consistently larger at head dimension 128 than at 64 (e.g., forward pass without causal mask at 8k: 3.11× vs. 1.85×). This is explained by the warp-level work partitioning change: FlashAttention's split-K inter-warp communication volume scales with d (each warp writes Br × d elements to shared memory), so eliminating this communication saves proportionally more at larger d. The paper does not explore whether even larger head dimensions (e.g., 256) would show further speedup amplification or whether register pressure would eventually limit the benefit.

Causal mask vs. no causal mask: The paper reports both settings for all configurations, revealing that FlashAttention-2's speedup over FlashAttention is similar with and without causal masking (e.g., Figure 4: 1.62× without causal mask at 8k head dim 64 vs. 1.86× with causal mask). The causal mask overhead (lower absolute TFLOPs/s) is approximately 5–15% for FlashAttention-2 (e.g., forward pass head dim 64 at 8k: 192→183 TFLOPs/s, a 4.7% reduction), which is expected since approximately half the FLOPs are skipped and the remaining computation includes mask application overhead. The fact that the relative speedup vs. FlashAttention is similar in both settings suggests that the work partitioning improvements do not interact negatively with the causal masking logic.

Sequence length scaling (512 to 16k): The paper sweeps sequence lengths from 512 to 16k in powers of 2. All methods show throughput that increases with sequence length up to a saturation point and then flattens. For FlashAttention-2 without causal mask and head dim 128 (Figure 5b), the forward pass scales from 127 TFLOPs/s at 512 to 224 TFLOPs/s at 8k, then remains flat to 16k. For FlashAttention, the scaling is 69→72→73 TFLOPs/s over the same range — it saturates much earlier and at a much lower level. This confirms that FlashAttention-2's sequence-length parallelization (Section 3.2) is effective at maintaining high utilization even at short sequences where FlashAttention suffers from low occupancy: at 512, FlashAttention-2's forward throughput is already 127 TFLOPs/s vs. 69 for FlashAttention (1.84×), showing the occupancy benefit matters even at moderate lengths.

Block size tuning (mentioned in Section 3.3, not ablated): The paper states that block sizes are chosen from {64, 128} × {64, 128} depending on head dimension and shared memory, and are manually tuned. There is no sensitivity analysis showing how performance varies with different block size choices, how often register spilling or shared memory overflow constrains the choice, or whether auto-tuning would find better configurations. This is acknowledged as future work. For reproducibility, the specific block sizes used for each configuration in Figures 4–7 are not reported, which is a weakness.

H100 without vs. with new hardware features: The H100 results (Figure 7) use the same kernel code as A100 with no special H100 instructions. The paper speculates that using TMA, 4th-gen Tensor Cores, and FP8 could yield "another 1.5x-2x speedup." This is not an ablation (since the optimized H100 kernel does not exist yet) but a projected expectation. The H100 results serve primarily to demonstrate that FlashAttention-2's improvements are not A100-specific — they transfer to newer hardware with similar relative speedups. The absolute H100 throughput (up to 338 TFLOPs/s) demonstrates that FlashAttention-2 can already exploit newer hardware's higher raw performance without code changes.

Multi-query and grouped-query attention (mentioned in Section 3.1.2, no experiments): The paper describes how FlashAttention-2 supports MQA and GQA by implicit head index manipulation and gradient summation in the backward pass, but provides no benchmarks for these variants. Given that MQA and GQA are increasingly common (especially for inference), this is a notable omission — it is unclear whether the 2× speedup carries over to these attention variants or whether the different K/V sharing patterns affect the work partitioning benefits.

Numerical correctness (not ablated, asserted): The paper does not report any numerical error measurements (e.g., max absolute error vs. PyTorch reference, mean squared error). It asserts that FlashAttention-2 computes exact attention "with no approximation" and references the same proof as FlashAttention (Dao et al., 2022, Theorem 1). While this is theoretically sound — the unscaled accumulator and logsumexp storage are mathematically equivalent transformations — floating-point arithmetic is not associative, and the different order of operations (rescaling only at the end vs. rescaling every iteration) could in principle produce slightly different numerical results. The absence of any numerical validation is a gap, though in practice the differences are likely negligible for FP16/BF16 training.

Critical Assessment

Claim 1: "FlashAttention-2 yields around 2× speedup compared to FlashAttention" (abstract, Section 3). The micro-benchmark evidence strongly supports this claim, but with important qualification about when the 2× figure applies. The speedup varies substantially by configuration: it is as low as 1.45× (forward+backward, no causal mask, head dim 64, sequence length 512; Figure 4a) and as high as 3.11× (forward only, no causal mask, head dim 128, sequence length 8k; Figure 5b). The "around 2×" characterization is a reasonable central tendency, but a practitioner looking at a specific configuration needs to consult the specific figure. The speedup is systematically larger for head dimension 128 than 64, systematically larger for forward-only than backward-only, and relatively stable across sequence lengths above 1k. The paper would be stronger if it provided a simple table summarizing the speedup range across all configurations rather than requiring the reader to extract numbers from bar charts.

Claim 2: "reaching 50-73% of the theoretical maximum FLOPs/s on A100" (abstract). This claim is supported but the range "50-73%" obscures that the 73% figure applies only to the forward pass without causal mask and with head dimension 128 (Figure 5b: 224/312 = 71.8%, or possibly 230 TFLOPs/s from a slightly different configuration = 73.7%). The backward pass peaks at 63% (196/312, Figure 6b). The forward+backward combined peaks at approximately 65% (203/312 for head dim 128 without causal mask at 16k, Figure 4b). The lower bound of "50%" appears to correspond to configurations with causal masking or head dimension 64. This is still substantially better than FlashAttention's typical 25–40% utilization, but the abstract's broad range should not be misinterpreted as FlashAttention-2 always achieving >70% — that is a best-case forward-pass-only number.

Claim 3: "getting close to the efficiency of GEMM operations" (abstract). This claim is difficult to evaluate because the paper does not benchmark GEMM throughput under identical conditions (same matrix dimensions, same GPU, same precision). The paper states that "optimized GEMM can reach up to 80-90% of the theoretical maximum device throughput" (Section 1), and FlashAttention-2's 73% forward-pass peak is indeed "close" to the lower end of that range. But the gap from 73% to 80% is still meaningful — it represents approximately 10% additional headroom. More importantly, the backward pass peak of 63% is not particularly close to GEMM efficiency, and since training spends more time in the backward pass than forward (due to the 5-vs-2 matmul ratio), the overall training-time attention efficiency is weighted more toward the backward pass performance. The claim is directionally true but somewhat optimistic for the training-relevant case.

Claim 4: "reaches training speed of up to 225 TFLOPs/s per A100 GPU (72% model FLOPs utilization)" (abstract, Table 1). This is the most practically meaningful claim, and it is supported by the end-to-end training numbers in Table 1. However, it is only demonstrated for one model configuration (GPT3-2.7B, 8k context). The 1.3B model at 2k context achieves only 196 TFLOPs/s (63% utilization), showing that the high utilization depends on long sequences making attention the dominant bottleneck. Additionally, the model FLOPs utilization metric (72%) depends on the FLOP counting formula (Section 4.2), which the paper notes is the Megatron-LM convention that does not halve the attention FLOPs for causal masking. If attention FLOPs were halved (reflecting that only half the attention matrix is computed), the utilization percentage would be lower. The paper's choice to follow the established convention is defensible for comparability, but readers should understand that "72% utilization" is relative to a specific (somewhat generous) FLOP counting methodology.

What the experiments do not show:

1. No comparison to approximate attention methods. The paper benchmarks only exact attention implementations (PyTorch, FlashAttention, xformers, Triton). There is no comparison to approximate methods (Performer, Reformer, Linformer, Big Bird, etc.) in terms of either speed or downstream model quality. The paper argues these are not widely used in large-scale training (Section 1), but providing even a single representative comparison would strengthen the case that exact attention with FlashAttention-2 is now fast enough to make approximations unnecessary for most use cases.

2. No multi-GPU scaling results. All benchmarks are single-GPU (or 8-GPU for end-to-end, but reported as per-GPU throughput). There is no analysis of how FlashAttention-2 interacts with model parallelism, tensor parallelism, or pipeline parallelism — all of which are standard in large-scale training. Does sequence-length parallelization conflict with tensor parallelism's partitioning of attention heads? This matters for the large training runs the paper is targeting.

3. No memory usage measurements. The paper focuses entirely on speed (TFLOPs/s) and does not report memory consumption. FlashAttention's headline feature was linear memory scaling in sequence length. FlashAttention-2 inherits this property (the algorithm is mathematically equivalent, so memory requirements are unchanged), but the additional sequence-length parallelization means more thread blocks are active simultaneously, potentially increasing the total SRAM and register footprint. The paper does not verify that memory usage is unchanged or report peak memory for the configurations tested.

4. No FP8 or BF16-specific optimizations beyond what Tensor Cores provide automatically. The paper uses FP16/BF16 inputs but does not explore FP8, which the H100 supports natively and would provide additional throughput. The H100 results (Figure 7) explicitly do not use FP8. This is acknowledged as future work.

5. No validation on non-NVIDIA GPUs. All results are on A100 and H100. The paper mentions AMD GPUs as future work in Section 5 but provides no data. The work partitioning principles (split-Q, sequence-length parallelization) should in principle transfer to AMD's GPU architecture, but the specific throughput gains depend on AMD's warp/wavefront size, shared memory bandwidth, and matrix multiply throughput, which differ from NVIDIA's.

6. Statistical rigor. No error bars, no multiple-run averaging reported, no variance estimates. For GPU kernel benchmarks where variance is typically low, this is common practice, but it means the exact speedup ratios (e.g., "1.68×") should not be taken as precise to two decimal places. A more careful paper would report measurement methodology (number of runs, warmup iterations, whether the first kernel launch is excluded to avoid JIT compilation overhead).

Missing experiment: Isolating individual contributions. The paper presents FlashAttention-2 as a single combined system with three improvements, but provides no ablation that isolates the contribution of each. We cannot determine whether the 2× speedup comes primarily from the algorithm tweaks (unscaled accumulator, logsumexp storage), the sequence-length parallelization, or the warp-level split-Q partitioning — or whether all three are necessary to achieve the full gain. An ablation that tested, for example, FlashAttention with just the split-Q change (keeping the original algorithm and no sequence-length parallelization) would reveal how much each improvement matters independently. The paper's claim that all three are important is plausible (they operate at different levels of the GPU hierarchy and should be roughly multiplicative), but unverified.

Missing baseline: Vanilla attention with torch.compile. The PyTorch baseline is described as "standard attention implementation in PyTorch" but it is unclear which PyTorch version is used and whether torch.compile or torch.nn.functional.scaled_dot_product_attention (which includes a FlashAttention backend in recent PyTorch versions) is enabled. If the baseline uses an older, unoptimized code path, the 3–10× speedup over "PyTorch" may be overstated relative to what a practitioner using current PyTorch would experience.

Where claims hold conditionally:

• The 2× speedup over FlashAttention holds broadly but is larger for larger head dimensions, for the forward pass specifically, and for moderate-to-long sequences (2k+). At very short sequences (512), the speedup is closer to 1.5×.
• The 73% peak utilization holds only for forward pass, large head dimension (128), and without causal masking. The backward pass peaks at 63%, and causal masking reduces both by 5–15%.
• The end-to-end 1.29× training speedup over FlashAttention holds at 8k context but drops to 1.04–1.08× at 2k context, meaning the practical benefit is concentrated in the long-context regime that the paper is targeting. For short-context training, the upgrade from FlashAttention to FlashAttention-2 yields marginal improvement.

Overall assessment: The experiments are thorough for a GPU kernel paper — they cover four attention configurations (with/without causal mask × head dim 64/128) across six sequence lengths on two GPU architectures, plus end-to-end training at two model sizes. The consistent 1.5–3× speedup over FlashAttention across all micro-benchmark configurations leaves little doubt that the combined improvements are real and substantial. The main weaknesses are the lack of ablation to attribute the speedup to specific changes, the absence of memory usage measurements, and the reliance on a single implementation (no independent reproduction or diverse hardware targets). These are not fatal — the paper's primary audience is practitioners who will evaluate FlashAttention-2 by running it on their own workloads, and the open-source release and wide adoption since publication provide de facto validation that the claimed speedups materialize in practice.

6. Limitations and Trade-offs

6.1 No Isolation of Individual Contributions — the 2× Speedup Is a Combined Effect of Three Changes Whose Relative Importance Is Unknown

The assumption or constraint. FlashAttention-2 presents three independent improvements — the unscaled accumulator algorithm tweak (Section 3.1.1), sequence-length parallelization (Section 3.2), and split-Q warp partitioning (Section 3.3) — but evaluates them only as a single combined system. The paper provides no ablation that isolates how much of the 2× speedup comes from each change. The authors do not report, for example, "FlashAttention with only the unscaled accumulator," or "FlashAttention with only sequence-length parallelization," or "FlashAttention with only split-Q partitioning."

The consequence. A practitioner cannot determine which improvement to prioritize if they can only implement a subset, or which change is most responsible for the speedup in their specific regime. This matters concretely: implementing the algorithm tweak (unscaled accumulator) requires modifying only the inner loop arithmetic and is relatively self-contained, while adding sequence-length parallelization requires restructuring the CUDA kernel launch grid and may interact with existing parallelism strategies (tensor parallelism, pipeline parallelism). If the unscaled accumulator alone provides 1.5× and the sequence-length parallelization adds only 1.1×, a team might implement just the algorithm change and skip the more invasive parallelism changes. Conversely, if sequence-length parallelization is the dominant factor, teams training at small batch sizes should prioritize that and might deprioritize the unscaled accumulator. Without ablation, there is no way to make this judgment from the paper. The three changes operate at different GPU hierarchy levels and the paper implies they are multiplicative, but this is unverified — they could be partially redundant (e.g., both the unscaled accumulator and split-Q reduce shared memory traffic in overlapping ways).

What evidence exists in the paper. The paper provides no ablation evidence whatsoever. All benchmark results (Figures 4–6) and the end-to-end training results (Table 1) compare the full FlashAttention-2 against FlashAttention and other baselines, never against partial variants. The consistent observation that speedup is larger for head dimension 128 than 64 (e.g., 2.45× vs. 1.62× forward+backward at 8k without causal mask; Figure 4) is indirect evidence that the warp-level split-Q change matters — since shared memory communication volume in FlashAttention's split-K scheme scales with head dimension — but this is correlational, not causal isolation. Similarly, the larger speedup at longer sequences (where occupancy matters more) is indirect evidence for the sequence-length parallelization benefit, but again does not isolate it.

Mitigation status. Not addressed at all. The paper does not acknowledge this as a limitation, nor does it suggest ablation experiments as future work. This is standard practice in GPU kernel papers (where implementations are monolithic and partial variants are rarely tested), but it weakens the paper's diagnostic claims: the paper asserts that three specific defects explain FlashAttention's inefficiency, but never verifies that fixing each one individually produces the expected benefit.

6.2 Backward Pass Remains Substantially Less Efficient Than Forward — 63% vs. 73% of Peak, and the Gap Limits End-to-End Training Gains

The assumption or constraint. FlashAttention-2 achieves up to 73% of theoretical peak FLOPs/s in the forward pass (230 TFLOPs/s, head dim 128, no causal mask; Figure 5b) but only 63% in the backward pass (196 TFLOPs/s; Figure 6b). The paper acknowledges this asymmetry implicitly by reporting forward-only, backward-only, and combined benchmarks separately, but does not analyze why the backward pass lags or whether the specific optimizations in FlashAttention-2 are less effective for the backward pass. Section 3.3 notes that the backward pass "still requires some synchronization due to the more complicated dependency" but does not quantify how much performance is left on the table by these residual synchronization costs.

The consequence. Training — which is the primary use case the paper targets (long-context model training) — spends more time in the backward pass than the forward pass because the backward pass involves 5 matrix multiplies compared to 2 in the forward pass. The paper's FLOP accounting reflects this: backward FLOPs = 2.5 × forward FLOPs (Section 4.1). The lower backward pass utilization means that the training-weighted attention efficiency is closer to the backward pass number (63%) than the forward pass number (73%). Concretely, for the forward+backward combined metric that matters for training throughput, FlashAttention-2 peaks at approximately 65% of theoretical peak (203 TFLOPs/s at head dim 128, no causal mask, 16k; Figure 4b), not the 73% headline number. This leaves a substantial gap to GEMM's 80–90% range. The end-to-end training results reflect this: FlashAttention-2 achieves 72% model FLOPs utilization at 8k context (Table 1), which accounts for non-attention operations, but the attention-specific utilization during training is lower than the forward-only micro-benchmark would suggest.

What evidence exists in the paper. Figure 6 systematically shows lower backward pass throughput than Figure 5 for forward pass, for all configurations. At the best-case configuration (head dim 128, no causal mask, 8k), the forward pass reaches 224 TFLOPs/s (Figure 5b) while the backward pass reaches 196 TFLOPs/s (Figure 6b) — a 14% gap. At the less favorable configuration (head dim 64, no causal mask, 8k), forward is 192 TFLOPs/s (Figure 5a) vs. backward 170 TFLOPs/s (Figure 6a) — a 13% gap. The speedup over FlashAttention is also systematically smaller in the backward pass: at head dim 64, no causal mask, 8k, the forward speedup is 1.85× (192 vs. 104 TFLOPs/s) while the backward speedup is 1.52× (170 vs. 112 TFLOPs/s). This confirms that the FlashAttention-2 optimizations are less effective for the backward pass, but the paper does not analyze which specific backward-pass operations are the bottleneck or whether the split-Q partitioning, algorithm tweaks, or sequence-length parallelization is less impactful there.

Mitigation status. Not addressed. The paper notes the backward pass complexity (Section 3.3: "it still requires some synchronization") but offers no analysis of the residual bottlenecks, no profiling data showing where backward pass time is spent, and no plan for closing the forward-backward gap. The paper's future work (Section 5) focuses on H100 features and FP8 support, not on backward pass optimization. This is a meaningful gap because closing the backward pass from 63% to even 75% utilization would likely yield a larger end-to-end training improvement than further forward pass optimization.

6.3 Difficulty Estimation Overhead for Choosing Block Sizes and Launch Configurations Is Manual and Unaccounted For

The assumption or constraint. FlashAttention-2 requires manual tuning of block sizes Br and Bc for each head dimension d and each GPU architecture. Section 3.3 states: "We manually tune for each head dimensions since there are essentially only 4 choices for block sizes, but this could benefit from auto-tuning to avoid this manual labor." Additionally, the kernel launch configuration — specifically the grid dimensions that implement the sequence-length parallelization (B × H × Tr for forward, B × H × Tc for backward) — must be computed based on sequence length, batch size, number of heads, and chosen block sizes. The paper does not discuss the cost of this pre-computation or whether it adds measurable overhead.

The consequence. For a practitioner integrating FlashAttention-2 into a training framework, the manual tuning is a deployment friction: they must determine the optimal block sizes for their specific head dimension and GPU, either by trusting the paper's defaults (which are not explicitly reported for each configuration) or by running their own micro-benchmarks. This is not a one-time cost if they train models with different architectures (different head dimensions) on different GPU types. The paper mentions only 4 choices for block sizes ({64, 128} × {64, 128}), but the interaction with sequence length, batch size, and the sequence-length parallelization grid dimensions creates a larger configuration space. A poor block size choice can cause either register spilling (too large) or excessive HBM transfers (too small), potentially negating a substantial fraction of the reported speedup. The end-to-end training results (Table 1) presumably use well-tuned block sizes, but the paper does not specify which sizes were used for which model, making the results harder to reproduce without guesswork.

What evidence exists in the paper. The paper explicitly acknowledges the manual tuning in Section 3.3, calling it out as a limitation that auto-tuning could address. However, no sensitivity analysis is provided to show how much performance degrades with suboptimal block sizes. For example, if a practitioner uses Br = Bc = 64 instead of the optimal 128 for head dim 128, does throughput drop by 5% or 30%? Without this information, the robustness of the reported numbers to block size choices is unknown. The H100 results (Figure 7) use the same block sizes as A100, and the paper notes that H100-specific tuning could yield additional speedup, but does not quantify the gap. This implies that even across GPU generations, optimal block sizes may differ, increasing the tuning burden.

Mitigation status. The paper suggests auto-tuning as future work (Section 3.3), which is a reasonable direction. However, auto-tuning itself has a cost — it requires running many kernel configurations and measuring their performance, which takes time and GPU resources. The paper does not discuss whether auto-tuning overhead would be amortizable over a long training run or whether it would need to be repeated for each model configuration. Modern frameworks like Triton and TVM include auto-tuning infrastructure, but integrating FlashAttention-2's CUDA kernels into such frameworks is non-trivial. The limitation is partially mitigated by the small search space (4 choices), but the interaction with the sequence-length parallelization grid dimensions suggests the effective search space may be larger than implied.

6.4 No Multi-GPU or Model Parallelism Evaluation — Interaction with Tensor, Pipeline, and Data Parallelism Is Uncharacterized

The assumption or constraint. All experiments in the paper are conducted on single GPUs (micro-benchmarks) or on 8 GPUs with what appears to be data parallelism only (end-to-end training, Table 1). The paper does not evaluate FlashAttention-2 in the context of tensor parallelism (where attention heads are split across GPUs), pipeline parallelism (where layers are split across GPUs), or sequence parallelism (where the sequence dimension is partitioned across GPUs, as in e.g., Megatron-LM's sequence parallelism or DeepSpeed-Ulysses). Section 4.2 notes that the end-to-end training experiments use 8×A100 GPUs, but the per-GPU throughput numbers (e.g., 225 TFLOPs/s per GPU) are reported as averages, suggesting data-parallel training where each GPU processes a different batch — not model-parallel training where GPUs cooperate on the same batch.

The consequence. Large-scale training runs — the very use case the paper targets with its long-context motivation — almost always use some form of model parallelism. GPT-4-scale models use tensor parallelism within nodes, pipeline parallelism across nodes, and data parallelism across both. The interaction between FlashAttention-2's sequence-length parallelization (which launches B × H × Tr thread blocks) and tensor parallelism (which partitions H across GPUs, reducing the per-GPU H) is particularly concerning: tensor parallelism with degree T reduces the per-GPU head count to H/T, which in turn reduces the thread block count from B × H × Tr to B × (H/T) × Tr. Since sequence-length parallelization was specifically introduced to compensate for small B × H (Section 3.2), reducing H via tensor parallelism works against this fix. The paper does not analyze whether the resulting thread block count remains sufficient to occupy all SMs when tensor parallelism is used.

Additionally, sequence parallelism (partitioning the sequence length dimension across GPUs) would directly conflict with FlashAttention-2's sequence-length parallelization: both attempt to parallelize the same dimension but at different granularities (GPU-level vs. SM-level). Whether these can coexist, and whether the combined parallelism yields additional speedup or causes resource contention, is unknown. The paper's silence on these interactions means that practitioners using model parallelism cannot predict FlashAttention-2's performance in their actual training setup from the reported numbers.

What evidence exists in the paper. None. The paper does not mention tensor parallelism, pipeline parallelism, or sequence parallelism. The end-to-end experiments (Table 1) describe "GPT-style models on 8×A100 GPUs" but do not specify the parallelism strategy. The fact that throughput is reported as "TFLOPs/s/GPU" (per GPU, not aggregate) is consistent with data parallelism where each GPU independently processes its batch, but this is not stated. For GPT3-2.7B at 8k context, 8 GPUs with data parallelism and a batch size of 1 per GPU would give a global batch size of 8 — which is plausible but unconfirmed, and the paper does not discuss whether this batch size fits in memory or requires gradient accumulation.

Mitigation status. Not addressed. The paper does not flag multi-GPU interaction as a limitation or suggest it as future work. This is a significant omission given the paper's explicit framing around long-context training, where model parallelism is often necessary because even a single sequence's activations may exceed a single GPU's memory. The open-source release of FlashAttention-2 may allow practitioners to test these interactions themselves, but the paper provides no guidance or baseline expectations.

6.5 Numerical Error Analysis Under Floating-Point Non-Associativity Is Absent — the Unscaled Accumulator and Logsumexp Reorder Operations, Potentially Changing Bits

The assumption or constraint. FlashAttention-2 computes exact attention in the mathematical sense — the algorithm is algebraically equivalent to standard softmax attention. However, floating-point arithmetic is not associative: (a + b) + c ≠ a + (b + c) in finite precision. FlashAttention-2 changes the order of operations relative to both standard attention and FlashAttention in two ways: (1) the unscaled accumulator defers the division by ℓ to the end, meaning that the partial contributions e^{S - m} V are summed over many column blocks before being divided, potentially accumulating larger intermediate values that lose precision; (2) the logsumexp L = m + log(ℓ) replaces separate storage of m and ℓ, combining them via a logarithm that introduces its own rounding error. The paper does not report any numerical error measurements — no comparison of FlashAttention-2 outputs against a double-precision reference, no max absolute error, no mean squared error, no cosine similarity between attention outputs.

The consequence. For the typical use case (FP16/BF16 training), small numerical differences are unlikely to affect model quality — training is robust to the level of noise introduced by non-associativity. However, there are scenarios where bit-exact reproducibility matters: debugging (comparing against a reference implementation to isolate bugs), high-precision inference (some applications use FP32 attention for numerical stability), and scientific computing workloads where attention is used as a primitive in non-neural-network contexts. In these scenarios, the absence of any numerical characterization means the practitioner cannot assess whether FlashAttention-2's outputs are "close enough" for their tolerance. More subtly, different operation ordering could interact with the particular pattern of values in attention matrices (which often have extreme dynamic range due to the exponential in softmax), potentially producing larger errors than naive operation counting would suggest.

What evidence exists in the paper. None. The paper does not report any error metrics. The correctness claim rests entirely on the algebraic proof in Dao et al. (2022, Theorem 1), which the paper states "is almost the same" for FlashAttention-2. But the original proof assumes real arithmetic, not floating-point, and the specific operation reorderings in FlashAttention-2 (especially the unscaled accumulator) are not covered by that proof's analysis. This is a gap: the paper claims "no approximation" (which is true algebraically) but does not verify that the floating-point realization is numerically accurate.

Mitigation status. Not addressed. The paper does not acknowledge floating-point non-associativity as a concern. In practice, the FlashAttention-2 open-source implementation has been widely deployed in production training pipelines (GPT-4, Llama, etc.) without reported numerical issues, providing de facto validation that the errors are negligible for FP16/BF16 training. However, the paper itself provides no such evidence. A minimal mitigation would be to report the max absolute error vs. a double-precision reference for a few representative tensor sizes — a standard practice in numerical linear algebra papers that this paper omits.

6.6 Latency vs. Throughput Trade-off from Sequence-Length Parallelization Is Uncharacterized — More Parallelism May Increase Inference Latency for a Single Query

The assumption or constraint. The paper evaluates attention exclusively in terms of throughput (TFLOPs/s), which is the appropriate metric for training where many sequences are processed in parallel. The sequence-length parallelization (Section 3.2) increases parallelism by launching more thread blocks, which improves throughput by saturating the GPU's SMs. However, for inference — particularly single-query inference where a user sends one prompt and waits for the response — the relevant metric is latency (wall-clock time to process one attention operation), not throughput. The paper does not report latency numbers for any configuration. The benchmark setup (Section 4.1) fixes the total number of tokens to 16k and varies sequence length, meaning that shorter sequences have larger batch sizes — this is a throughput-optimizing benchmark design, not a latency-measuring one.

The consequence. During autoregressive decoding (the dominant inference paradigm for LLMs), attention is computed one query token at a time (or one chunk of query tokens) against the growing KV cache. The sequence length N grows from 1 to the full context length over the course of generation. At short sequence lengths (early tokens in generation), Tr (the number of row blocks) is small (e.g., at N=128 with Br=64, Tr=2). The sequence-length parallelization benefit is minimal at best, and the overhead of launching many near-empty thread blocks could increase latency compared to a kernel without sequence-length parallelization. Additionally, the kernel launch overhead itself (CPU→GPU kernel launch latency, typically 5–10 microseconds) becomes a larger fraction of total time for short sequences. At longer sequence lengths (late tokens in generation), sequence-length parallelization becomes beneficial, but the latency-critical inference scenario often involves many short-sequence steps (one per token) rather than a single long-sequence step.

The paper's emphasis on "long-context training" is appropriate for its scope, but FlashAttention is also widely used for inference (it reduces the KV cache memory footprint and speeds up the attention computation during decoding). A practitioner evaluating FlashAttention-2 for inference needs latency numbers — especially for the common deployment pattern of processing one prompt with batch size 1 — and the paper provides none. It is possible that FlashAttention-2's additional parallelism increases single-query latency compared to FlashAttention due to kernel launch overhead and thread block scheduling overhead, even while improving throughput for batched processing.

What evidence exists in the paper. The micro-benchmarks (Figures 4–6) are structured throughput experiments: they fix total tokens at 16k and adjust batch size as sequence length varies (e.g., at seqlen=512, batch=32; at seqlen=16k, batch=1). Reported speeds are TFLOPs/s, which is throughput, not time-per-operation. The lowest sequence length tested is 512 — typical autoregressive inference spends most tokens at lengths much shorter than 512 (e.g., generating 100 tokens from a 2k prompt means most attention operations are at lengths 2000–2100). The paper provides no data for very short sequences (N=1, 2, 4, 8, 16, 32, 64, 128, 256) that would characterize the latency-critical regime of early decoding. The paper's causal masking results (Figures 4c–d, 5c–d, 6c–d) are relevant to autoregressive inference, but they still use the same throughput-optimizing batch-size adjustment.

Mitigation status. Not addressed. The paper does not discuss inference latency, does not provide single-query latency measurements, and does not flag this as a limitation. The abstract and introduction frame the work around training ("scaling Transformers to longer sequence lengths," "train models with 16k longer context"), which is the primary intended use case. However, the paper also states that "FlashAttention-2 will also speed up training, finetuning, and inference of existing models" (Section 5), making a claim about inference speedup that is not supported by latency measurements. Practitioners deploying FlashAttention-2 for inference must benchmark latency themselves — the paper provides no guidance.

7. Implications and Future Directions

How This Work Changes the Landscape

FlashAttention-2 changes the landscape not by introducing a new attention algorithm, but by demonstrating that exact attention can be implemented at a throughput approaching optimized matrix multiplication — the gold standard for GPU kernel efficiency. This is a diagnostic and engineering contribution, not a conceptual paradigm shift. But its practical implications are substantial because it removes the remaining throughput penalty that made exact attention feel "expensive" relative to other model operations.

The methodological shift is in how the field should think about attention kernel optimization. Before FlashAttention-2, the narrative was: FlashAttention solved the memory bottleneck, and the remaining 25–40% utilization was a tolerable gap — attention would always be somewhat slower than GEMM because of the element-wise softmax operations and complex data dependencies. FlashAttention-2 refutes this: by systematically diagnosing three specific parallelism defects (arithmetic intensity imbalance from overuse of slow CUDA cores, occupancy starvation at small batch sizes, and unnecessary inter-warp shared memory communication from split-K partitioning), it demonstrates that attention can reach 73% of peak in the forward pass and 63% in the backward pass — close enough to GEMM's 80–90% that the remaining gap is now in the "engineering refinement" category, not the "fundamental limitation" category.

This reframes the tradeoff between exact and approximate attention. For years, approximate attention methods (Performer, Reformer, Linformer, Big Bird) were motivated by the $O(N^2)$ cost of exact attention. FlashAttention already reduced the constant factor and memory footprint, making exact attention viable for longer sequences than before. FlashAttention-2 doubles down: it makes exact attention fast enough that the quality-cost tradeoff tilts further toward exact attention for most practical sequence lengths. If FlashAttention made approximate methods less necessary for memory reasons, FlashAttention-2 makes them less necessary for throughput reasons. The paper does not benchmark against approximate methods, so it does not directly settle the question of whether a 2×-faster exact attention kernel eliminates the use case for, say, a 4×-faster approximate kernel that loses 1% accuracy. But the direction is clear: the bar for approximate methods is now higher. An approximate method must not only be faster than exact attention, but faster by enough to justify any accuracy degradation, and the "fast enough" threshold has moved from FlashAttention's 25–40% utilization to FlashAttention-2's 50–73% utilization.

A concrete reconciliation the paper enables: prior work on efficient attention was split between "exact but I/O-aware" (FlashAttention) and "approximate but asymptotically cheaper" (Performer, et al.). The two families addressed different bottlenecks (memory vs. compute complexity) and their relative merits depended on sequence length. FlashAttention-2 blurs this line by showing that the exact approach can be pushed much further in compute efficiency than previously demonstrated. For sequence lengths where FlashAttention-2 runs at 70%+ utilization (roughly 2k–16k on current hardware), an approximate method would need to be >1.4× faster than FlashAttention-2 just to match it in throughput — and then still justify its accuracy loss. This is a higher bar than when FlashAttention ran at 30% utilization and an approximate method only needed to be 2–3× faster than that baseline.

The paper also shifts attention in the systems-for-ML community from "what clever tiling or sparsity pattern can we invent?" toward "how should we map the existing tiling algorithm onto the GPU's parallelism hierarchy?" This is a less glamorous research direction — it involves understanding warp scheduling, shared memory bank conflicts, and register allocation rather than inventing new attention mechanisms — but FlashAttention-2's 2× gain demonstrates that this "boring" engineering can produce speedups comparable to algorithmic innovations. The implication for researchers: before proposing a new approximate attention method, first ensure that your exact attention baseline is as optimized as FlashAttention-2. Many papers comparing approximate methods against "standard attention" were comparing against implementations that leave 60–75% of the GPU idle — a weak baseline that makes approximate methods look better than they are.

Finally, the paper's three-level taxonomy (algorithm-level arithmetic intensity, thread-block-level occupancy, warp-level communication) provides a diagnostic framework that transfers to other GPU kernels. Future work on optimizing other transformer components (MLP layers, layer norm, positional encodings, cross-attention, Mixture-of-Experts routing) can apply the same structured analysis: (1) are you spending too much time on non-Tensor-Core operations? (2) do you have enough thread blocks to occupy all SMs, especially at the batch sizes typical for your workload? (3) within each thread block, are your warps communicating unnecessarily through shared memory when they could be working independently on disjoint output regions? This is not a new paradigm, but it is a useful checklist that the paper implicitly validates by showing that addressing all three yields a 2× cumulative speedup.

Follow-Up Research This Work Enables

Ablation study isolating the three improvements to quantify their individual contributions. The paper presents FlashAttention-2 as a combined system and does not report the speedup from each change independently. A direct follow-up would implement four variants: (a) FlashAttention with only the unscaled accumulator algorithm tweak, (b) FlashAttention with only sequence-length parallelization, (c) FlashAttention with only the split-Q warp partitioning, and (d) the full FlashAttention-2. Benchmarking these on the same A100 configurations used in Figures 4–6 would reveal whether the speedup is dominated by one change (e.g., split-Q accounting for 1.6× of the 2×) or whether all three are necessary (each contributing a roughly multiplicative factor). This matters for implementation prioritization: if the unscaled accumulator alone yields 1.4×, a team could adopt just that change with minimal code modification, while if sequence-length parallelization is the dominant factor, the more invasive kernel launch changes are worth the engineering effort. The experiment would also reveal interactions — e.g., whether split-Q's shared memory savings are partially redundant with the unscaled accumulator's reduction in rescaling operations, or whether the gains are independent and combine multiplicatively.

Numerical error characterization across precisions and sequence lengths. The paper asserts mathematical exactness but provides no floating-point error analysis. A careful follow-up would measure the element-wise error of FlashAttention-2 outputs relative to a double-precision (FP64) reference implementation, for FP16, BF16, and FP32 precisions, across sequence lengths from 512 to 32k. The key question: does the unscaled accumulator (which sums potentially large unnormalized values over many column blocks before dividing by ℓ) cause loss of precision relative to FlashAttention's per-iteration normalization? This would manifest as larger errors at long sequence lengths where the unscaled accumulator accumulates more terms before the final division. The logsumexp substitution (storing L = m + log(ℓ) instead of m and ℓ separately) could also introduce error: log(ℓ) loses precision when ℓ is very large or very small, and the reconstruction P = exp(S - L) compounds any error in L. A follow-up should report max absolute error, mean relative error, and cosine similarity between FlashAttention-2 and the FP64 reference, with particular attention to whether errors accumulate differently for the backward pass (which uses L to reconstruct P, adding reconstruction error to gradient computation). A negative result (e.g., error grows with sequence length for FP16 but stays bounded) would be practically important for users deciding between precisions.

Latency-optimized variant for single-query autoregressive inference. The paper evaluates only throughput (TFLOPs/s) with batch sizes adjusted to keep total tokens constant. For inference latency — particularly single-query autoregressive decoding where one token attends to a growing KV cache — the optimal kernel configuration likely differs. A follow-up would benchmark FlashAttention-2's per-step latency (wall-clock microseconds) for the autoregressive decoding pattern: query length 1, KV cache length growing from 1 to, say, 32k tokens. Key questions: (1) At what KV cache length does sequence-length parallelization (which launches B × H × Tr thread blocks) become beneficial rather than harmful due to kernel launch overhead? (2) Does split-Q warp partitioning help or hurt at query-length-1 where the Q tile is tiny (Br = 1)? (3) Can the kernel dynamically switch between a "latency mode" (single thread block, no sequence-length parallelization) for short sequences and a "throughput mode" for long sequences? (4) How does FlashAttention-2 interact with KV cache storage formats (contiguous vs. paged)? The paper's H100 results hint that even without H100-specific features, throughput is high; a latency-focused follow-up would complete the inference picture. A strong result would be a single kernel that matches or beats FlashAttention's latency for short sequences (where it is currently competitive) while providing the full 2× throughput gain for long sequences.

Interaction with model parallelism strategies for large-scale training. The paper's end-to-end experiments use 8 GPUs but do not specify the parallelism strategy. A systematic follow-up would benchmark FlashAttention-2 under the parallelism configurations used in large-scale training: tensor parallelism (TP) of varying degrees, pipeline parallelism (PP), and sequence parallelism. The key tension: FlashAttention-2's sequence-length parallelization multiplies the available thread blocks by Tr (number of row blocks), but tensor parallelism divides the number of heads H across GPUs, reducing thread blocks from B × H × Tr to B × (H/TP) × Tr. At what TP degree does the per-GPU head count become small enough that occupancy drops again, negating the sequence-length parallelization benefit? For a model with H=32 and sequence length 8k (Tr=63 with Br=128), TP=8 reduces per-GPU heads to 4, giving B × 4 × 63 = 252B thread blocks — still abundant if B≥1. But at H=16 and TP=8, per-GPU heads = 2, yielding 126B thread blocks — marginal for occupying 108 SMs. A strong follow-up would map out the "safe operating region" in (batch_size, seq_length, TP_degree, num_heads) space where FlashAttention-2 maintains high occupancy. Additionally, sequence parallelism (partitioning the sequence dimension across GPUs) directly conflicts with FlashAttention-2's SM-level sequence partitioning — the follow-up should measure whether combining them yields additive benefit or causes contention.

H100-specific optimization exploiting TMA, 4th-gen Tensor Cores, and FP8. The paper's H100 results (Figure 7) explicitly use no H100-specific features, and the authors project "1.5x-2x speedup" from using them. A direct follow-up would implement: (1) TMA (Tensor Memory Accelerator) for asynchronous data movement between HBM and shared memory, which could overlap memory transfers with computation and reduce the latency of loading K and V tiles in the inner loop; (2) 4th-gen Tensor Cores with higher FP16/BF16 throughput; and (3) FP8 support, which halves memory bandwidth requirements and doubles effective compute throughput for matmuls. The concrete benchmark: FlashAttention-2 on H100 with these features enabled, measuring forward+backward throughput on the same configurations as Figure 7, targeting the projected 500–670 TFLOPs/s range. A key design question: does FP8's reduced precision interact badly with the unscaled accumulator (which accumulates large values over many blocks before the final division)? FP8 has very limited dynamic range, and the unscaled accumulator values can grow large, potentially causing overflow. The follow-up should measure numerical error in FP8 and potentially propose a mixed-precision variant (FP8 for matmuls, FP16/BF16 for accumulator and softmax statistics).

Extension to non-NVIDIA GPU architectures and cross-platform portability. The paper reports results only on A100 and H100, with AMD GPUs flagged as future work. A follow-up would port FlashAttention-2 to AMD's CDNA2/CDNA3 architecture (MI250X, MI300X) and Intel's Ponte Vecchio. The key questions: (1) Does the split-Q warp partitioning transfer to AMD's wavefront model (64-thread wavefronts vs. NVIDIA's 32-thread warps), or does the different granularity require rethinking the partitioning? (2) AMD's matrix core throughput and shared memory bandwidth differ from NVIDIA's — does the 16× matmul-to-non-matmul throughput ratio hold, or do the algorithm tweaks need different prioritization? (3) Does the sequence-length parallelization yield similar occupancy benefits, or do architectural differences in thread block scheduling change the optimal Tr multiplier? A negative result (e.g., split-Q doesn't help on AMD because of different shared memory latency characteristics) would refine the understanding of which optimizations are general principles vs. NVIDIA-specific tuning. A strong result would be a single codebase that achieves >60% of theoretical peak on both NVIDIA and AMD GPUs, establishing FlashAttention-2 as a genuinely cross-platform optimization rather than an NVIDIA-specific one.

Practical Applications and Downstream Use Cases

Long-context language model training at reduced cost. This is the paper's primary intended use case, and the numbers in Table 1 make the benefit concrete. Training GPT3-2.7B on 8k context with FlashAttention-2 achieves 225 TFLOPs/s per A100 GPU, compared to 175 TFLOPs/s with FlashAttention — a 1.29× speedup. For a training run that would take 30 days with FlashAttention, FlashAttention-2 reduces it to approximately 23 days, saving 7 days of 8-GPU compute time. At 8k context, the gap to the no-FlashAttention baseline is even more dramatic: 225 vs. 80 TFLOPs/s, a 2.81× speedup — meaning FlashAttention-2 trains the same model in 11 days that would take 30 days without it. For teams training models on 16k or 32k context (where the paper shows FlashAttention-2's throughput stays flat while FlashAttention's declines), the relative benefit grows. The key operational implication: for a fixed training budget, FlashAttention-2 enables either training on 2× longer sequences at the same wall-clock time, or training the same model in roughly 77% of the time (1/1.29) compared to FlashAttention.

High-resolution image and video transformer training. Vision transformers (ViT) and video transformers use attention over patches or frames, where the "sequence length" is the number of patches (e.g., a 1024×1024 image with 16×16 patches has 4096 patches — a sequence length comparable to language models). FlashAttention-2's 2× throughput improvement at sequence lengths 2k–16k directly speeds up training for high-resolution vision models. For video, where attention may be applied across both spatial and temporal dimensions (potentially tens of thousands of tokens), FlashAttention-2's flat throughput at long sequences (Figure 4b: 203 TFLOPs/s at 16k, the same as at 8k) means the quadratic attention cost is being handled efficiently — the GPU stays compute-bound rather than becoming memory-bandwidth-bound. A research team training a video diffusion model with spatio-temporal attention on 16k tokens can expect roughly 2× faster attention computation compared to FlashAttention, which directly reduces training time or enables processing more frames per GPU.

Cost-efficient large-scale inference serving with long contexts. While the paper's benchmarks are throughput-oriented, the 2× attention speedup also benefits inference serving where multiple requests are batched together. A serving system handling long-context queries (document summarization, codebase understanding, multi-turn conversations with long history) processes attention as part of each forward pass. With FlashAttention-2 achieving 73% of peak in the forward pass (Figure 5), the attention portion of each inference step runs substantially faster than with FlashAttention. For a deployment serving GPT-style models at 8k context with batch processing, the reduced attention latency per batch means higher throughput (more requests per second per GPU) and lower per-request cost. The H100 results (up to 338 TFLOPs/s without H100-specific features; Figure 7) suggest even larger gains on next-generation hardware, making long-context inference economically viable at larger scale. The 1.3–1.5× forward pass speedup over the Triton implementation (Figure 5) means teams using Triton-based serving frameworks can get an immediate throughput improvement by switching to FlashAttention-2's CUDA kernels, without changing their model architecture or serving infrastructure.

---

When to Prefer This Method

The paper positions FlashAttention-2 as a direct replacement for FlashAttention, not as one option among alternatives. It computes the same mathematical function (exact attention) with approximately 2× higher throughput on A100 GPUs, with no accuracy tradeoff. The decision rule is therefore whether the practitioner's hardware, precision, and parallelism configuration benefit from the specific optimizations:

• Prefer FlashAttention-2 when training or running inference on NVIDIA A100 or H100 GPUs with sequence lengths ≥ 512 and FP16/BF16 precision. The throughput improvement over FlashAttention is robust across all configurations tested (Figures 4–6), ranging from 1.45× to 3.11× depending on head dimension, causal masking, and sequence length. The benefit is largest with head dimension 128, without causal masking, and at sequence lengths 2k–16k.

• Prefer FlashAttention-2 particularly when training at long sequences (≥ 2k) with small batch sizes. The sequence-length parallelization (Section 3.2) specifically addresses the low-occupancy problem that occurs when batch_size × num_heads is small — exactly the regime for long-context training. The end-to-end training results (Table 1) show a 1.29× speedup at 8k context vs. only 1.04× at 2k context, confirming that the benefits are weighted toward the long-sequence regime.

• The paper does not propose a tradeoff against approximate attention methods, so no decision rule can be extracted. FlashAttention-2 computes exact attention and should match standard attention's outputs bit-for-bit (modulo floating-point non-associativity). The paper makes no claims about when one should prefer exact vs. approximate attention, and does not benchmark against approximate methods. A practitioner choosing between FlashAttention-2 and, say, a linear-complexity approximation like Performer faces a quality-vs-speed tradeoff that this paper does not inform — FlashAttention-2 is faster than the exact baselines, but an approximate method may still be faster at very long sequences (32k+) at some accuracy cost. That tradeoff must be evaluated with application-specific accuracy benchmarking that the paper does not provide.

• The paper also does not explicitly position FlashAttention-2 against FlashAttention in a way that suggests conditional preference. Since FlashAttention-2 strictly dominates FlashAttention in throughput on the tested hardware without changing the output, there is no scenario described where one would prefer the original FlashAttention. The only implicit exception is that FlashAttention-2's sequence-length parallelization might increase kernel launch overhead for very short sequences (N < 512), a regime the paper does not benchmark. If latency at very short sequences is critical, the original FlashAttention (or a latency-optimized variant of FlashAttention-2) might be preferable, but the paper provides no data to confirm or refute this.