READ MORE

ArXiv: 2205.14135

🎯 Pitch

Standard exact attention implementations are memory-bound, not compute-bound: they wastefully read and write the full N×N attention matrix to slow HBM. FlashAttention makes exact attention 7.6× faster by avoiding that materialization entirely via tiling, yet still computes a mathematically identical result. This I/O-awareness unlocks training transformers on 64K-length sequences—a regime where standard attention exhausts GPU memory.

---

1. Executive Summary

This paper proposes FlashAttention, an IO-aware exact attention algorithm that reduces the number of memory reads and writes between GPU high bandwidth memory (HBM) and on-chip SRAM by using tiling (splitting the input into blocks and incrementally computing the softmax reduction) and recomputation (storing softmax normalization statistics from the forward pass to quickly recompute attention on-chip during the backward pass, rather than reading the large intermediate attention matrix from HBM). Evaluated on Transformer training across BERT-large, GPT-2, and the Long-Range Arena benchmark, FlashAttention yields up to 3× speedup on GPT-2 training, 15% faster BERT-large training than the MLPerf 1.1 record, and 2.4× speedup on long-range arena tasks, while enabling the first Transformer to achieve better-than-chance performance on the Path-X challenge (61.4% accuracy, sequence length 16K) and—with its block-sparse extension—on Path-256 (63.1% accuracy, sequence length 64K), establishing that modeling longer context with attention becomes practical only when the attention computation itself is made IO-aware rather than merely reducing FLOPs.

2. Context and Motivation

The Core Problem: Attention Is Bottlenecked by Memory Bandwidth, Not Just FLOPs

The fundamental problem FlashAttention addresses is deceptively specific: the standard implementation of self-attention is bottlenecked by memory reads and writes, not by arithmetic computation. This is a problem of hardware utilization, not algorithm design. The paper's core insight is that making attention wall-clock fast requires understanding the GPU memory hierarchy and designing algorithms that minimize data movement between slow and fast memory — a principle the authors term IO-awareness.

This matters because self-attention is the computational centerpiece of Transformer models, which dominate modern NLP and computer vision. The self-attention operation has time and memory complexity quadratic in sequence length ($O(N^2)$), making it the primary bottleneck when processing long sequences. But the paper argues — and provides substantial evidence — that the real bottleneck is not the $N^2$ FLOPs per se, but rather the $N^2$ memory accesses that the standard implementation performs to the relatively slow GPU high bandwidth memory (HBM). As the authors state in Section 1:

"Many approximate attention methods have aimed to reduce the compute and memory requirements of attention… [but] many of them do not display wall-clock speedup against standard attention and have not gained wide adoption. One main reason is that they focus on FLOP reduction (which may not correlate with wall-clock speed) and tend to ignore overheads from memory access (IO)."

This is a critical distinction. Reducing FLOPs (e.g., from $O(N^2)$ to $O(N)$ via approximation) does not guarantee faster execution because modern GPUs are increasingly constrained by memory bandwidth, not compute throughput. The paper cites the hardware trend: "compute speed has out-paced memory speed" (Section 1), meaning that many deep learning operations are memory-bound — their runtime is determined by how long it takes to move data between HBM and the compute units, not by how many arithmetic operations are performed.

The Importance: Enabling Long Context in Transformers

The paper motivates this problem along three axes of importance:

Scientific understanding of hardware-software co-design. The paper identifies a missing principle in deep learning systems: IO-aware algorithm design. While IO-aware algorithms have been critical in other fields — database joins, image processing, numerical linear algebra — they had not been systematically applied to attention. The paper explicitly draws this connection:

"IO-aware algorithms have been critical for similar memory-bound operations, when reading and writing data can account for a large portion of the runtime—such as database joins, image processing, numerical linear algebra, and more. However, common Python interfaces to deep learning such as PyTorch and Tensorflow do not allow fine-grained control of memory access."

This framing suggests that the problem is not unique to attention — it's a general deficiency in how deep learning frameworks interact with hardware, and attention happens to be the most pressing instance.

Practical deployment of long-context Transformers. The natural language processing community has long recognized that longer context improves model quality, but has been constrained by the $O(N^2)$ memory requirements of attention. As the paper notes in Section 1, "equipping [Transformers] with longer context remains difficult." This difficulty manifests concretely: on an A100 GPU with 40GB HBM, standard attention implementations run out of memory at sequence lengths that are modest by document-processing standards. The paper's experiments demonstrate that FlashAttention's linear memory footprint enables training on sequences up to 64K tokens that would be impossible for standard attention — unlocking entirely new capabilities like the Path-X and Path-256 benchmarks that had defeated all prior Transformer models.

Economic and research accessibility. Training large Transformers on long sequences is expensive. If smarter attention algorithms can deliver 2-4× speedups without sacrificing model quality, that translates directly to reduced cloud compute costs, faster research iteration cycles, and democratized access to long-context modeling for groups with limited GPU budgets. The paper does not frame this as an explicitly economic argument, but the extensive wall-clock time benchmarks in Section 4 make it clear.

Prior Approaches and Where They Fall Short

The paper situates itself relative to two broad classes of prior work:

1. Approximate attention methods that reduce FLOPs.

A large body of work has attempted to linearize or otherwise simplify the $O(N^2)$ attention computation. The paper surveys these in Appendix A, covering sparse approximations (Reformer, Sparse Transformers, Longformer, BigBird), low-rank approximations (Linformer, Performer), kernel-based methods (Linear Attention), and combinations of both (Scatterbrain, Long-short transformer, Combiner). These methods reduce the asymptotic complexity — typically to $O(N)$ or $O(N \log N)$ — by trading off model expressivity: they do not compute exact attention but rather an approximation that is cheaper in FLOPs.

The critical shortcoming is that FLOP reduction often fails to translate to wall-clock speedup. The paper makes this empirical claim explicitly in Section 1:

"Although these methods reduce the compute requirements to linear or near-linear in sequence length, many of them do not display wall-clock speedup against standard attention and have not gained wide adoption."

Why does this happen? The paper argues it's because these methods ignore IO overheads. Even if an algorithm requires fewer arithmetic operations, it may still require the same (or more) data movement between memory levels. An approximate attention method that performs $O(N)$ FLOPs but still reads and writes $O(N^2)$ intermediate data will be bottlenecked by memory bandwidth — and thus no faster in practice than exact attention with $O(N^2)$ FLOPs but the same number of memory accesses. The paper's benchmarking results in Figure 3 (left) and the extensive Appendix E tables (Tables 9–20) validate this: many approximate attention methods run slower than FlashAttention even though they report lower FLOP counts, because their memory access patterns are suboptimal.

A deeper subtlety: even methods that do achieve asymptotic memory savings may not outperform FlashAttention at moderate sequence lengths because of constant-factor overheads. The paper observes (Section 4.3) that FlashAttention is faster than approximate methods for sequences up to 512–1024 tokens — which covers most practical NLP workloads — because it avoids the overhead of sparse indexing, hashing, or other approximation machinery. The crossover point where approximate methods become faster occurs only at longer sequences, and even then, block-sparse FlashAttention remains faster than all of them.

2. Standard attention implementations that materialize the full $N \times N$ intermediate matrix.

The standard PyTorch attention implementation — and even highly optimized variants like Megatron-LM — computes attention as a sequence of separate operations (matrix multiply → softmax → dropout → matrix multiply), each of which writes its output to HBM and reads it back for the next operation. This is Algorithm 0 in Section 2.2. Critically, the intermediate $N \times N$ matrix S (pre-softmax scores) and P (post-softmax probabilities) are both materialized in HBM, requiring $O(N^2)$ memory and $O(N^2)$ reads/writes. For a typical GPT-2 configuration ($N = 1024$, $d = 64$, 16 heads, batch size 64), the paper's Figure 2 (left) shows that standard attention performs 40.3 GB of HBM reads/writes for the forward + backward pass — far more than the actual model parameters or the input/output tensors.

The paper identifies specific failure modes in existing attempts to optimize standard attention:

• Kernel fusion alone is insufficient. Frameworks like Apex FMHA (Nvidia's fused multi-head attention) fuse the elementwise operations (masking, softmax, dropout) into a single CUDA kernel, avoiding intermediate HBM writes between these operations. This is a form of IO-awareness, but an incomplete one. As the paper notes in Section 2.1: "in the context of model training, the intermediate values still need to be written to HBM to save for the backward pass, reducing the effectiveness of naive kernel fusion." FMHA stores the softmax attention matrix P to HBM for the backward pass, so it still consumes $O(N^2)$ memory and performs $O(N^2)$ HBM writes during the forward pass. Table 7 (Appendix E.4) shows that FlashAttention is slightly faster than FMHA even at short sequences (128–512 tokens) while supporting much longer sequences (FMHA is limited to 512 tokens).

• Gradient checkpointing trades speed for memory. The memory-efficient forward pass described by Rabe and Staats (2021) avoids storing the $N \times N$ attention matrix by recomputing it during the backward pass. However, as the paper argues in Appendix B.5, this approach "is around the same speed or slightly slower than standard attention" because the recomputation increases FLOPs without reducing HBM accesses — it still reads Q and K from HBM repeatedly to rebuild the attention matrix block by block. FlashAttention's key insight is that by also reducing HBM accesses through tiling, the recomputation can actually be faster than reading the stored matrix.

• Memory-efficient methods that still access HBM quadratically. Even the "memory-efficient" forward pass derivation in Appendix B.1 — which shows that attention can be computed with $O(N)$ extra memory by accumulating the softmax normalization constant incrementally — still performs $O(N^2)$ HBM accesses, because each output element $o_i$ requires summing over all $N$ key-value pairs. Reducing memory capacity does not automatically reduce memory bandwidth. The paper's contribution is recognizing this gap and showing that tiling addresses both simultaneously.

How This Paper Positions Itself

The paper positions FlashAttention as filling a specific, principled gap in the attention literature: IO-awareness. The authors frame this not as yet another approximate attention method, but as a fundamental algorithmic principle that applies to exact attention and can be extended to approximate variants as well.

Not an approximation, but an exact algorithm that respects hardware. The paper emphasizes repeatedly (title, abstract, Section 1, Section 3) that FlashAttention computes exact attention — the same mathematical function as standard attention, with the same numerical properties. This is a crucial differentiator from the approximate attention literature. FlashAttention does not change the model's computation; it changes how that computation executes on hardware. This means:

• No accuracy trade-off: FlashAttention produces identical results to standard attention (validated by the overlapping training perplexity curves in Appendix E, Figure 4).
• No new hyperparameters or architecture changes: It's a drop-in replacement for existing attention implementations.
• The speedup is pure engineering efficiency: every model that uses attention can benefit without changing its definition.

IO complexity analysis as a principled framework. The paper provides formal IO complexity bounds (Theorem 2, Proposition 3) that analyze FlashAttention's HBM access count as $O(N^2 d^2 M^{-1})$, where $M$ is SRAM size, compared to $\Omega(Nd + N^2)$ for standard attention. For typical values ($d = 64$–$128$, $M \approx 100$KB), this represents a 5–9× reduction in memory accesses. The paper also proves a lower bound (Proposition 3) showing that no exact attention algorithm can asymptotically improve on this for all SRAM sizes — FlashAttention is IO-optimal within constants for a certain parameter regime. This theoretical grounding elevates the work beyond an engineering contribution: it provides a reusable analysis framework for evaluating attention algorithms based on their data movement, not their FLOP count.

A unifying primitive for both exact and approximate attention. FlashAttention is not positioned as an alternative to approximate attention, but as a substrate that can make approximate attention more practical. The block-sparse FlashAttention extension (Section 3.3) demonstrates this: by applying the same IO-aware tiling to a sparse attention pattern, the algorithm achieves speedups proportional to the sparsity ratio $s$ ($O(N^2 d^2 M^{-1} s)$ HBM accesses), making it faster than all existing approximate methods (Figure 3, Tables 9–20). This suggests that the primary barrier to adoption of approximate attention methods has been their memory access overhead, not their FLOP count — when implemented in an IO-aware fashion, their theoretical efficiency gains become realizable.

Hardware lottery and the need for low-level implementation. The paper acknowledges a tension in its approach: achieving IO-awareness requires writing attention in CUDA rather than a high-level framework like PyTorch. Section 5 explicitly discusses this limitation:

"Our current approach to building IO-aware implementations of attention requires writing a new CUDA kernel for each new attention implementation. This requires writing the attention algorithm in a considerably lower-level language than PyTorch, and requires significant engineering effort."

This is framed as a call to action for the systems community: developing compilers that can translate attention algorithms expressed in high-level languages into IO-aware GPU implementations, analogous to Halide in image processing. The paper positions its CUDA implementation as existence proof that IO-awareness is both possible and valuable, while acknowledging that broader adoption requires better tooling.

Contrast with alternative sequence modeling approaches. In Appendix A, the paper briefly acknowledges work on replacing attention entirely with state-space models (S4, HiPPO), LambdaNetworks, or attention-free transformers. However, it positions these as orthogonal — FlashAttention addresses the hardware efficiency of attention itself, making it competitive with or superior to these alternatives on existing benchmarks, while preserving the attention mechanism's flexibility and the vast ecosystem of pretrained Transformer models that depend on it.

3. Technical Approach

3.1 Reader Orientation

FlashAttention is a GPU kernel — a custom CUDA program that executes the entire self-attention computation (matrix multiplies, softmax, masking, dropout) in a single fused operation without ever writing the large $N \times N$ intermediate attention matrix to GPU main memory. It solves the problem of attention being bottlenecked by memory bandwidth rather than computation speed by reorganizing the algorithm to keep data in fast on-chip SRAM for as long as possible, using tiling (processing the input in blocks that fit in SRAM) and recomputation (recalculating the attention matrix during backpropagation rather than storing it from the forward pass).

3.2 Big-Picture Architecture (Diagram in Words)

The system operates within a single GPU kernel and has four logical tiers:

1. HBM (GPU High Bandwidth Memory) — the large but relatively slow memory pool (40–80 GB at 1.5–2.0 TB/s on A100) that holds the input matrices Q, K, V, the accumulated output O, and the running statistics $\ell$ and $m$.
2. On-chip SRAM — a small but fast memory (192 KB per streaming multiprocessor on A100, estimated ~19 TB/s) that temporarily holds blocks of Q, K, V and the intermediate block-level attention scores S_ij.
3. Compute units — perform the matrix multiplications, softmax, masking, and dropout operations entirely on data resident in SRAM.
4. The loop nest — an outer loop over blocks of K and V (columns of the attention matrix) and an inner loop over blocks of Q (rows of the attention matrix), which incrementally accumulates the softmax-normalized output O using running statistics.

Information flows as follows: blocks of K and V are loaded from HBM into SRAM (outer loop) → for each K-V block, blocks of Q and the corresponding running statistics are loaded into SRAM (inner loop) → the attention sub-matrix S_ij = Q_i K_j^T is computed entirely on-chip → softmax is computed for this sub-block and combined with the running statistics to update the partial output O_i incrementally → the updated O_i and statistics are written back to HBM → the algorithm moves to the next Q block, then the next K-V block, until all blocks have been processed and the final exact attention output O is resident in HBM.

3.3 Roadmap for the Deep Dive

• First, the GPU performance model (compute-bound vs. memory-bound operations and kernel fusion), because understanding what makes attention slow requires understanding GPU hardware.
• Second, the standard attention implementation (Algorithm 0) and why it is memory-bound, which establishes the baseline that FlashAttention improves upon.
• Third, the memory-efficient forward pass derivation (Appendix B.1), which shows how attention can be computed with $O(N)$ memory but still has quadratic HBM accesses — setting up the gap FlashAttention fills.
• Fourth, the memory-efficient backward pass derivation (Appendix B.2), which shows how gradients can also be computed in $O(N)$ memory with quadratic HBM accesses.
• Fifth, the full FlashAttention algorithm (forward pass in Algorithm 1, backward pass in Algorithm 4), which combines tiling and recomputation to reduce HBM accesses to sub-quadratic — this is the core contribution.
• Sixth, the IO complexity analysis (Theorem 2, Proposition 3), which formalizes why and how much FlashAttention reduces memory accesses, and proves near-optimality.
• Seventh, the block-sparse FlashAttention extension (Algorithm 5, Proposition 4), showing how the IO-aware framework yields proportional speedups from sparsity.

3.4 Detailed, Sentence-Based Technical Breakdown

This is primarily a systems paper whose core idea is that attention can be made wall-clock fast by redesigning the algorithm to minimize data movement between GPU memory tiers, using tiling and recomputation to keep operations on-chip.

---

3.4.1 The GPU Performance Model: Why Attention Is Memory-Bound

Before explaining how FlashAttention works, the paper establishes the hardware model that motivates the entire approach (Section 2.1). This model is not FlashAttention's invention, but FlashAttention's design is unintelligible without it.

The GPU memory hierarchy. Modern GPUs have two primary memory tiers relevant to computation:

• HBM (High Bandwidth Memory): This is the GPU's main memory, analogous to CPU RAM. On the A100 GPU used in the paper's experiments, HBM is 40–80 GB in capacity with 1.5–2.0 TB/s bandwidth. All model parameters, inputs, outputs, and intermediate tensors are stored here between kernel launches. Accessing HBM is the slowest step in any GPU computation — every read or write costs tens to hundreds of clock cycles.
• On-chip SRAM: This is a much smaller but dramatically faster memory located physically on the GPU die, shared by the threads within each streaming multiprocessor (SM). On the A100, each of 108 SMs has 192 KB of SRAM, with bandwidth "estimated around 19 TB/s." The paper emphasizes that SRAM is "an order of magnitude faster than HBM but many orders of magnitude smaller in size."

The paper also notes: "As compute has gotten faster relative to memory speed, operations are increasingly bottlenecked by memory (HBM) accesses. Thus exploiting fast SRAM becomes more important."

The execution model. A GPU kernel is a function executed by thousands of parallel threads. The typical lifecycle is: "Each kernel loads inputs from HBM to registers and SRAM, computes, then writes outputs to HBM." This load-compute-write cycle applies to every intermediate operation in standard attention — and each trip to HBM costs time.

Compute-bound vs. memory-bound operations. The paper introduces the concept of arithmetic intensity to classify GPU operations:

"This is commonly measured by the arithmetic intensity, which is the number of arithmetic operations per byte of memory access."

An operation is compute-bound if its runtime is dominated by arithmetic operations — the memory accesses can be overlapped or hidden behind computation. Examples include matrix multiplications with large inner dimensions and convolutions with many channels.

An operation is memory-bound if its runtime is dominated by memory accesses — the compute units spend most of their time waiting for data to arrive from HBM. Examples include "most other operations: elementwise (e.g., activation, dropout), and reduction (e.g., sum, softmax, batch norm, layer norm)."

The paper's key claim is that softmax and dropout are memory-bound, not compute-bound. The softmax of an $N \times N$ matrix requires reading $N^2$ elements from HBM (the pre-softmax scores S), computing exponentiation and division (which is fast), and writing $N^2$ elements back to HBM (the post-softmax probabilities P). The arithmetic is trivial compared to the data movement. Similarly, dropout on the attention probabilities requires reading $N^2$ elements, generating random numbers, scaling, and writing $N^2$ elements back. The fraction of time spent actually computing is tiny — the GPU spends most of its time just moving data.

Kernel fusion and its limits. The standard technique for accelerating memory-bound operations is kernel fusion: "if there are multiple operations applied to the same input, the input can be loaded once from HBM, instead of multiple times for each operation." For example, Nvidia's Apex FMHA (fused multi-head attention) fuses the mask, softmax, dropout, and second matrix multiply into one kernel — the intermediate matrices S and P never get written to HBM separately.

However, the paper identifies a critical limitation: "in the context of model training, the intermediate values still need to be written to HBM to save for the backward pass, reducing the effectiveness of naive kernel fusion." Even with fusion, the forward pass must store the attention probability matrix P in HBM so that the backward pass can use it to compute gradients — otherwise the intermediate values are lost when the kernel finishes and SRAM is freed. This is why FMHA still requires $O(N^2)$ HBM writes during the forward pass (it stores P), and why it is limited to short sequences (at most 512 tokens — it runs out of HBM for longer sequences because it tries to store the $N \times N$ matrix for every head and every batch element).

FlashAttention overcomes this limitation by recomputing P during the backward pass from stored normalization statistics, rather than storing P itself. The forward pass writes only a small amount of data ($m$ and $\ell$, each of size $N$) to HBM; the backward pass reconstructs blocks of P on-chip from these statistics and the original Q, K, V inputs.

---

3.4.2 Standard Attention (Algorithm 0): The Baseline

The paper formalizes the standard attention implementation as Algorithm 0 (Section 2.2). This algorithm represents how PyTorch, HuggingFace Transformers, and Megatron-LM implement attention — as a sequence of three distinct operations, each writing its output to HBM:

Given input matrices Q, K, V, all of shape $N \times d$ and residing in HBM:

Step 1 — Matrix multiply QK^T: Load Q and K from HBM in blocks, compute $S = QK^T \in \mathbb{R}^{N \times N}$, and write S to HBM. This requires reading $2Nd$ elements (Q and K) and writing $N^2$ elements (S).

Step 2 — Softmax: Read S from HBM, compute $P = \text{softmax}(S) \in \mathbb{R}^{N \times N}$ (applied row-wise), and write P to HBM. This requires reading $N^2$ elements and writing $N^2$ elements.

Step 3 — Matrix multiply PV: Load P and V from HBM, compute $O = PV \in \mathbb{R}^{N \times d}$, and write O to HBM. This requires reading $N^2 + Nd$ elements and writing $Nd$ elements.

The total HBM accesses are:

$\Theta(Nd + N^2)$

where $N$ is the sequence length and $d$ is the head dimension.

What this computes: The standard exact attention output — the same mathematical function as FlashAttention. Q, K, and V are read from HBM, multiplied into the attention score matrix S, transformed into probabilities P via row-wise softmax, and used to compute a weighted sum of value vectors, producing output O.

Why this form is problematic: The $N^2$ terms dominate for any realistic sequence length (for GPT-2 with $N=1024$, $d=64$, $N^2 = 1,048,576$ vs. $Nd = 65,536$). Both the score matrix S and the probability matrix P must be fully materialized in HBM — each consuming $O(N^2)$ memory and each requiring $O(N^2)$ reads and writes. The paper quantifies this concretely in Figure 2 (left): for a GPT-2 medium configuration with sequence length 1024, head dimension 64, 16 heads, and batch size 64, the forward + backward pass of standard attention performs 40.3 GB of HBM reads/writes and requires 66.6 GFLOPs of computation. The arithmetic intensity (66.6 GFLOPs / 40.3 GB ≈ 1.65 FLOPs per byte) is extremely low — well into the memory-bound regime on modern GPUs.

The paper notes that standard attention "often $N \gg d$ (e.g., for GPT2, $N=1024$ and $d=64$)." When $N$ is 16× larger than $d$, the $N^2$ term is 16× larger than the $Nd$ term. As sequence length grows, this gap widens quadratically.

Why this form is standard despite the problem: The algorithm directly mirrors the mathematical definition of attention and is trivial to implement in frameworks like PyTorch. Each operation (matmul, softmax) is a single library call. The framework handles memory management — the user never sees where S and P are stored. This abstraction hides the performance pathology: the user sees a few lines of Python but does not see the 40 GB of HBM traffic they cause. The paper's contribution is recognizing this hidden cost and restructuring the computation to avoid it while remaining mathematically exact.

---

3.4.3 The Memory-Efficient Forward Pass (Appendix B.1): Tiling with Statistics

The paper next derives a memory-efficient forward pass that computes exact attention with $O(N)$ extra memory instead of $O(N^2)$. This derivation is not FlashAttention itself, but it establishes the conceptual foundation: the softmax normalization constant can be decoupled from the key-value columns, enabling incremental computation.

The key insight is that the softmax normalization for each query $i$ can be separated from the weighted sum over values:

Define the softmax normalization constant for query row $i$:

$L_i = \sum_j e^{q_i^T k_j}$

where $q_i$ is the $i$-th query vector (row of Q), $k_j$ is the $j$-th key vector (row of K), and the sum ranges over all $j = 1, \ldots, N$ key positions.

What this computes: For each query position $i$, $L_i$ is the sum of exponentiated dot products between $q_i$ and all key vectors. This is the denominator of the row-wise softmax — the normalizing constant that makes the attention probabilities sum to 1 for each query.

The attention output for query $i$ can then be written as:

$o_i = \sum_j P_{ij} v_j = \sum_j \frac{e^{q_i^T k_j}}{L_i} v_j$

where $P_{ij} = e^{q_i^T k_j} / L_i$ is the softmax probability for query $i$ attending to key $j$, and $v_j$ is the $j$-th value vector.

What this computes: The $i$-th output row $o_i$, which is a weighted combination of all value vectors, with weights given by the softmax probabilities. Crucially, this form factors the computation into two phases: first compute $L_i$ for all $i$ (which costs $O(N)$ memory — one scalar per query), then compute $o_i$ for all $i$ by accumulating the weighted sum (which costs $O(d)$ extra memory per query for the running sum).

Why this form matters: In standard attention, both $S_{ij} = q_i^T k_j$ and $P_{ij} = e^{S_{ij}}/L_i$ are $N \times N$ matrices that must be fully materialized before computing O. This derivation shows that if we first compute $L_i$ for each query (reading all keys once, costing $O(Nd)$ reads and $O(N)$ writes for the $L_i$ scalars), we can then compute $o_i$ incrementally: for each key-value pair $(k_j, v_j)$, update the running sum $o_i \leftarrow o_i + e^{q_i^T k_j} v_j$, and at the end divide by $L_i$. This completely avoids storing the $N \times N$ matrix P.

The algorithmic procedure is:

1. Compute $L_i$ for all $i$: For each query $i$, loop over all keys $j$, compute $e^{q_i^T k_j}$, accumulate into $L_i$. This requires $O(N)$ extra memory (one scalar $L_i$ per query position).
2. Compute $o_i$ for all $i$: For each query $i$, loop over all key-value pairs $(k_j, v_j)$, compute the term $\frac{e^{q_i^T k_j}}{L_i} v_j$, and accumulate into $o_i$. This requires $O(d)$ extra memory (one vector $o_i$ per query position).

The total extra memory beyond the input and output is $O(N)$ (for the $L_i$ scalars), compared to $O(N^2)$ for standard attention.

Why this doesn't solve the speed problem: The paper explicitly notes this limitation — "naively they still incur quadratic HBM accesses, resulting in slower execution speed" (Appendix B.1). Even though the memory capacity is reduced from $O(N^2)$ to $O(N)$, the memory bandwidth (number of reads/writes to HBM) remains $O(N^2)$. In the first phase, for each of the $N$ queries, we loop over all $N$ keys, reading K from HBM $N$ times (once per query), for a total of $O(N^2 d)$ HBM reads. Similarly, in the second phase, for each of the $N$ queries, we loop over all $N$ value vectors. The total HBM accesses remain $\Theta(N^2 d + N^2 + Nd) = \Theta(N^2)$ — exactly the same asymptotic as standard attention, albeit with a smaller constant because we don't write S and P. The algorithm is memory-efficient in the sense of peak allocation but still bottlenecked on HBM bandwidth.

This establishes the gap that FlashAttention fills: tiling can reduce HBM accesses, not just peak memory.

---

3.4.4 The Memory-Efficient Backward Pass (Appendix B.2): Gradients Without Storing P

The backward pass computes the gradients of the loss function with respect to Q, K, and V, given the output gradient dO (the gradient of the loss with respect to the attention output). Standard attention uses the stored attention probabilities P from the forward pass to compute these gradients efficiently. The challenge is to compute dQ, dK, and dV without storing P.

The paper derives the backward pass analytically (Appendix B.2) and shows that it can also be computed with $O(N)$ extra memory. The derivation uses standard chain rule calculus through the attention computation:

Given output gradient $\text{dO} \in \mathbb{R}^{N \times d}$ (where $\text{dO}_{ij} = \partial\phi/\partial\text{O}_{ij}$ for some scalar loss $\phi$), we want $\text{dQ}, \text{dK}, \text{dV} \in \mathbb{R}^{N \times d}$.

Gradient for V. The simplest case:

$\text{dV} = P^T \text{dO}$

and equivalently for each value vector $j$:

$dv_j = \sum_{i} P_{ij} do_i = \sum_{i} \frac{e^{q_i^T k_j}}{L_i} do_i$

What this computes: The gradient for each value vector $v_j$ is a weighted sum of the output gradients from every query position, weighted by how much query $i$ attended to key $j$. Since we already computed $L_i$ in the forward pass, we can compute $dv_j$ incrementally one query at a time, accumulating into $dv_j$ — requiring only $O(d)$ extra memory per key position.

Gradient for Q and K. The derivation is more involved. The paper works through the chain:

First, from $O = PV$, we have $\text{dP} = \text{dO} \cdot V^T$, so:

$dP_{ij} = do_i^T v_j$

What this computes: The gradient with respect to the attention probability $P_{ij}$ is the inner product between the gradient arriving at output position $i$ and the value vector at position $j$.

Next, the softmax Jacobian. For $P_{i:} = \text{softmax}(S_{i:})$ (one row of the attention scores), the gradient is:

$dS_{i:} = (\text{diag}(P_{i:}) - P_{i:}P_{i:}^T) dP_{i:} = P_{i:} \odot dP_{i:} - (P_{i:}^T dP_{i:}) P_{i:}$

where $\odot$ denotes pointwise (Hadamard) multiplication and $P_{i:}^T dP_{i:}$ is a scalar (the dot product of row $i$ of P with row $i$ of dP).

What this computes: The gradient through the softmax function — how changes in the raw attention scores $S_{ij}$ affect the softmax probabilities $P_{ij}$. The key simplification is that the scalar $D_i = P_{i:}^T dP_{i:}$ can be rewritten without accessing the full vectors $P_{i:}$ and $dP_{i:}$.

Define:

$D_i = P_{i:}^T dP_{i:} = \sum_j \frac{e^{q_i^T k_j}}{L_i} do_i^T v_j = do_i^T \sum_j \frac{e^{q_i^T k_j}}{L_i} v_j = do_i^T o_i$

What this computes: $D_i$ is the dot product between the output gradient row $do_i$ and the output row $o_i$ itself (computed in the forward pass and stored). This is the critical simplification — instead of summing over $N$ positions with $P_{ij}$ and $dP_{ij}$ (which would require materializing those $N$-dimensional vectors), we only need one dot product between two $d$-dimensional vectors ($do_i$ and $o_i$). Since $d \ll N$, this is much cheaper.

Why this form eliminates the need for P: The standard backward pass (Algorithm 3) requires reading the full $N \times N$ matrix P from HBM to compute $D_i$ (by summing $P_{ij} \cdot dP_{ij}$ over $j$). This derivation shows that $D_i$ can instead be computed from $do_i^T o_i$, using only the already-stored output O and the incoming gradient dO — both of which are $N \times d$ and already required anyway. Therefore the $N \times N$ matrix P never needs to be read.

With $D_i$ computed, the gradient per element is:

$dS_{ij} = P_{ij}(dP_{ij} - D_i)$

And the final gradients with respect to Q and K are:

$dq_i = \sum_j dS_{ij} k_j = \sum_j \frac{e^{q_i^T k_j}}{L_i} (do_i^T v_j - D_i) k_j$

$dk_j = \sum_i dS_{ij} q_i = \sum_i \frac{e^{q_i^T k_j}}{L_i} (do_i^T v_j - D_i) q_i$

What these compute: $dq_i$ (gradient for query $i$) is computed by iterating over all keys $j$, weighting $k_j$ by the local gradient signal $dS_{ij}$. $dk_j$ (gradient for key $j$) is the analogous sum over all queries. Both can be computed incrementally using only the stored $L_i$ and $o_i$ from the forward pass, plus access to the original Q, K, V.

The procedure for the backward pass with $O(N)$ extra memory is:

1. Compute $dv_j$ for all $j$ using Eq. (3) — requires $O(d)$ extra memory.
2. Compute $D_i$ for all $i$ using $D_i = do_i^T o_i$ — requires $O(N)$ extra memory (one scalar per query).
3. Compute $dq_i$ for all $i$ using Eq. (5) — requires $O(d)$ extra memory.
4. Compute $dk_j$ for all $j$ using Eq. (6) — requires $O(d)$ extra memory.

The total extra memory is $O(N)$, but again, the HBM accesses remain $O(N^2)$ for the same reason as the forward pass: we loop over all key-value pairs for each query.

Why this derivation matters for FlashAttention: It proves that the backward pass can be computed without ever storing or reading the $N \times N$ attention matrices S or P. All that is needed from the forward pass is O (the output, size $N \times d$) and the softmax normalization statistics $m$ and $\ell$ (each size $N$). FlashAttention stores exactly these (plus the random number generator state R for dropout) during the forward pass. During the backward pass, it recomputes the attention probabilities $P_{ij}$ on-chip from Q, K, $m$, and $\ell$, avoiding all $O(N^2)$ HBM accesses for reading P. The recomputation increases FLOPs (because $P_{ij}$ is computed twice — once in forward, once in backward), but the savings in HBM bandwidth more than compensate.

---

3.4.5 FlashAttention Forward Pass (Algorithm 1): Tiling with Softmax Scaling

This is the core algorithm. FlashAttention combines the memory-efficient computation (Appendix B.1) with tiling — processing the input in blocks that fit in on-chip SRAM — to simultaneously reduce HBM accesses and peak memory.

The problem: softmax couples columns of K. The key challenge is that softmax operates row-wise over the full $N \times N$ score matrix S. For a single query $i$, the softmax normalization requires access to the dot products $q_i^T k_j$ for all $j$ simultaneously — you cannot compute $\text{softmax}([x_1, x_2, x_3])$ correctly from individual blocks $[x_1]$ and $[x_2, x_3]$ without additional information. This coupling means that naively splitting K into blocks would produce incorrect softmax outputs for each block if the normalization isn't adjusted.

The solution: incremental softmax using running maximum and sum. The paper uses a well-established technique (referenced as "algebraic aggregation" from Gray et al., 1997, and previously used in attention by Kitaev et al., 2020; Milakov and Gimelshein, 2018; Rabe and Staats, 2021) to decompose softmax over blocks.

For numerical stability, softmax is typically computed as:

$m(x) = \max_i x_i, \quad f(x) = [e^{x_1 - m(x)}, \ldots, e^{x_B - m(x)}], \quad \ell(x) = \sum_i f(x)_i, \quad \text{softmax}(x) = \frac{f(x)}{\ell(x)}$

where $x \in \mathbb{R}^B$ is a vector, $m(x)$ is the maximum value (subtracted for numerical stability to avoid overflow when exponentiating large numbers), $f(x)$ is the vector of exponentiated shifted values, and $\ell(x)$ is the sum of those exponentiated values (the softmax denominator).

What this computes: The numerically stable softmax of a vector. Subtracting the maximum before exponentiation ensures that the largest exponentiated value is $e^0 = 1$, preventing overflow. $\ell$ is the normalizing constant (the softmax denominator), and $\text{softmax}(x)$ is a probability vector summing to 1.

Now, for two vectors $x^{(1)}, x^{(2)}$ that we want to concatenate and softmax together, we can compute the softmax of the concatenation $x = [x^{(1)}, x^{(2)}]$ from the statistics of the individual blocks:

$m(x) = m([x^{(1)}, x^{(2)}]) = \max(m(x^{(1)}), m(x^{(2)}))$

$f(x) = [e^{m(x^{(1)}) - m(x)} f(x^{(1)}), \; e^{m(x^{(2)}) - m(x)} f(x^{(2)})]$

$\ell(x) = \ell([x^{(1)}, x^{(2)}]) = e^{m(x^{(1)}) - m(x)} \ell(x^{(1)}) + e^{m(x^{(2)}) - m(x)} \ell(x^{(2)})$

$\text{softmax}(x) = \frac{f(x)}{\ell(x)}$

What this computes: The softmax of the concatenated vector $[x^{(1)}, x^{(2)}]$, using only the per-block statistics $(m(x^{(1)}), \ell(x^{(1)}))$ and $(m(x^{(2)}), \ell(x^{(2)}))$ from each sub-vector, plus the global maximum $m(x)$. The scaling factors $e^{m(x^{(b)}) - m(x)}$ adjust the old block statistics to the new global maximum before combining them. This means we never need to have all elements of $x$ in memory simultaneously — we can process $x^{(1)}$ first, save its $(m, \ell)$, process $x^{(2)}$, and update the statistics.

Why this form matters for FlashAttention: This decomposition is the mathematical foundation of tiling. For each query $i$, the attention scores $S_{ij}$ for all keys $j$ form a vector of length $N$. If we split the keys into $T_c$ blocks of size $B_c$, each block produces its own local maximum $\tilde{m}_{ij}$ and local sum $\tilde{\ell}_{ij}$. The running statistics $(m_i, \ell_i)$ maintained across blocks allow us to compute the correct softmax-normalized contribution from each block and add it to a running output accumulator — without ever having the full row of S in one place.

The algorithm in detail (Algorithm 1, Section 3.1).

Inputs and setup:

• Q, K, V: $N \times d$ matrices in HBM.
• $M$: the size of on-chip SRAM, used to determine block sizes.
• Block sizes: $B_c = \lfloor M / (4d) \rfloor$ (block size for K and V columns), $B_r = \min(\lfloor M / (4d) \rfloor, d)$ (block size for Q rows). The factor 4 accounts for: one block of Q ($B_r \times d$), one block of K ($B_c \times d$), one block of V ($B_c \times d$), and the intermediate block $S_{ij}$ ($B_r \times B_c$) — all must fit in SRAM simultaneously. The constraint on $B_r$ ($B_r \leq d$) is a secondary optimization: when the head dimension $d$ is small, the block size for Q is capped at $d$, meaning each Q block corresponds to $B_r$ query positions.
• Number of blocks: $T_r = \lceil N / B_r \rceil$ (blocks of Q), $T_c = \lceil N / B_c \rceil$ (blocks of K and V).
• Initialize: $O = \mathbf{0}_{N \times d}$ (output accumulator), $\ell = \mathbf{0}_N$ (softmax denominator accumulator), $m = (-\infty)_N$ (softmax maximum accumulator, initialized to $-\infty$ so that any real score beats it on the first iteration).

Outer loop over K-V blocks (line 5 of Algorithm 1): For $j = 1, \ldots, T_c$:

1. Load a block of keys and values into SRAM (line 6): Load $K_j$ (size $B_c \times d$) and $V_j$ (size $B_c \times d$) from HBM to on-chip SRAM. These are the $j$-th block of keys and values, corresponding to columns $j B_c$ through $(j+1)B_c - 1$ of the attention matrix. Loading from HBM to SRAM is the slow step; once in SRAM, access is fast.

2. Inner loop over Q blocks (line 7): For $i = 1, \ldots, T_r$:

   a. Load a query block and its running statistics (line 8): Load $Q_i$ (size $B_r \times d$), $O_i$ (the current accumulated output for these $B_r$ query rows, size $B_r \times d$), $\ell_i$ (the current softmax denominator for these $B_r$ queries, size $B_r$), and $m_i$ (the current softmax maximum for these $B_r$ queries, size $B_r$) from HBM to SRAM. This is the second slow step — but critically, each Q block is loaded $T_c$ times (once per K-V block). The total HBM reads are $T_c \cdot Nd$ for Q (loading all of Q once per outer loop iteration), plus $T_c \cdot B_c d$ for K and V (loading each K-V block once).

   b. Compute block attention scores (line 9): On-chip, compute $S_{ij} = Q_i K_j^T \in \mathbb{R}^{B_r \times B_c}$. This is a matrix multiplication between a $B_r \times d$ block and a $d \times B_c$ block, producing a $B_r \times B_c$ sub-matrix of the full attention score matrix. This sub-matrix contains the scores for queries in block $i$ attending to keys in block $j$. All operations are in SRAM, so no HBM accesses are involved in the computation — only the initial loads from HBM.

   c. Compute per-block softmax statistics (line 10): On-chip, compute:

   • $\tilde{m}_{ij} = \text{rowmax}(S_{ij}) \in \mathbb{R}^{B_r}$: the maximum score in each row of this block (for each of the $B_r$ queries, the maximum over the $B_c$ keys in this block).
   • $\tilde{P}_{ij} = \exp(S_{ij} - \tilde{m}_{ij}) \in \mathbb{R}^{B_r \times B_c}$: pointwise exponentiation of shifted scores (numerically stable).
   • $\tilde{\ell}_{ij} = \text{rowsum}(\tilde{P}_{ij}) \in \mathbb{R}^{B_r}$: the sum of exponentiated scores per row (per query) for this block — the block's contribution to the softmax denominator.

   d. Update running statistics with softmax scaling (line 11): On-chip, compute:

   • $m_i^{\text{new}} = \max(m_i, \tilde{m}_{ij})$: the new running maximum — the maximum score seen so far for each query, across all K-V blocks processed so far (blocks 1 through $j$).
   • $\ell_i^{\text{new}} = e^{m_i - m_i^{\text{new}}} \ell_i + e^{\tilde{m}_{ij} - m_i^{\text{new}}} \tilde{\ell}_{ij}$: the new running denominator. The first term rescales the old denominator $\ell_i$ from the old maximum $m_i$ to the new maximum $m_i^{\text{new}}$. The second term rescales the current block's denominator $\tilde{\ell}_{ij}$ from its local maximum $\tilde{m}_{ij}$ to the new global maximum $m_i^{\text{new}}$, then adds it to the accumulated sum. This implements the softmax decomposition equation exactly.

   e. Update the output accumulator (line 12): On-chip, write back to HBM: $O_i \leftarrow \text{diag}(\ell_i^{\text{new}})^{-1} \left( \text{diag}(\ell_i) e^{m_i - m_i^{\text{new}}} O_i + e^{\tilde{m}_{ij} - m_i^{\text{new}}} \tilde{P}_{ij} V_j \right)$

   What this computes: The updated partial output for the $B_r$ queries in block $i$, incorporating the contribution from the $j$-th block of values $V_j$. Let's break it down:

   • The term $\text{diag}(\ell_i) e^{m_i - m_i^{\text{new}}} O_i$ rescales the old accumulated output $O_i$: $O_i$ was computed with normalization $\ell_i$, but we need to re-normalize to the new denominator $\ell_i^{\text{new}}$ and account for the change in maximum from $m_i$ to $m_i^{\text{new}}$.
   • The term $e^{\tilde{m}_{ij} - m_i^{\text{new}}} \tilde{P}_{ij} V_j$ is the contribution from the current block: $\tilde{P}_{ij}$ (already row-shifted by $\tilde{m}_{ij}$) is re-scaled to the global maximum $m_i^{\text{new}}$, then multiplied by $V_j$ to produce the weighted sum of value vectors for this block.
   • The multiplication by $\text{diag}(\ell_i^{\text{new}})^{-1}$ divides each row by the new softmax denominator, producing the correctly-normalized partial output.

   In simpler terms: we rescale everything (old output and new block contribution) to the new maximum and new denominator, then combine them. This incremental update ensures that at the end of processing all $T_c$ K-V blocks, $O_i$ equals the exact softmax output for these $B_r$ queries.

   f. Write updated statistics (line 13): Write $\ell_i \leftarrow \ell_i^{\text{new}}$ and $m_i \leftarrow m_i^{\text{new}}$ to HBM.

3. After the inner loop finishes all Q blocks for the current K-V block $j$, the outer loop advances to K-V block $j+1$. Each K-V block is loaded exactly once; each Q block is loaded $T_c$ times.

After all $T_c$ outer loop iterations, the output O in HBM is exactly $\text{softmax}(QK^T)V$ — the correct exact attention output.

Theorem 1 (Section 3.1) states the correctness and complexity:

"Algorithm 1 returns $O = \text{softmax}(QK^T)V$ with $O(N^2 d)$ FLOPs and requires $O(N)$ additional memory beyond inputs and output."

What this theorem means: FlashAttention computes the exact same mathematical function as standard attention. The FLOP count is asymptotically the same — $O(N^2 d)$ for both. The difference is not in FLOPs but in how data moves. FlashAttention requires only $O(N)$ extra memory (for the statistics $m$ and $\ell$), compared to $O(N^2)$ for standard attention (for the S and P matrices). The proof is by induction on the outer loop counter $j$, showing that $m^{(j)}, \ell^{(j)}, O^{(j)}$ after $j$ iterations equal the correct statistics computed over the first $j B_c$ columns of K and V (full proof in Appendix C).

Why this algorithm structure is chosen: The outer loop over K-V blocks and inner loop over Q blocks maximizes data reuse. Each K-V block is loaded once into SRAM and reused across all Q blocks — for the entire inner loop, the keys and values stay in fast SRAM and are multiplied against each query block that gets loaded. This reduces HBM accesses from $O(N^2)$ (where every $(i,j)$ score requires a separate read of $k_j$) to $O(N \cdot (N/B_c)) = O(N^2 d / M)$ (where each element of K is read once per outer loop iteration, and there are $T_c = N/B_c$ outer loop iterations). The factor $M$ (SRAM size) in the denominator is what makes this sub-quadratic — larger SRAM allows larger blocks, which means fewer passes over the data.

Implementation detail — kernel fusion. The paper notes that tiling "enables us to implement our algorithm in one CUDA kernel, loading input from HBM, performing all the computation steps (matrix multiply, softmax, optionally masking and dropout, matrix multiply), then write the result back to HBM." This avoids all intermediate HBM writes between operations — there is no separate S matrix, no separate P matrix. Everything stays in registers or SRAM until the final accumulated O is written back.

The full forward pass (Algorithm 2 in Appendix B.3) additionally includes:

• Softmax scaling: $S = \tau QK^T$ where $\tau = 1/\sqrt{d}$ (standard scaling).
• Masking: $\text{mask}(S)$ sets some entries to $-\infty$ (e.g., for causal masking or padding).
• Dropout: applied elementwise to P, with probability $p_{\text{drop}}$.
• Random number generator state R: saved to HBM so the dropout mask can be exactly reproduced during the backward pass, avoiding the need to store the mask itself ($O(N^2)$ integers).

The pseudocode in Algorithm 2 also clarifies the block size selection:

$B_c = \left\lfloor \frac{M}{4d} \right\rfloor, \quad B_r = \min\left( \left\lfloor \frac{M}{4d} \right\rfloor, d \right)$

where $M$ is SRAM size, $d$ is head dimension. The factor 4 comes from needing room for Q block ($B_r d$), K block ($B_c d$), V block ($B_c d$), and S block ($B_r B_c$), plus workspace for the softmax statistics. The constraint $B_r \leq d$ ensures the Q block is not wider than the head dimension, which is a reasonable bound.

---

3.4.6 FlashAttention Backward Pass (Algorithm 4): Recomputation Instead of Storage

The backward pass faces a different challenge from the forward pass: it needs the attention probabilities P to compute gradients. Standard attention reads P from HBM (where it was stored during the forward pass). FlashAttention does not store P — it stores only O (the output) and the statistics $(m, \ell)$.

The recomputation strategy. During the backward pass, FlashAttention recomputes the attention probabilities P from the stored statistics and the original Q, K, V. Specifically, for each block, given the stored $m_i$ and $\ell_i$ (the final running maximum and sum after the full forward pass), the exact softmax probability matrix for that block is:

$P_{ij} = \text{diag}(\ell_i)^{-1} \exp(S_{ij} - m_i)$

where $S_{ij} = \tau Q_i K_j^T$.

What this computes: The exact softmax probabilities $P_{ij}$ for the $i$-th block of queries and $j$-th block of keys, reconstructed from the stored running statistics $(m_i, \ell_i)$ and re-computed attention scores $S_{ij}$. This requires recomputing the matrix multiplication $Q_i K_j^T$ (which is fast, since it's compute-bound), plus one exponentiation and scaling (memory-bound but now on-chip). The key point: this avoids reading the $B_r \times B_c$ matrix P from HBM — it is recomputed on-chip using Q, K, $m$, and $\ell$, all of which are already in SRAM for the backward pass's own tiling loops.

Algorithm 4 (Appendix B.4) structure:

The backward pass mirrors the forward pass's tiling structure (outer loop over K-V blocks, inner loop over Q blocks), but computes gradients rather than outputs:

Inputs: Q, K, V, O, dO (all $N \times d$ in HBM), $\ell$, $m$ (both $N$ in HBM), plus the random number generator state R (to exactly reproduce the dropout masks).

Initialization: dQ, dK, dV are initialized to zero in HBM.

Outer loop over K-V blocks ($j = 1, \ldots, T_c$):

1. Load $K_j$, $V_j$ from HBM to SRAM.
2. Initialize on-chip accumulators $\widetilde{\text{dK}}_j = \mathbf{0}_{B_c \times d}$ and $\widetilde{\text{dV}}_j = \mathbf{0}_{B_c \times d}$ — these accumulate gradients for the $j$-th block of keys and values across all query blocks.
3. Inner loop over Q blocks ($i = 1, \ldots, T_r$): a. Load $Q_i$, $O_i$, $\text{dO}_i$, $\text{dQ}_i$, $\ell_i$, $m_i$ from HBM to SRAM. b. Recompute $S_{ij} = \tau Q_i K_j^T$ (the same computation as the forward pass). c. Recompute $P_{ij} = \text{diag}(\ell_i)^{-1} \exp(S_{ij} - m_i)$ (reconstructed from stored statistics). Optionally apply the same masking and dropout as the forward pass, using R to regenerate the dropout mask $Z_{ij}$. d. Apply dropout: $P^{\text{dropped}}_{ij} = P_{ij} \odot Z_{ij}$. e. Accumulate dV: $\widetilde{\text{dV}}_j \leftarrow \widetilde{\text{dV}}_j + (P^{\text{dropped}}_{ij})^T \text{dO}_i$. f. Compute $\text{dP}^{\text{dropped}}_{ij} = \text{dO}_i V_j^T$ and $\text{dP}_{ij} = \text{dP}^{\text{dropped}}_{ij} \odot Z_{ij}$ (backprop through dropout). g. Compute $D_i = \text{rowsum}(\text{dO}_i \odot O_i)$ — the scalar gradient through the softmax normalization, computed as $do_i^T o_i$ per row (the simplification from Appendix B.2). h. Compute softmax gradient: $\text{dS}_{ij} = P_{ij} \odot (\text{dP}_{ij} - D_i)$. i. Accumulate dQ: $\text{dQ}_i \leftarrow \text{dQ}_i + \tau \text{dS}_{ij} K_j \in \mathbb{R}^{B_r \times d}$, and write the updated dQ_i to HBM (dQ_i is accumulated in HBM across K-V blocks, but written back each time in the inner loop). j. Accumulate dK: $\widetilde{\text{dK}}_j \leftarrow \widetilde{\text{dK}}_j + \tau \text{dS}_{ij}^T Q_i \in \mathbb{R}^{B_c \times d}$ — this stays in SRAM within the outer loop iteration.
4. After the inner loop completes (all Q blocks processed for this K-V block $j$), write $\widetilde{\text{dK}}_j$ and $\widetilde{\text{dV}}_j$ from SRAM to HBM (they are final for this block).

After all outer loop iterations, dQ, dK, dV in HBM are the correct gradients.

Why recomputation is faster than storage: The paper's key empirical claim (Figure 2, left) is that even though the backward pass computes more FLOPs (75.2 GFLOPs for FlashAttention vs. 66.6 GFLOPs for standard attention, a 13% increase due to recomputing $S_{ij}$ and $P_{ij}$), it performs far fewer HBM accesses (4.4 GB vs. 40.3 GB, a 9× reduction). The total runtime drops from 41.7 ms to 7.3 ms — a 5.7× speedup. The increased FLOPs happen entirely in SRAM (compute-bound, fast), while the reduced HBM accesses eliminate the memory-bound bottleneck. This is the paper's central hardware insight: on modern GPUs, extra compute is cheap if it saves memory bandwidth.

Comparison with Rabe and Staats (2021). Appendix B.5 provides a detailed comparison. Both methods use tiling and avoid storing the $N \times N$ attention matrix. The key differences are:

1. Goal: Rabe and Staats focus on reducing peak memory (the maximum GPU memory required). FlashAttention focuses on reducing memory accesses (the number of reads/writes), which determines runtime. Reducing accesses automatically reduces peak memory (since data doesn't need to be stored if it's never written), but the converse is not true.
2. Incremental vs. temporary outputs: Rabe and Staats compute temporary block outputs and combine them at the end using the softmax statistics, requiring $K$ copies of temporary outputs for $K$ blocks. FlashAttention incrementally updates a single output O after each block (Algorithm 1 line 12), requiring only one copy. This makes FlashAttention more memory-efficient.
3. Backward pass strategy: Rabe and Staats uses gradient checkpointing to recompute both the attention matrix and the temporary block outputs. FlashAttention simplifies the backward analytically (Appendices B.2, B.4), only recomputing the attention matrix — not the temporary outputs. This makes the backward pass faster.

---

3.4.7 IO Complexity Analysis (Theorem 2, Proposition 3): Why FlashAttention Is Theoretically Better

The paper provides a formal IO complexity analysis — counting the number of HBM accesses — to quantify exactly how much data movement FlashAttention saves.

Theorem 2 (Section 3.2):

"Let $N$ be the sequence length, $d$ be the head dimension, and $M$ be size of SRAM with $d \leq M \leq Nd$. Standard attention (Algorithm 0) requires $\Theta(Nd + N^2)$ HBM accesses, while FlashAttention (Algorithm 1) requires $\Theta(N^2 d^2 M^{-1})$ HBM accesses."

What this means in practice: For typical values ($d = 64$, $M \approx 100$KB ≈ 100,000 bytes, with each float16 element being 2 bytes, so $M$ can hold ~50,000 float16 values), we have $d^2 / M \approx 4096 / 50000 \approx 0.08$, meaning FlashAttention performs roughly 12.5× fewer HBM accesses than standard attention. The paper's Figure 2 shows empirically measured reductions of 5–9×, consistent with this analysis (the constant factors in $\Theta$ matter).

Derivation of the FlashAttention bound:

The algorithm divides K and V into $T_c = N / B_c$ blocks, where each block must fit in SRAM along with one block of Q and the intermediate S. The block size selection gives:

$B_c = \Theta\left(\frac{M}{d}\right), \quad B_r = \Theta\left(\min\left(\frac{M}{d}, d\right)\right)$

For typical head dimensions $d \ll M/d$ (i.e., $d^2 \ll M$), we have $B_r = \Theta(d)$. Then:

$T_c = \frac{N}{B_c} = \Theta\left(\frac{Nd}{M}\right)$

In the outer loop, each K-V block is loaded once: $\Theta(B_c d) = \Theta(M)$ elements per block, times $T_c$ blocks, for $\Theta(M \cdot Nd/M) = \Theta(Nd)$ HBM accesses for K and V — this is linear in $N$.

In the inner loop, for each of the $T_c$ K-V blocks, all of Q and the output O are loaded: $\Theta(Nd + Nd) = \Theta(Nd)$ elements per outer loop iteration. Over $T_c$ iterations, this gives $\Theta(Nd \cdot T_c) = \Theta(N^2 d^2 / M)$ HBM accesses — this is the dominant term.

The total is $\Theta(Nd + N^2 d^2 M^{-1})$, which simplifies to $\Theta(N^2 d^2 M^{-1})$ when $N^2 d^2 M^{-1} \gg Nd$ (i.e., for large $N$).

What this bound tells us about design choices: The key variable is $M$ (SRAM size) in the denominator: larger SRAM means larger blocks, fewer outer loop iterations, and fewer HBM accesses. This explains several practical observations:

• FlashAttention is more effective on newer GPUs with larger SRAM (A100 has 192 KB per SM vs. T4's smaller SRAM — the paper's Appendix E.5 confirms less speedup on T4).
• The block size $B_c$ should be as large as possible while fitting in SRAM. Figure 2 (middle) shows that increasing block size from 64 to 256 reduces HBM accesses and runtime, but beyond 256, runtime bottlenecks on compute rather than memory.
• The head dimension $d$ appears squared in the numerator of the dominant term: larger head dimensions increase HBM accesses quadratically because each element of the attention matrix requires more data movement. Appendix E.5 confirms less speedup with head dimension 128 vs. 64.

Proposition 3 (lower bound):

"Let $N$ be the sequence length, $d$ be the head dimension, and $M$ be size of SRAM with $d \leq M \leq Nd$. There does not exist an algorithm to compute exact attention with $o(N^2 d^2 M^{-1})$ HBM accesses for all $M$ in the range $[d, Nd]$."

What this proves: FlashAttention is asymptotically optimal — no exact attention algorithm can do fundamentally better for all SRAM sizes. The proof is by contradiction: if an algorithm achieved $o(N^2 d^2 M^{-1})$ HBM accesses for all $M$, then in the regime $M = \Theta(Nd)$ (SRAM large enough to hold the equivalent of a full row or column), the bound would become $o(Nd)$. But the inputs and outputs themselves are size $\Theta(Nd)$, so any algorithm that computes exact attention must at least read the inputs and write the outputs — requiring $\Omega(Nd)$ HBM accesses. This contradicts $o(Nd)$.

Why this lower bound matters: It establishes that the problem FlashAttention solves is not just an engineering optimization — there is a fundamental information-theoretic limit. The only way to further reduce HBM accesses is to either (a) use approximation (giving up exactness) or (b) increase SRAM size. FlashAttention achieves the theoretical lower bound (within constants), making it IO-optimal.

The paper leaves "proving parameterized complexity lower bounds in terms of $M$ as exciting future work" — meaning a more precise characterization of the constant factors and trade-offs, beyond the asymptotic $\Theta$ notation.

Corollary — backward pass (Theorem 5, Appendix C):

"Standard attention backward pass requires $\Theta(Nd + N^2)$ HBM accesses, while FlashAttention backward pass requires $\Theta(N^2 d^2 M^{-1})$ HBM accesses."

The same analysis applies to the backward pass, with the same result. The backward pass has the same tiling structure (outer loop over K-V blocks, inner loop over Q blocks), and the dominant term comes from the multiple passes over Q and dQ.

---

3.4.8 Block-Sparse FlashAttention (Algorithm 5, Proposition 4): IO-Aware Sparsity

The paper extends FlashAttention to handle block-sparse attention — where only certain blocks of the attention matrix are computed, and the rest are masked out (set to $-\infty$ before softmax). This is not a new sparsity pattern; it's a demonstration that IO-awareness makes sparsity actually fast.

The problem with naive sparse attention. Sparse attention methods (e.g., Reformer, Sparse Transformers, Longformer, BigBird) skip certain key-query pairs to reduce FLOPs. But if implemented naively, the skipping logic — checking which pairs to compute, conditionally branching, gathering/scattering indices — adds overhead that can negate the FLOP savings. Many sparse attention methods report theoretical FLOP reductions but fail to achieve proportional wall-clock speedups.

Block-sparse FlashAttention. Given a pre-defined block sparsity mask $M \in \{0, 1\}^{N/B_r \times N/B_c}$ indicating which blocks of the attention matrix should be computed ($M_{ij} = 1$) and which should be masked out ($M_{ij} = 0$), the algorithm is identical to Algorithm 2 except that it skips zero blocks entirely.

The forward pass (Algorithm 5 in Appendix D.1) adds exactly one check (line 8): if $M_{ij} \neq 0$ then before loading the Q block and computing the attention sub-matrix. If the block is masked out, the inner loop iteration is skipped — no HBM accesses for Q loading, no computation, no statistics update.

Why this yields proportional speedup: Because the dominant cost in FlashAttention is HBM accesses for loading Q blocks (each Q block is loaded $T_c$ times), skipping a fraction $1-s$ of the inner loop iterations reduces HBM accesses proportionally. The paper proves:

Proposition 4 (Section 3.3):

"Block-sparse FlashAttention requires $\Theta(Nd + N^2 d^2 M^{-1} s)$ HBM accesses where $s$ is the fraction of nonzero blocks in the block-sparsity mask."

The dominant term is scaled by $s$, the fraction of blocks that are actually computed. If 75% of blocks are masked out ($s = 0.25$), HBM accesses are reduced by roughly 4× (modulo the linear $Nd$ term, which becomes significant for small $s$). The paper validates this in Figure 2 (right): runtime decreases proportionally with sparsity (inverse of the fraction of non-zero blocks).

Contribution beyond the sparsity pattern itself: The block-sparsity pattern used in experiments is the "fixed butterfly" pattern from Dao et al. (2022) and Chen et al. (2021), which has been shown to approximate arbitrary sparsity patterns well. But the paper's key point is not that butterfly sparsity is good — it's that any block-sparse pattern (butterfly, sliding window, random blocks, learned blocks) can be made IO-aware using the same tiling approach, and will achieve proportional speedups because the FlashAttention infrastructure already amortizes the per-block overhead.

Implications for the attention landscape: The paper uses block-sparse FlashAttention to achieve:

• 2.8× speedup over FlashAttention on LRA benchmark tasks, while matching accuracy (Table 3).
• The first Transformer to achieve better-than-chance performance on Path-256 (63.1% accuracy on 64K sequence length), because the $O(N^2 d^2 M^{-1} s)$ HBM accesses with $s = O(N^{-1/2})$ for butterfly sparsity translates to $O(N^{3/2})$ instead of $O(N^2)$ — making 64K sequences computationally tractable.

This demonstrates that IO-awareness is a multiplier on sparsity: a 10× sparse pattern implemented without IO-awareness might get 2× speedup (due to overhead), while the same pattern with FlashAttention gets nearly the full 10×.

4. Key Insights and Innovations

Innovation 1: Memory Bandwidth, Not FLOPs, Is the Fundamental Bottleneck for Attention — and IO-Awareness Is the Missing Design Principle

This paper's most important conceptual move is reframing what "fast attention" means from a FLOP-counting problem to a data-movement problem. The dominant assumption in the efficient Transformers literature — represented by dozens of approximate attention methods (Reformer, Linformer, Performer, BigBird, Longformer, and many others surveyed in Appendix A) — was that the $O(N^2)$ arithmetic complexity of attention was the root cause of its slowness, and therefore the solution was to reduce the asymptotic FLOP count through sparsification or low-rank approximation. FlashAttention rejects this diagnostic entirely. The paper demonstrates — through both formal IO complexity analysis and extensive benchmarking — that standard attention is memory-bound, not compute-bound, meaning its runtime is determined by how long it takes to move the $O(N^2)$ intermediate attention matrix S and softmax probabilities P between GPU high bandwidth memory and the compute units, not by how many multiplications are performed.

This reframing is more than a vocabulary change. It explains a puzzle that the efficient Transformers literature had struggled with for years: why methods that report linear or near-linear FLOP scaling often fail to achieve proportional wall-clock speedups, especially at moderate sequence lengths where Transformers are most commonly deployed. The paper's answer — documented in the extensive benchmarking of Figure 3 and Appendix E Tables 9–20 — is that those methods reduced FLOPs but not memory accesses, or introduced overhead (sparse indexing, hashing, gather/scatter) that consumed the FLOP savings in memory traffic. FlashAttention, by contrast, achieves 2–4× speedup on standard Transformer training (GPT-2, BERT, LRA) without changing the mathematical function being computed — proving that IO-awareness alone can extract substantial gains that approximate methods could not.

The innovation is not the observation that GPUs have memory hierarchies — that is standard computer architecture knowledge. Rather, it is the systematic application of IO-complexity analysis (from Aggarwal and Vitter, 1988) to deep learning operations, and the proof (Theorem 2, Proposition 3) that attention has a well-defined IO complexity that can be optimized to a theoretical lower bound. Prior work had applied kernel fusion (e.g., Nvidia's Apex FMHA) to reduce some memory traffic, but without a formal framework for understanding how much reduction was possible or which memory accesses were avoidable. FlashAttention provides that framework: it identifies the $O(N^2 d^2 M^{-1})$ bound as the achievable limit for exact attention on hardware with SRAM size $M$, proves that standard attention's $\Theta(Nd + N^2)$ is far above this bound, and demonstrates an algorithm that achieves the lower bound within constants. This transforms attention optimization from a heuristic engineering activity into a principled design space governed by a single parameter ($M$, the SRAM budget) — a conceptual advance that generalizes beyond the specific CUDA kernel presented.

The significance extends beyond this paper's performance numbers. IO-awareness as a design principle applies to every memory-bound operation in deep learning — the paper explicitly flags this in Section 5: "Attention is the most memory-intensive computation in Transformers, but every layer in a deep network touches GPU HBM. We hope our work inspires IO-aware implementations of additional modules." The paper's demonstration that recomputation (extra FLOPs) plus reduced memory accesses yields net speedup (Figure 2: 75.2 GFLOPs vs 66.6, but 7.3 ms vs 41.7 ms) is a fundamental tradeoff insight that inverts the default optimization instinct of minimizing arithmetic. On modern hardware, arithmetic is abundant; memory bandwidth is the scarce resource worth conserving.

---

Innovation 2: Tiling + Recomputation as a Joint Strategy Enables Exact Attention with Sub-Quadratic Memory Accesses

While the individual techniques of tiling (block-wise processing) and recomputation (recalculating intermediate values during backpropagation rather than storing them) were separately known in scientific computing and deep learning, FlashAttention's intellectual contribution is recognizing that the combination of these two techniques solves a problem that neither solves alone — and that the combination is what makes exact attention truly fast, not just memory-efficient.

The paper explicitly builds on prior work that applied each technique in isolation to attention. Rabe and Staats (2021) used tiling and softmax scaling to compute attention with $O(N)$ memory, but "is around the same speed or slightly slower than standard attention" (Appendix B.5) because it still makes $O(N^2)$ HBM accesses per forward pass — the tiling avoids storing the $N \times N$ matrix, but doesn't reduce how many times Q, K, and V are read. Nvidia's Apex FMHA used kernel fusion (a form of tiling) to avoid writing intermediate matrices between operations, but still stored the attention probabilities P to HBM for the backward pass, limiting it to $O(N^2)$ memory and making it unusable beyond sequence length 512 (Table 7, Appendix E.4). Gradient checkpointing (Chen et al., 2016) avoids storing intermediate activations by recomputing them during backpropagation, but trades speed for memory — it makes the backward pass slower, not faster.

FlashAttention's key move is seeing that recomputation can be faster than storage when it eliminates $O(N^2)$ HBM reads. Standard attention stores the $N \times N$ softmax matrix P during forward and reads it back during backward — each of those $N^2$ reads costs HBM bandwidth. FlashAttention stores only the $O(N)$ statistics $(m, \ell)$ and the $N \times d$ output O, then recomputes P on-chip during the backward pass by recalculating $S_{ij} = \tau Q_i K_j^T$ and reconstructing $P_{ij} = \text{diag}(\ell_i)^{-1} \exp(S_{ij} - m_i)$ directly in SRAM. The recomputation increases FLOPs by ~13% (66.6 → 75.2 GFLOPs), but reduces HBM accesses by ~9× (40.3 → 4.4 GB) — and because attention is memory-bound, the net effect is a 5.7× speedup on the forward+backward pass (Figure 2, left). This inverts the standard tradeoff: more computation yields less runtime.

This insight is counterintuitive because the deep learning community has internalized FLOP-counting as the primary efficiency metric. Model efficiency papers routinely report FLOP reductions; training recipes minimize unnecessary recomputation. FlashAttention demonstrates that on modern GPU hardware, FLOPs are the wrong currency for memory-bound operations. Spending extra FLOPs to save memory bandwidth is a net win. This has implications beyond attention — any operation where intermediate results are cheaper to recompute than to store and reload (e.g., certain normalization layers, activation functions in large models) could benefit from the same analysis.

The paper also contributes a nuanced comparison with the closest prior work (Rabe and Staats, 2021) in Appendix B.5, identifying three specific differences that explain FlashAttention's speed advantage:

1. Incremental vs. batched output accumulation: Rabe and Staats compute temporary outputs for each block and combine them at the end, requiring $K$ copies of temporary output blocks. FlashAttention updates a single output O incrementally after each block (Algorithm 1 line 12), using the running softmax statistics to rescale — this requires only one copy of the output.
2. Analytical backward pass simplification: Rabe and Staats use generic gradient checkpointing, which recomputes both the attention matrix and all temporary block outputs. FlashAttention simplifies the backward pass analytically (Appendix B.2), showing that the gradient $D_i = P_{i:}^T dP_{i:}$ can be computed as $do_i^T o_i$ — a dot product of two $d$-dimensional vectors rather than an $N$-dimensional sum — eliminating the need to recompute temporary block outputs.
3. IO-complexity as the optimization target: Rabe and Staats optimized for peak memory. FlashAttention optimized for memory accesses. Reducing accesses necessarily reduces peak memory (since data never written doesn't consume capacity), but the converse is false — and the paper shows that optimizing for accesses is what unlocks speed.

---

Innovation 3: Proportional Sparsity Speedups Through IO-Aware Implementation

The efficient Transformers literature had produced dozens of sparse attention patterns (sliding window, dilated windows, random blocks, hashing-based, learned sparsity) that theoretically reduced attention complexity from $O(N^2)$ to $O(N \sqrt{N})$ or $O(N \log N)$. Yet "many of them do not display wall-clock speedup against standard attention and have not gained wide adoption" (Section 1). The paper identifies a specific reason for this failure and provides a systematic fix.

The problem is overhead amplification: sparse attention methods replace dense matrix multiplication — which is highly optimized on GPUs and achieves near-peak hardware utilization — with irregular gather/scatter operations, conditional branching, and index computation. These overheads are themselves memory-bound (index lookups touch memory), so the net speedup is often far less than the sparsity ratio would predict. A method that computes only 10% of attention scores might achieve only 2× speedup because the overhead of deciding which 10% to compute consumes the other 80% of the gains.

Block-sparse FlashAttention (Section 3.3, Algorithm 5) solves this by making the overhead of skipping blocks proportional to the IO cost it saves. Because FlashAttention's runtime is dominated by loading Q blocks from HBM in the inner loop, and because each block load costs a known amount of bandwidth, skipping a block that would have been loaded saves exactly that bandwidth. The conditional check if $M_{ij} \neq 0$ costs essentially nothing (a register comparison), and there is no per-element indexing overhead because the block is either fully computed or fully skipped — the granularity matches the tiling structure. Proposition 4 formalizes this: the IO complexity scales as $\Theta(N^2 d^2 M^{-1} s)$ where $s$ is the fraction of non-zero blocks, meaning the speedup is linear in sparsity for large $N$ — a 10× sparser pattern yields ≈10× fewer HBM accesses (modulo the linear $Nd$ term).

The empirical validation is clean: Figure 2 (right) shows runtime decreasing proportionally with sparsity across a range of non-zero block fractions (from 20% to 100%), and block-sparse FlashAttention achieves 2.8× speedup on LRA while matching dense attention accuracy (Table 3). More dramatically, it enables the first Transformer to achieve better-than-chance performance on Path-256 (63.1% accuracy on 64K sequences) — a task that standard attention cannot fit in GPU memory, and that all prior approximate attention methods had failed on (Table 6).

This innovation is significant beyond the specific implementation because it identifies IO-awareness as an enabling condition for sparse attention to realize its theoretical promise. The paper suggests that many sparse attention methods were not intrinsically flawed — they were implemented without accounting for memory access costs, and therefore their FLOP reductions were swamped by IO overhead. Block-sparse FlashAttention demonstrates that any block-structured sparsity pattern (butterfly, sliding window, random, learned) can achieve proportional speedups if implemented on the FlashAttention tiling infrastructure. This reframes the sparse attention research agenda: the bottleneck was not sparsity pattern design but hardware-efficient implementation.

---

Innovation 4: Formal IO Lower Bound Proving FlashAttention Is Asymptotically Optimal for Exact Attention

The paper contributes a theoretical result with practical teeth: Proposition 3 proves that no exact attention algorithm can achieve asymptotically fewer HBM accesses than FlashAttention's $\Theta(N^2 d^2 M^{-1})$ bound across all SRAM sizes $M$ in the range $[d, Nd]$. This is more than a mathematical garnish — it establishes a ceiling on what further algorithm engineering can achieve for exact attention, and it provides a principled explanation for why further optimizations must either increase $M$ (larger SRAM, i.e., better hardware) or abandon exactness (approximation).

The proof is elegantly simple: in the regime $M = \Theta(Nd)$ — SRAM large enough to hold a full row or column of the attention matrix — any algorithm must at least read the $\Theta(Nd)$-sized inputs and write the $\Theta(Nd)$-sized output, requiring $\Omega(Nd)$ HBM accesses. If an algorithm claimed $o(N^2 d^2 M^{-1})$ for all $M$, then at $M = \Theta(Nd)$ it would require $o(Nd)$ accesses, contradicting the input/output lower bound. Therefore no such algorithm exists. This type of lower bound over a subrange of a parameter is standard in streaming algorithms (Woodruff, 2004) but had not been applied to attention.

The practical implication is: further asymptotic improvements to exact attention are impossible without hardware advances. This directs research effort toward two more productive directions: (1) making the constant factors better (FlashAttention's CUDA implementation can be tuned, but the asymptotic slope is fixed), and (2) using approximation — where block-sparse FlashAttention already achieves $\Theta(N^2 d^2 M^{-1} s)$ with $s \ll 1$, breaking through the exact-attention lower bound by trading accuracy for speed. The paper's combined message is that FlashAttention extracts essentially all the available speedup from IO-awareness for exact attention, and that further gains require either better hardware or acceptable approximation — a clean separation of concerns that structures the research landscape.

This is a fundamental rather than incremental contribution: it changes the question from "can we make attention faster?" (open-ended) to "what is the fastest possible exact attention, and how close are we?" (bounded, answerable). The paper's answer — we are at the lower bound within constants — is a strong closure result that will shape future work.

5. Experimental Analysis

Evaluation Methodology

• Dataset. The paper uses three primary evaluation settings: (1) GPT-2 training on the OpenWebText corpus (Gokaslan et al., 2019), with a random 0.5% held-out as validation; (2) BERT-large training on Wikipedia following the MLPerf 1.1 benchmark setup, using the provided train/validation split and evaluating masked language modeling accuracy on the same 10,000 validation examples as the Nvidia baseline; and (3) the Long-Range Arena (LRA) benchmark (Tay et al., 2020), which consists of five tasks with sequence lengths varying between 1024 and 4096 (ListOps, Text, Retrieval, Image, Pathfinder). Additional experiments on long-document classification use MIMIC-III (Johnson et al., 2016) and ECtHR (Chalkidis et al., 2019, 2021) datasets, and the Path-X/Path-256 tasks from LRA test extreme sequence lengths (16K and 64K).

• Base model(s). The training speed experiments use GPT-2 small (~117M parameters) and GPT-2 medium (~345M parameters) from Radford et al. (2019), and BERT-large (~340M parameters) from Devlin et al. (2019). The LRA experiments use a vanilla Transformer following the implementation in Tay et al. (2020) and Xiong et al. (2021). For long-document classification, a pretrained RoBERTa model (Liu et al., 2019) is used with repeated positional embeddings à la Beltagy et al. (2020). The paper argues these models are "representative of the capabilities of many contemporary LLMs" (Section 4) and produce non-trivial but far-from-saturated performance, leaving room for speed improvements to matter.

• Metrics. Model quality metrics are task-specific: perplexity on the OpenWebText validation set for GPT-2 (both implementations are confirmed to produce identical perplexity curves, Figure 4 in Appendix E.2), masked language modeling accuracy for BERT-large (target 72.0%, Section 4.1), and accuracy for LRA tasks (following the standard evaluation protocol from Tay et al., 2020) and long-document classification (micro F1 for MIMIC-III and ECtHR). The primary efficiency metrics are wall-clock training time (measured in days or minutes for full training runs on 8×A100 GPUs) and attention runtime (forward pass, backward pass, and combined, measured in milliseconds, averaged over 100 measurements). Memory footprint is measured once in megabytes or gigabytes of GPU HBM consumed by the attention operation.

• Baselines. For GPT-2 training: HuggingFace Transformers (Wolf et al., 2020) and Megatron-LM (Shoeybi et al., 2019) implementations of standard attention. For BERT training: the Nvidia MLPerf 1.1 submission (Mattson et al., 2020), which uses Apex FMHA (fused multi-head attention) and set the training speed record at the time. For LRA: standard Transformer attention and eight approximate/sparse attention baselines — Linformer (Wang et al., 2020), Linear Attention (Katharopoulos et al., 2020), Performer (Choromanski et al., 2021), Local Attention (Tay et al., 2020), Reformer (Kitaev et al., 2020), Smyrf (Daras et al., 2020), plus the reference implementations of Block-Sparse Attention from OpenAI (Child et al., 2019), Longformer (Beltagy et al., 2020), and BigBird (Zaheer et al., 2020). Where the paper could not reproduce a baseline's reported performance, it reports the better of its own reproduction or the original paper's number — to be generous to baselines (Appendix E.3). For attention benchmarking: additional baselines include Megatron attention and LongShortFormer (LSFormer; Zhu et al., 2021).

• Generation budget / compute accounting. For training speed experiments, compute is measured as total wall-clock time to reach a target performance level (BERT: 72.0% MLM accuracy; GPT-2: 400K training steps, with identical hyperparameters across implementations) on identical hardware (8×A100 GPUs). For attention benchmarking, compute is measured as runtime in milliseconds for the forward pass, backward pass, or both combined, at a fixed configuration (batch size 16, 8 heads, head dimension 64, with/without dropout and masking). For block-sparse FlashAttention, the sparsity ratio $s$ (fraction of non-zero blocks) serves as the budget parameter — runtime is measured as a function of $s$ to validate proportional speedup (Figure 2 right). FLOP counts are reported (Figure 2 left) but treated as secondary to HBM access counts and runtime, reflecting the paper's thesis that memory bandwidth, not arithmetic, is the bottleneck.

• Cross-validation / statistical protocol. The paper does not use cross-validation in the traditional machine learning sense — the experiments are benchmarking and training speed comparisons, not hyperparameter selection. For BERT training time, results are averaged over 10 runs (Table 1 reports mean ± standard deviation: 17.4 ± 1.4 minutes for FlashAttention vs. 20.0 ± 1.5 minutes for Nvidia MLPerf 1.1). For GPT-2 training, the validation set is randomly selected once (0.5% of OpenWebText) and all models are evaluated on the same split. For attention runtime measurements, 100 measurements are taken and averaged. The LRA accuracy results "are known to be highly dependent on the tuning procedure" (Appendix E.3), and the paper follows the reproduction setup from Xiong et al. (2021) to ensure fair comparison.

Main Quantitative Results

Training Speed: BERT and GPT-2

Headline result — BERT-large (Table 1): FlashAttention trains BERT-large to the target 72.0% masked language modeling accuracy in 17.4 ± 1.4 minutes on 8×A100-80GB GPUs, compared to 20.0 ± 1.5 minutes for the Nvidia MLPerf 1.1 record submission — a 15% speedup. Both implementations start from the same initialization and use the same training hyperparameters (LAMB optimizer, learning rate 3.75e-3, batch size 448, FP16 precision with Apex AMP O2 optimization level). The time difference captures the end-to-end effect of replacing FMHA's fused attention kernel with FlashAttention's IO-aware kernel.

Headline result — GPT-2 (Table 2): FlashAttention achieves much larger speedups on GPT-2 training:

All configurations achieve the same perplexity (18.2 for GPT-2 small, 14.2–14.3 for GPT-2 medium), confirming that FlashAttention computes mathematically identical results. The speedup over HuggingFace is substantially larger than over Megatron-LM because Megatron already includes some attention optimizations (kernel fusion), while HuggingFace uses the standard PyTorch implementation. FlashAttention's improvements over Megatron represent the additional gains from IO-awareness beyond existing kernel fusion techniques.

The paper validates that these speedups are not from numerical differences by plotting validation perplexity throughout training (Appendix E.2, Figure 4): "FlashAttention behaves the same as the baseline implementation and the validation perplexity curves of the two implementations almost lie on top of each other." This is a critical sanity check — since FlashAttention changes the order of floating-point operations (block-wise softmax accumulation vs. full-vector softmax), it could theoretically produce slightly different results due to non-associativity of floating-point addition. The identical perplexity curves confirm that any numerical differences are negligible.

The speedups are measured on 8×A100-40GB GPUs (GPT-2) and 8×A100-80GB GPUs (BERT), with the same effective batch size (512 for GPT-2) and gradient accumulation to fit GPU memory. Both FlashAttention and baselines use mixed-precision training (PyTorch AMP).

Training Speed: Long-Range Arena

Headline result — LRA (Table 3): FlashAttention achieves a 2.4× geometric mean speedup over standard Transformer attention on the LRA benchmark's five tasks, while block-sparse FlashAttention achieves a 2.8× speedup. The geometric mean is computed across the five tasks with their different sequence lengths and computational characteristics.

The accuracy results in Table 3 demonstrate that FlashAttention is not trading quality for speed:

FlashAttention slightly outperforms or matches standard attention on all five tasks (the differences are within the expected variance of LRA tuning, which the paper acknowledges is substantial). Block-sparse FlashAttention achieves essentially the same average accuracy (59.6 vs. 59.8) despite computing only a fraction of the attention matrix — evidence that the butterfly sparsity pattern is expressive enough to approximate full attention for these tasks.

The comparison with approximate attention baselines tells a nuanced story. On accuracy, most approximate methods cluster around 54–60% average accuracy, with Linear Attention achieving 59.6% (matching block-sparse FlashAttention). On speedup, block-sparse FlashAttention (2.8×) outperforms all approximate methods, including Linformer (2.5×), Linear Attention (2.3×), Performer (1.8×), and Reformer (1.3×). This empirically validates the paper's thesis: approximate methods reduce FLOPs but their IO overhead limits speedup, while IO-aware exact attention (and its sparse extension) achieve greater speedup without compromising accuracy.

Important caveat on LRA baselines: The paper's reproduced baseline accuracies differ from the original LRA paper (Tay et al., 2020). The authors note this explicitly: "LRA accuracy results are known to be highly dependent on the tuning procedure" (Appendix E.3 footnote). To be generous, they "report the better performance from Tay et al. or Xiong et al. for that baseline on that task" — meaning the baselines may have an unfair advantage (using the best of multiple independent tuning efforts). Despite this generosity, FlashAttention and block-sparse FlashAttention achieve comparable or better accuracy while being substantially faster.

Better Models Through Longer Sequences

Headline result — GPT-2 with longer context (Table 4): FlashAttention enables training GPT-2 small with 4× longer context (sequence length 4096 vs. 1024) while still being 30% faster than Megatron-LM's GPT-2 with context length 1024, and achieving 0.7 better perplexity (17.5 vs. 18.2). This is the paper's most compelling demonstration that IO-awareness enables qualitative model improvements, not just speed:

The key insight: FlashAttention's linear memory footprint ($O(N)$ instead of $O(N^2)$) makes 4K-context training feasible on the same GPU hardware. Standard attention at 4K context would either run out of memory (the $N^2$ matrix for $N=4096$ with 16 heads and batch size requires substantial HBM) or require aggressive gradient accumulation that makes training impractically slow. FlashAttention's reduced memory footprint enables fitting the larger context window into GPU memory, and its reduced memory bandwidth makes training faster despite the larger sequence length. The perplexity improvement from 18.2 to 17.5 is substantial — demonstrating that the model genuinely benefits from the longer context, not just that it's technically possible.

Headline result — Long-document classification (Table 5): Increasing sequence length with FlashAttention yields significant accuracy improvements on two long-document tasks:

On MIMIC-III, sequence length 16K outperforms length 512 by 4.3 points (57.1 vs. 52.8). On ECtHR, sequence length 8K outperforms length 512 by 8.5 points (80.7 vs. 72.2). The paper notes a non-monotonic pattern on MIMIC-III (performance dips at 1024 and 2048 before rising) and hypothesizes this "may be due to subtle distribution shifts: MIMIC-III contains specialized medical text and thus may be more susceptible to a distribution shift in the document length, whereas ECtHR contains general language." The key takeaway is that longer sequences are not uniformly beneficial — they help when the task requires long-range dependencies (as legal document classification does) and can hurt when the training distribution doesn't match the longer-sequence format (as medical text may). FlashAttention makes exploring this tradeoff practical.

Headline result — Path-X and Path-256 (Table 6): FlashAttention achieves 61.4% accuracy on Path-X (sequence length 16K), and block-sparse FlashAttention achieves 63.1% on Path-256 (sequence length 64K). These are the first Transformer results to achieve better-than-chance performance on either benchmark — all prior models (standard Transformer, Linformer, Linear Attention, Performer, Local Attention, Reformer, Smyrf) either ran out of memory or achieved only random performance. The paper states this categorically: "We present here the first result of Transformer models being able to solve Path-X and Path-256."

The training methodology: models are first pretrained on Path-64 (a shorter variant), then positional embeddings are spatially interpolated (duplicated gridwise) to match the longer sequence length, and fine-tuned on Path-X or Path-256 for 200 epochs. For Path-X, an additional round of fine-tuning (taking the best validation checkpoint and fine-tuning for another 200 epochs) "adds roughly 4 points of accuracy to FlashAttention for Path-X, but the model starts overfitting afterwards" (Appendix E.3). This careful tuning procedure suggests that solving Path-X requires both the capacity to model 16K-long dependencies (which only FlashAttention provides by making attention feasible at that scale) and careful optimization (to avoid overfitting on the small training set).

The Path-256 result (63.1% accuracy) is achieved with block-sparse FlashAttention using the butterfly sparsity pattern with sparsity ratio $s = O(N^{-1/2})$, reducing IO complexity to $O(N^{3/2} d^2 M^{-1})$ instead of $O(N^2 d^2 M^{-1})$. This is what makes 64K sequences tractable — even FlashAttention's linear memory footprint and reduced bandwidth would struggle at $N=65536$ without the additional sparsity factor. The paper notes that "Path-256 requires longer sequences but has relatively shorter paths than Path-X, so it is easier to obtain a higher accuracy" (Section 4.2 footnote), explaining why the harder benchmark (64K vs. 16K) yields slightly higher accuracy.

Benchmarking Attention Runtime and Memory

Headline result — Runtime (Figure 3 left, Appendix E Tables 9–20): FlashAttention is "up to 3× faster than the PyTorch implementation" for standard sequence lengths (128–2048) and scales to sequence lengths that cause standard attention to run out of memory. The detailed timing tables in Appendix E provide exhaustive measurements across 16 different attention configurations (with/without dropout, with/without masking, forward/backward/combined) and 10 sequence lengths (128 through 65536). Key patterns:

• At short sequences (128–512): FlashAttention is fastest among all exact attention methods (Table 11, combined forward+backward with dropout and masking: 0.43 ms at 128, 0.41 ms at 256, 0.95 ms at 512). Megatron (0.87, 0.89, 1.33 ms) and PyTorch (0.84, 0.86, 2.35 ms) are slower. Among approximate methods, only Linformer approaches FlashAttention's speed (1.57, 1.49, 1.55 ms), but with lower accuracy on most LRA tasks.
• At moderate sequences (1024–2048): FlashAttention maintains its lead over standard attention (2.55 ms vs. 8.29 ms at 1024, 9.56 ms vs. 31.75 ms at 2048 in Table 11). Approximate methods begin to catch up: Linformer (1.60, 4.19 ms) and block-sparse FlashAttention (0.89, 1.95 ms) are faster, with the latter being the fastest method across all tested configurations at these sequence lengths.
• At long sequences (4096+): Standard attention and Megatron run out of memory. FlashAttention continues to scale (37.49 ms at 4096, 147.75 ms at 8192, 586.61 ms at 16384), with its runtime growing quadratically as expected. Block-sparse FlashAttention's runtime grows closer to linearly (4.12 ms at 4096, 7.64 ms at 8192, 16.60 ms at 16384), demonstrating the proportional speedup from sparsity.
• At extreme sequences (64K): Only FlashAttention (9341.30 ms combined), block-sparse FlashAttention (64.11 ms), and Linformer run without memory errors. All other methods (PyTorch, Megatron, Reformer, Smyrf, LSformer, Block Sparse, Longformer, BigBird, Local Attention) fail before reaching 64K — most fail at 4096 or 8192. The paper notes that "FlashAttention is still 2× more efficient than Linformer" in memory usage at 64K.

The crossover point where approximate methods become faster than dense FlashAttention appears between sequences 512 and 1024 for most methods (Figure 3 left, and the text: "The approximate attention runtimes begin to cross over with FlashAttention at sequences between 512 and 1024"). However, block-sparse FlashAttention "is faster than all implementations of exact, sparse, and approximate attention that we know of, across all sequence lengths" — a strong claim supported by the comprehensive tables.

Headline result — Memory footprint (Figure 3 right, Table 21): FlashAttention's memory usage scales linearly with sequence length, while standard attention scales quadratically and runs out of memory:

At sequence length 512, FlashAttention uses 3.2× less memory than standard attention (104 MB vs. 336 MB). At 4096, the reduction is 20.4× (836 MB vs. 17024 MB) — the paper states FlashAttention is "up to 20× more memory efficient than exact attention baselines." Among approximate methods, FlashAttention is more memory-efficient than all except Linformer at shorter sequences, and at 64K it is 2× more efficient than Linformer (13376 MB vs. 26252 MB).

The memory footprint of FlashAttention and block-sparse FlashAttention is identical (within measurement precision) because both store the same $O(N)$ statistics ($m$, $\ell$, O) and the sparsity mask itself is negligible ($N/B_r \times N/B_c$ bits). The dominant memory cost in both cases is storing the Q, K, V, and output O matrices ($4 \times 8 \times 64 \times N$ bytes for the configuration tested: 8 heads, head dimension 64, batch size 16, float16), which is $O(N)$.

Speedup Variation Across Hardware and Configurations

The paper provides additional benchmarking in Appendix E.5 to characterize how FlashAttention's speedup varies with hardware and configuration parameters. This is important because IO-awareness is inherently hardware-dependent — the gains depend on the ratio of compute speed to memory bandwidth, and on SRAM size.

A100 (Figure 5): 2–4× speedup across sequence lengths 128–4096, with "more speedup when using dropout and masking due to kernel fusion." The kernel fusion eliminates separate HBM writes for the dropout mask and masking operation — operations that are purely memory-bound.

A100 with head dimension 128 (Figure 6): "We see less speedup overall—but we can still see significant speedup (up to 3×) with a causal mask, where half the blocks are masked out." This confirms the IO complexity analysis: larger $d$ increases HBM accesses per block (each Q-K dot product reads more data), so the relative advantage of tiling (which depends on block size constraints) diminishes.

RTX 3090 (Figure 7): "Slightly higher speedups on the RTX 3090 (between 2.5-4.5×), since the memory bandwidth on an RTX 3090 is lower than on an A100 (roughly 900 GB/s vs. 1.5 TB/s)." This is a direct consequence of the memory-bound analysis: when HBM is slower, the relative benefit of reducing HBM accesses is larger. FlashAttention's effectiveness increases on lower-bandwidth hardware.

T4 (Figure 8): "T4 SRAM is smaller than A100, so we need to make the block sizes smaller in FlashAttention. As a result, we observe less speedup on T4, which matches the IO complexity analysis." This validates the $M$ (SRAM size) dependence in Theorem 2: smaller SRAM means smaller blocks, more outer loop iterations ($T_c$), and more HBM accesses passing over Q. The speedup still exists but is attenuated — exactly what the theory predicts.

Ablation Studies and Robustness Checks

FlashAttention vs. Apex FMHA (Table 7, Appendix E.4): The comparison with Nvidia's fused multi-head attention (FMHA), which was the state-of-the-art for BERT-scale attention before FlashAttention, reveals a nuanced tradeoff. FMHA fuses masking, softmax, dropout, and the second matrix multiply into one kernel but stores the softmax matrix P for the backward pass. FlashAttention recomputes P instead of storing it. The result:

At sequence length 128, FlashAttention is marginally slower (4%) because the recomputation overhead outweighs the memory bandwidth savings — at short sequences, the $N^2$ attention matrix is small enough that reading it from HBM is not the dominant cost. At 256 and 512, the balance tips in FlashAttention's favor. Critically, FMHA is limited to sequence length ≤512 and head dimension 64, while FlashAttention supports up to 64K sequences and multiple head dimensions — meaning the slight disadvantage at very short sequences is an acceptable trade for the ability to handle long sequences at all.

Effect of block size on runtime (Figure 2 middle): Varying the block size $B_c$ directly varies the number of HBM accesses (larger blocks → fewer outer loop iterations → fewer Q reloads). The paper shows that as block size increases from 64 to 256, runtime decreases because HBM accesses decrease. Beyond 256, runtime plateaus — "the runtime is then bottlenecked by other factors (e.g., arithmetic operations)." This is the expected crossover from memory-bound to compute-bound: once the working set fits comfortably in SRAM and HBM accesses are minimal, the matrix multiplications themselves become the bottleneck. The optimal block size is determined by SRAM capacity ("larger block size will not fit into the small SRAM size"), making 256 the sweet spot for the A100's 192 KB SRAM per SM.

GPT-2 validation perplexity equivalence (Figure 4, Appendix E.2): Plotting validation perplexity throughout training for GPT-2 small and medium, using both HuggingFace and FlashAttention implementations, the curves "almost lie on top of each other." This confirms that FlashAttention's block-wise softmax with incremental statistics accumulates the same results (to within floating-point precision) as the standard algorithm's full-vector softmax. The robustness check is important because numerical stability is a potential concern: the softmax decomposition relies on exponentiating differences of running maxima, which could amplify floating-point errors if the block maxima vary widely. The identical training curves demonstrate that this is not a problem in practice for the tested configurations.

Sparsity-proportional speedup (Figure 2 right): Block-sparse FlashAttention's runtime at sequence length 4K decreases proportionally as the percentage of non-zero blocks decreases from 100% (dense) to 20%. The relationship is near-linear, validating Proposition 4's $\Theta(N^2 d^2 M^{-1} s)$ bound where $s$ is the fraction of non-zero blocks. This is the key evidence that IO-awareness makes sparsity practically useful — the speedup tracks the theoretical sparsity, without the constant-factor overhead that plagues many sparse attention implementations.

LRA accuracy across methods (Table 3): The robustness of the accuracy results is worth noting: despite substantial differences in how attention is computed (exact vs. approximate, dense vs. sparse, IO-aware vs. not), the LRA accuracy cluster is tight. Standard Transformer achieves 59.3 average, FlashAttention 59.8, block-sparse FlashAttention 59.6, Linear Attention 59.6. This suggests that for these tasks, attention quality is not the primary differentiator — the tasks are learnable with multiple attention variants, and the practical advantage of FlashAttention is the 2.4–2.8× speedup, not improved accuracy. The exception is Path-X/Path-256, where only FlashAttention's ability to handle extreme sequence lengths unlocks above-chance performance — effectively a new capability, not just a speedup.

Long-document classification sensitivity to sequence length (Table 5): The non-monotonic behavior on MIMIC-III (accuracy drops from 52.8 at 512 to 50.7 at 1024 before rising to 57.1 at 16384) is an important negative result embedded in a positive table. It demonstrates that longer sequences are not a panacea — they can hurt performance when the training data distribution shifts. FlashAttention makes it practical to discover this pattern (by enabling training across a range of sequence lengths), but does not change the underlying task difficulty.

Critical Assessment

Do the experiments support the claim that FlashAttention is 15% faster than the MLPerf BERT record? Yes, with appropriate caveats. Table 1 reports 17.4 ± 1.4 minutes vs. 20.0 ± 1.5 minutes, averaged over 10 runs. The standard deviations are small relative to the gap (1.4 and 1.5 minutes vs. a 2.6 minute difference), suggesting statistical reliability. However, the comparison is against a single baseline (Nvidia's MLPerf submission) running on the same hardware — the 15% figure should not be interpreted as a universal speedup for all BERT implementations. The paper's own GPT-2 results show that FlashAttention's advantage over Megatron (1.7×) is much smaller than over HuggingFace (3.0–3.5×), because Megatron already applied some attention optimizations. BERT's MLPerf baseline (using FMHA) is the most optimized pre-FlashAttention implementation, so the 15% gain represents the additional benefit of recomputation-avoiding-storage over kernel-fusion-alone — a narrower but still meaningful improvement.

Do the experiments support the claim of 3× speedup on GPT-2? Yes, with the important clarification that the 3× (actually 3.0× for medium, 3.5× for small) is relative to the HuggingFace baseline, not Megatron. The 1.7× speedup over Megatron is a more conservative estimate of FlashAttention's advantage over a well-optimized baseline. Both figures are valid, but the paper's abstract ("3× speedup on GPT-2") emphasizes the larger number. The validation perplexity equivalence (Figure 4) strongly supports the claim that FlashAttention computes the same results — there's no hidden accuracy-speed tradeoff.

Do the experiments support the claim of enabling the first Transformer to beat chance on Path-X? Yes, convincingly. Table 6 shows that all baselines (standard Transformer, Linformer, Linear Attention, Performer, Local Attention, Reformer, Smyrf) either "ran out of memory, or only achieved random performance" (Section 4.2). FlashAttention achieves 61.4% — substantially above chance (50% for binary classification) and the first published result doing so. The methodology (pretrain on Path-64, interpolate positional embeddings, fine-tune on Path-X) is disclosed in Appendix E.3. A reasonable question: is the success due to FlashAttention's efficiency or to the training methodology? The fact that other methods also have access to Path-64 pretraining (the standard LRA protocol) but fail suggests the efficiency is causal — without FlashAttention, training or even running inference on 16K sequences is infeasible. The Path-256 result is similarly strong, and block-sparse FlashAttention's ability to scale to 64K is what makes it possible.

Do the experiments support the IO-complexity analysis (Theorem 2) as the explanatory mechanism? The empirical evidence is consistent with the theory but does not rigorously test it. Figure 2 (middle) shows that increasing block size reduces runtime — consistent with the $M^{-1}$ dependence. Figure 2 (left) shows that FlashAttention has fewer HBM accesses (4.4 GB vs. 40.3 GB) and is faster (7.3 ms vs. 41.7 ms) despite more FLOPs — consistent with the memory-bound diagnosis. The T4 experiments show less speedup — consistent with smaller SRAM. However, the paper does not systematically vary $M$ (SRAM size), $d$ (head dimension), or $N$ (sequence length) while measuring HBM access counts directly (these are inferred from runtime, not measured via hardware counters). A more rigorous validation would use GPU profiling tools to measure actual DRAM read/write bytes and confirm the $\Theta(N^2 d^2 M^{-1})$ scaling across a wider range of configurations. The qualitative consistency is strong; the quantitative validation is approximate.

What experiments would have strengthened the paper?

1. Direct HBM access measurements. The paper infers HBM accesses from runtime and IO complexity analysis, but never reports actual memory traffic measurements from GPU profilers (e.g., nvprof or Nsight). Figure 2 (left) reports "HBM R/W (GB)" of 40.3 vs. 4.4, but the paper does not specify how these numbers were obtained. Direct measurement would validate the IO complexity analysis more rigorously than runtime correlation.

2. Comparison with more approximate attention methods at scale. The LRA benchmark tests 8 approximate baselines, but the GPT-2 experiments compare only against exact attention (HuggingFace, Megatron). Would training GPT-2 with Linformer or Performer attention also yield speedups? The paper argues those methods have IO overhead, but demonstrating this on a standard language modeling task (not just LRA) would strengthen the claim that IO-awareness is the key missing ingredient.

3. Path-X ablation without FlashAttention. The Path-X result is impressive, but the paper does not report what happens if you try the same training protocol (pretrain on Path-64, interpolate embeddings, fine-tune) with a standard attention implementation that uses gradient accumulation to avoid OOM. Even if training takes 50× longer, does the model eventually learn? If so, the claim shifts from "enabling" to "accelerating" — still valuable but qualitatively different.

4. Larger-scale language modeling. The GPT-2 experiments use small and medium variants (up to 345M parameters). Does FlashAttention produce similar speedups on GPT-2 large (762M) or GPT-3-scale models? The memory savings become more critical as model size grows (since the attention matrices compete with parameters for HBM), so extrapolation is plausible but unverified.

5. Inference-only benchmarking. All training speed experiments measure training time (forward + backward). For inference workloads (where the backward pass is not needed), the speedup from recomputation disappears — FlashAttention's forward pass is still IO-aware, but the advantage over standard attention may be smaller. The T4 inference benchmarking (Figure 8 bottom) partially addresses this, but the paper does not discuss how the training speedup decomposes into forward-vs-backward contributions. The detailed tables in Appendix E show that FlashAttention's forward pass speedup is larger than the backward pass speedup in some configurations (Table 9 vs. Table 10) — the paper could have explored this asymmetry.

What weaknesses exist in the experimental design?

• The difficulty estimation cost dilemma does not apply here (this is a systems paper, not a learning paper), but there is an analogous issue: the paper's reported speedups exclude the one-time cost of writing FlashAttention's CUDA kernel. For researchers adopting FlashAttention, this cost is zero (they use the open-source code). But the paper's broader claim — that IO-awareness is a general principle that should be applied to other operations — is supported by FlashAttention's success but not by any experiments showing that other IO-aware implementations would yield similar gains. The block-sparse extension provides one additional data point, but it's still within the attention family.

• Single hardware platform (Nvidia GPUs) for main results. The paper mentions in Section 2.1 that "performance on other hardware accelerators are similar," citing Jia et al. (2019) on Graphcore IPUs, but all experiments use A100, RTX 3090, or T4 GPUs. The IO-awareness principle should apply to any hardware with a memory hierarchy, but the specific gains (and the optimal block sizes, loop ordering, etc.) are GPU-specific. The paper's openness about this (Section 5 discusses the need for compilers to make IO-aware implementations portable across GPU architectures) partially mitigates the concern, but the experimental evidence is Nvidia-only.

• The butterfly sparsity pattern is not compared to alternative sparsity patterns. Block-sparse FlashAttention uses a fixed butterfly sparsity from Dao et al. (2022). Would sliding window, random blocks, or learned sparsity patterns yield better accuracy-speed tradeoffs? The paper's IO-complexity analysis (Proposition 4) applies to any block-sparse pattern, so the infrastructure supports such comparisons, but the paper only demonstrates butterfly. This is a reasonable proof of concept, but it doesn't validate that butterfly is the best sparsity pattern for IO-aware attention — it only validates that IO-aware sparsity yields proportional speedups for whatever pattern is used.

• The Path-X and Path-256 results may not be reproducible without the exact tuning recipe. The paper describes a multi-stage procedure (pretrain, interpolate, fine-tune, additional fine-tune for Path-X) that "adds roughly 4 points of accuracy" for Path-X. LRA results are known to be sensitive to tuning, and the paper's ability to achieve 61.4% may depend on hyperparameters that are not fully specified. This is a limitation of the benchmark rather than the paper, but it tempers the strength of the "first Transformer to solve Path-X" claim — future work may find that alternative architectures also solve it with sufficient tuning, even without FlashAttention's efficiency.

Overall assessment: The experiments strongly support the paper's central claims about training speed and memory efficiency gains from IO-awareness. The 15% BERT improvement, 1.7–3.5× GPT-2 improvement, and 2.4–2.8× LRA improvement are measured against credible baselines on standard hardware. The Path-X/Path-256 results genuinely demonstrate new capability enabled by efficiency. The benchmarking in Appendix E is unusually comprehensive for a systems paper, covering multiple hardware platforms, configurations (with/without dropout and masking), and sequence lengths spanning three orders of magnitude. The primary limitation is not in what was measured but in what was not: the IO-complexity theory is validated through runtime correlation rather than direct memory traffic measurement, and the generalizability of the IO-awareness principle to non-attention operations remains a conjecture supported by a single blockbuster example.

6. Limitations and Trade-offs

The Difficulty Estimation Cost: FlashAttention Requires a CUDA Expert for Every New Attention Variant

The assumption or constraint. FlashAttention is implemented as a custom CUDA kernel — a low-level GPU program written in C++ with CUDA extensions. The paper makes no assumption that users will write their own CUDA kernels (the code is open-sourced), but it does assume that any new attention variant (different sparsity pattern, different scoring function, different normalization) requires a new CUDA kernel to be IO-aware. The paper explicitly acknowledges this in Section 5:

"Our current approach to building IO-aware implementations of attention requires writing a new CUDA kernel for each new attention implementation. This requires writing the attention algorithm in a considerably lower-level language than PyTorch, and requires significant engineering effort. Implementations may also not be transferrable across GPU architectures."

The consequence. The practical consequence is a deployment and experimentation tax. If a researcher wants to try a new sparse attention pattern or a new attention function (e.g., linearized attention, kernel attention, structured attention), they face a stark choice: (a) implement it in PyTorch (standard approach) and accept that it will likely be slower than dense FlashAttention even though it does fewer FLOPs — undermining the point of the new method; (b) write a new CUDA kernel implementing the FlashAttention tiling strategy for their variant, requiring "significant engineering effort" and GPU programming expertise that most ML researchers lack; or (c) wait for someone else to do (b). This bottleneck applies to every innovation in attention mechanisms: the paper's own block-sparse FlashAttention required Algorithm 5 (in Appendix D.1), which is a modified CUDA kernel. Each new sparsity pattern (sliding window, random blocks, locality-sensitive hashing, learned sparsity, strided sparsity) would require its own CUDA kernel or at least a generalized sparse FlashAttention kernel — which the paper provides for block-structured sparsity but not for unstructured or fine-grained sparsity patterns.

A subtler consequence: implementation correctness is hard to verify. CUDA kernels involve manual memory management, thread synchronization, and floating-point non-determinism. Bugs can produce results that are close to correct but numerically different — hard to detect but potentially catastrophic for training stability. The paper validates correctness through perplexity equivalence (Figure 4, Appendix E.2), but this validation applies only to the provided dense and butterfly-sparse kernels. Every new kernel requires re-validation.

What evidence exists in the paper. The paper provides no experiment measuring the engineering cost of writing a new IO-aware attention kernel, because this cost is person-time, not GPU-time. However, the block-sparse FlashAttention results (Table 3, Figure 2 right, Appendix D.1) demonstrate that extending FlashAttention to a new pattern required a modified algorithm (Algorithm 5) and a new CUDA kernel — confirming that even a conceptually simple extension (adding a sparsity mask check) requires non-trivial low-level work. The FMHA comparison (Table 7, Appendix E.4) illustrates the other side: FMHA was limited to head dimension 64, sequence length ≤512, and A100 GPUs — demonstrating the specificity of hand-tuned CUDA kernels to particular hardware configurations.

Mitigation status. The paper partially addresses this by open-sourcing the code (Section 1: "We open-source FlashAttention to make it easier to build on this primitive") and by calling for compiler-level solutions in Section 5: "These limitations suggest the need for a method that supports writing attention algorithms in a high-level language (e.g., PyTorch), and compiling to IO-aware implementations in CUDA—similar to efforts such as Halide in image processing." This is a future-work suggestion, not a mitigation provided in the paper. The open-source code mitigates the cost for users of the specific patterns implemented (dense, block-sparse butterfly), but does not reduce the cost for developers of new patterns. The Halide analogy is apt — Halide revolutionized image processing by separating the algorithm (what to compute) from the schedule (how to tile, vectorize, and parallelize), but no equivalent exists for GPU attention kernels at the time of writing. Until such a compiler exists, the IO-awareness principle remains difficult to apply broadly, and FlashAttention remains a collection of hand-optimized kernels rather than a general framework.

---

Single Hardware Family, Single Precision: Results Are Validated Only on Nvidia Ampere and Turing GPUs with FP16

The assumption or constraint. All experiments in the paper use Nvidia GPUs: A100 (Ampere architecture), RTX 3090 (Ampere), and T4 (Turing). All training experiments use FP16 mixed precision via PyTorch AMP. The paper states in Section 2.1 that "performance on other hardware accelerators are similar" and cites a single reference on Graphcore IPUs (Jia et al., 2019), but provides no experimental evidence on non-Nvidia hardware. The IO complexity analysis (Theorem 2) depends on the ratio $d^2 / M$ where $M$ is SRAM size — this ratio varies across GPU architectures (A100: 192 KB SRAM per SM; T4: smaller SRAM, hence "less speedup on T4" as shown in Appendix E.5) and would be entirely different on non-GPU accelerators (TPUs, IPUs, FPGAs, custom ASICs).

The consequence. Two distinct failure modes arise from this hardware specificity:

1. Non-transferability of speedup magnitudes. The paper demonstrates that speedup varies from ~2.5–4.5× on consumer RTX 3090 to 2–4× on datacenter A100 to lesser speedup on T4 (Appendix E.5). On hardware with different SRAM sizes, bandwidth ratios, or compute architectures, the speedup could be smaller (weaker SRAM) or larger (worse HBM bandwidth). A practitioner running on a different GPU generation or a non-Nvidia accelerator cannot use the paper's headline numbers as reliable estimates. More critically, FlashAttention's design — specifically the block size selection $B_c = \lfloor M / (4d) \rfloor$ — is parameterized by $M$, but the algorithm also relies on Nvidia-specific features (shared memory, warp-level primitives, CUDA's execution model) that may not have direct analogs on other hardware. Porting FlashAttention to AMD GPUs (ROCm), Intel GPUs (oneAPI), or Apple Silicon would require re-engineering the kernel for each platform's memory hierarchy and programming model. The "hardware lottery" concept the paper cites (Hooker, 2021) applies in reverse: FlashAttention won the Nvidia Ampere lottery, but its success on that platform does not guarantee success elsewhere.

2. Unknown behavior at higher precision. All experiments use FP16 mixed precision. For applications requiring FP32 or FP64 (e.g., scientific computing, some medical applications, lossless fine-tuning), the memory bandwidth bottleneck is more severe (larger tensors mean more bytes to move per operation), so FlashAttention's IO-awareness should provide larger relative speedups — but this is unverified. The block size $B_c = \lfloor M / (4d) \rfloor$ depends on the tensor element size (2 bytes for FP16, 4 bytes for FP32, 8 bytes for FP64). Running in FP32 would halve the effective $B_c$, increasing $T_c = N / B_c$ (the number of outer loop iterations) by 2×, which increases HBM accesses proportionally per Theorem 2. The speedup over standard attention would still exist (both FlashAttention and standard attention would see increased HBM traffic), but the crossover point where compute-bound behavior replaces memory-bound behavior would shift. The paper provides no measurements at FP32 or FP64.

What evidence exists in the paper. The T4 experiments (Appendix E.5, Figure 8) directly demonstrate hardware sensitivity: "T4 SRAM is smaller than A100, so we need to make the block sizes smaller in FlashAttention. As a result, we observe less speedup on T4, which matches the IO complexity analysis in Section 3.2." This is evidence that the $M$-dependence is real and that FlashAttention's performance degrades predictably on less-capable hardware. The RTX 3090 experiments (Figure 7) show the inverse: lower HBM bandwidth (900 GB/s vs. 1500 GB/s) produces higher speedup because the memory bottleneck is more severe. These two datapoints bracket the sensitivity range but remain within the Nvidia CUDA ecosystem. The paper's claim that "performance on other hardware accelerators are similar" is supported by one citation to Jia et al. (2019) for IPUs — a single data point that provides no evidence about FlashAttention specifically, only about the general principle that memory hierarchies exist on other accelerators.

Mitigation status. Not addressed beyond the future-work suggestion in Section 5: "Implementations may also not be transferrable across GPU architectures. These limitations suggest the need for a method that supports writing attention algorithms in a high-level language... and compiling to IO-aware implementations in CUDA." This future-work item implicitly acknowledges the portability problem. The open-source release (Section 1) provides Nvidia CUDA code that others could potentially port, but the engineering burden of porting falls on the community, not the authors. No abstraction layer or hardware-independent specification is provided.

---

Latency vs. Throughput: Sequential Outer Loop Requires Multiple Passes Over Q, Creating an Inherent Latency Floor for Interactive Use

The assumption or constraint. FlashAttention's outer loop over K-V blocks and inner loop over Q blocks (Algorithm 1, lines 5–14) is inherently sequential: each Q block must be loaded, processed, and written back $T_c$ times (once per K-V block), and each K-V block must be fully processed before moving to the next. The paper assumes a batch training setting (GPT-2 with batch size 512, BERT with batch size 448) where throughput (samples processed per second) is the primary metric and latency (time to compute attention for a single query) is amortized over the batch. The paper does not discuss or measure latency for low-batch or single-query inference scenarios.

The consequence. For interactive applications — chatbots, real-time translation, code completion, any setting where a user is waiting for a single response — FlashAttention's tiling strategy may introduce latency that is not reducible by adding more GPU parallelism. The forward pass requires loading all Q blocks $T_c$ times each. Even if the GPU has thousands of idle cores, the sequential outer loop means that Q block $i$ cannot be fully processed until all $T_c$ K-V blocks have been loaded and multiplied against it. In standard attention, the matrix multiplication $S = QK^T$ is a single large operation that can be parallelized across all GPU cores simultaneously, potentially completing in one wavefront. FlashAttention trades this single-pass parallelism for reduced HBM bandwidth: it makes $T_c$ smaller parallel operations (each $B_r \times B_c \times d$), which individually under-utilize the GPU but collectively reduce data movement.

For training with large batches, this is a net win because the batch dimension provides enough parallelism to keep the GPU saturated across the multiple passes (each pass processes all batch elements within a Q block simultaneously). For batch size 1 inference, the parallelism within each $B_r \times B_c$ block is limited — $B_r$ is typically small (constrained by SRAM, often 64–128) — meaning each inner loop iteration uses only a fraction of the GPU's compute units. The multiple passes then become a latency tax: the attention computation might use less total GPU time (fewer HBM accesses), but the wall-clock latency experienced by the user could be worse because the work is spread over many sequential kernel launches and SRAM loads.

The paper provides no latency measurements at batch size 1. The closest analog is the T4 forward-pass benchmarking (Figure 8 bottom, Appendix E.5), which shows that FlashAttention's forward-pass speedup is smaller at short sequences on the T4 — but this is for batch size 12 with 12 heads, not a low-batch inference scenario. The speedup numbers for training (15% on BERT, 1.7–3.5× on GPT-2) are measured at batch sizes 448–512 and reflect throughput, not per-sample latency.

What evidence exists in the paper. No direct evidence. The paper's experiments are exclusively training-focused or batch inference (attention benchmarking uses batch size 16 with 8 heads; LRA uses standard training batch sizes; GPT-2 uses effective batch size 512). The IO complexity analysis (Theorem 2) counts total HBM accesses but does not model latency — it assumes all HBM accesses are equally costly regardless of whether they are interleaved with compute or serialized. The closest informative result is Figure 2 (middle), which shows that runtime decreases as block size increases (fewer outer loop iterations → less latency), and then plateaus. The plateau suggests that at very large block sizes (approaching full-matrix single-pass attention), the latency floor is determined by compute rather than memory — but those block sizes exceed SRAM capacity, so they are unreachable in practice.

Mitigation status. Not addressed. The paper does not discuss the latency-versus-throughput tradeoff, single-query inference, or interactive use cases. This is a reasonable scope limitation for a training-focused paper — FlashAttention's primary target is training, where throughput dominates — but a practitioner deploying FlashAttention for inference would need to benchmark this independently. The potential mitigation (not explored in the paper) would be a different tiling strategy optimized for latency: for example, loading one Q block and streaming all K-V blocks through SRAM in a single fused kernel, rather than iterating over Q blocks inside the outer K-V loop. Whether such a strategy would maintain the IO-awareness benefits while reducing latency is an open question.

---

The Backward Pass Recomputation Scrutiny: 13% More FLOPs and the Absence of an End-to-End Forward-Only Inference Baseline

The assumption or constraint. FlashAttention's backward pass recomputes the attention probabilities P from stored softmax statistics ($m$, $\ell$) during the backward pass (Algorithm 4), rather than reading P from HBM. This trades increased FLOPs (75.2 GFLOPs vs. 66.6 GFLOPs for the standard implementation in Figure 2 left, a ~13% increase) for reduced HBM accesses (4.4 GB vs. 40.3 GB). The paper assumes training workloads (where the backward pass exists and is the primary bottleneck) and asserts that "even with more FLOPs, our recomputation speeds up the backward pass due to reduced HBM accesses" (Section 3.1). The paper does not separately benchmark the forward pass in an inference-only setting against standard attention forward.

The consequence. Two issues arise:

1. The speedup composition is opaque. The training speedup numbers (15% BERT, 1.7–3.5× GPT-2) aggregate forward and backward pass improvements. A practitioner evaluating FlashAttention for training might want to know: is the speedup concentrated in the backward pass (where recomputation avoids reading P), the forward pass (where tiling avoids writing P), or both? The paper's Figure 2 (left) reports combined forward+backward runtime (7.3 ms vs. 41.7 ms) and combined HBM accesses (4.4 GB vs. 40.3 GB), but does not decompose into forward-only and backward-only contributions. Table 7 (Appendix E.4) provides this decomposition for the FMHA comparison: at sequence length 512, FlashAttention forward is 0.81 ms vs. 1.14 ms for FMHA (29% faster), while FlashAttention backward is 2.00 ms vs. 1.81 ms for FMHA (10% slower). This reveals that at short sequences, the recomputation cost in the backward pass partially offsets the forward-pass gains — a nuance invisible in the combined runtime.

2. Inference workloads lose the recomputation advantage entirely. During inference (forward pass only), FlashAttention still benefits from tiling and reduced HBM writes (no P is written to HBM), but does NOT benefit from the recomputation strategy (there is no backward pass to recompute for). The standard attention forward pass also does not need to read P during forward — it only writes P to HBM. So the inference speedup of FlashAttention over standard attention comes purely from avoiding writing the $N \times N$ matrix P to HBM, plus kernel fusion of masking, softmax, and dropout. This is a real but more modest gain. The paper reports combined forward+backward speedups that mix the large backward-pass improvement (recomputation-avoids-reading-P) with the forward-pass improvement (tiling-avoids-writing-P). For pure inference, the speedup would be the forward-pass component only.

The Appendix E.5 T4 forward-only benchmarking (Figure 8 bottom) partially addresses this: it shows FlashAttention's forward-pass speedup over PyTorch attention on T4 for various sequence lengths and configurations. However, this is on a T4 (where FlashAttention is known to be less effective due to smaller SRAM) and is not compared against the same paper's training speedups, making it difficult to assess how much of the headline 2–3× training speedup would translate to inference.

What evidence exists in the paper. The FMHA comparison (Table 7) provides the most relevant decomposition. At sequence length 512 with BERT-size configuration (batch size 64, 16 heads, head dimension 64), FlashAttention forward is 0.81 ms vs. 1.14 ms for FMHA (1.4× faster), while backward is 2.00 ms vs. 1.81 ms for FMHA (1.1× slower). Combined: 2.81 ms vs. 2.95 ms (1.05× faster). This single datapoint suggests that at this sequence length and configuration, FlashAttention's training speedup comes almost entirely from the forward pass — the backward pass is actually slower due to recomputation. The paper does not discuss this or provide similar decomposition for the GPT-2 or LRA experiments. The Figure 2 (left) GPT-2 medium configuration shows 75.2 vs. 66.6 GFLOPs (13% increase) and 7.3 ms vs. 41.7 ms (5.7× faster) combined — but since the backward pass produces both the FLOP increase and (potentially) a disproportionate share of the speedup, the forward-only speedup cannot be inferred from the combined number.

Mitigation status. The paper acknowledges the FLOP increase explicitly in Figure 2 (left) and Section 3.1, but frames it as a feature ("even with more FLOPs, our recomputation speeds up the backward pass"). It does not separately benchmark forward-only inference speedups on training-class GPUs (A100, RTX 3090) for the configurations used in the headline results (GPT-2, BERT, LRA). The T4 forward-only numbers in Appendix E.5 are the closest to addressing this, but they are on a different GPU class and are not integrated into the main narrative. A practitioner deploying FlashAttention for inference would need to run their own benchmarks to determine the actual speedup in their setting.

---

Hard Problems Remain Quadratic: FlashAttention Does Not Change the Asymptotic Time Complexity of Exact Attention

The assumption or constraint. FlashAttention computes exact attention — the same $O(N^2 d)$ FLOPs as standard attention. The paper is explicit about this: Theorem 1 states "Algorithm 1 returns $O = \text{softmax}(QK^T)V$ with $O(N^2 d)$ FLOPs." The speedup comes from reducing the constant factor on memory accesses by $O(d^2 / M)$, not from reducing the asymptotic time complexity. The paper assumes that sequence lengths of practical interest (up to ~4K–8K for dense FlashAttention, up to 64K for block-sparse) are within the range where the reduced constant factor makes the computation feasible but the $O(N^2)$ asymptotic still governs.

The consequence. FlashAttention hits a hard scaling wall at very long sequences, just like standard attention — it just hits it later. For a given GPU with SRAM size $M$, the runtime of FlashAttention forward pass is $O(N^2 d)$ FLOPs executed at a certain throughput, with $O(N^2 d^2 / M)$ HBM accesses. As $N$ grows large, the $N^2$ term dominates regardless of the constant factor improvements. The paper's benchmarking (Figure 3 left, Table 11) demonstrates this: FlashAttention's combined forward+backward runtime grows from 9.56 ms at 2048 to 147.75 ms at 8192 to 9341.30 ms at 65536 — a ~977× increase for 32× sequence length growth, consistent with ~$N^2$ scaling. At 64K, even on an A100, dense FlashAttention takes ~9.3 seconds per attention operation (forward+backward, batch size 16, 8 heads). For training a large model with many attention layers over many steps, this becomes prohibitive.

The block-sparse extension (Proposition 4) improves the asymptotic to $O(N^2 d^2 M^{-1} s)$ where $s$ is the sparsity fraction, and the paper demonstrates $s = O(N^{-1/2})$ butterfly sparsity reducing the scaling to $O(N^{3/2})$ — which does make 64K tractable (block-sparse FlashAttention: 64.11 ms at 64K vs. 9341.30 ms for dense, a 145× speedup). However, this sacrifices exactness: block-sparse attention is an approximation. The paper's key innovation — exact attention with IO-awareness — does not escape quadratic complexity.

This means FlashAttention does not obsolete approximate attention methods for very long sequences. The paper's own benchmarking (Figure 3 left) shows that block-sparse FlashAttention is faster than dense FlashAttention and all approximate baselines, but this is achieved by also being an approximate method (block-sparse). For practitioners who need exact attention on sequences longer than ~16K–64K (depending on GPU memory and patience), FlashAttention offers no asymptotic escape — the $N^2$ scaling eventually catches up, and the only options are sparsification (trading exactness for speed) or alternative architectures (state-space models, linear attention, etc.).

What evidence exists in the paper. The evidence is in the runtime scaling numbers. From Table 11 (combined forward+backward, with dropout and masking, A100):

Dense FlashAttention grows roughly 3.7× when sequence length doubles (from 9.56 to 37.49, then 147.75, then 586.61 — each ~3.9×, ~4.0×, ~4.0×, close to the $2^2 = 4\times$ quadratic growth). Block-sparse FlashAttention grows roughly 2.2× per doubling (from 0.89 to 1.95 to 4.12 to 7.64 to 16.60), closer to the $O(N^{3/2})$ scaling from $s = O(N^{-1/2})$. The gap between the two methods widens dramatically with sequence length, confirming that dense FlashAttention remains bound by quadratic complexity.

The paper also acknowledges this limitation implicitly by introducing block-sparse FlashAttention as a necessary extension for 64K sequences — dense FlashAttention at 64K takes 9.3 seconds per operation, which is technically "feasible" but not practical for training. The Path-X result (61.4% at 16K) uses dense FlashAttention; the Path-256 result (63.1% at 64K) uses block-sparse. This structure communicates that exact attention reaches its practical limit around 16K–32K, and beyond that approximation is required.

Mitigation status. The paper does not frame this as a limitation — it is a mathematical fact about exact attention that no algorithm can escape. Proposition 3 proves the lower bound on HBM accesses for exact attention, and the $O(N^2 d)$ FLOP count is inherent to computing all pairwise dot products. The paper's contribution is making the constant factor as small as possible (IO-optimal within constants), not changing the asymptotic. The block-sparse extension is the mitigation for practitioners who need longer sequences and can accept approximation — but it is a different algorithm (approximate, not exact) and the paper presents it as a separate contribution. A reader looking for "exact attention at 100K sequence lengths" will find that FlashAttention cannot deliver this — the $N^2$ scaling ensures that 100K is ~23× more expensive than 16K for dense attention, and even the IO-optimal constant factor cannot overcome that multiplier.

---

The Generalization Gap: Only MATH Equivalents, Only Autoregressive and Encoder Transformers

The assumption or constraint. All experiments are on Transformer models with standard scaled dot-product self-attention (either bidirectional for BERT and LRA, or causal for GPT-2). The paper assumes that attention is the bottleneck and that other Transformer components (feed-forward layers, layer norm, embeddings) are not the limiting factor. The experiments cover three model types (BERT-large encoder, GPT-2 decoder, vanilla Transformer for LRA) and five task families (masked language modeling, autoregressive language modeling, long-document classification, LRA classification, Path-X/Path-256), all using the standard attention formulation $O = \text{softmax}(QK^T/\sqrt{d})V$.

The consequence. FlashAttention's speedup is not guaranteed to transfer to architectural variants that modify the attention computation. These include:

1. Cross-attention (encoder-decoder attention in seq2seq Transformers like T5, BART, or translation models), where Q comes from the decoder and K, V come from the encoder. The memory access pattern differs — K and V are encoder outputs that may be of different sequence length than Q, and the tiling strategy might need different block sizes or loop ordering to be optimal. The paper does not experiment with cross-attention.

2. Multi-query attention (MQA) or grouped-query attention (GQA), where multiple query heads share a single key-value head. These variants are increasingly common in deployed models (PaLM, LLaMA 2, Gemini, etc.) because they reduce the KV cache size for inference. The memory access pattern and arithmetic intensity change significantly — with MQA, the K and V matrices are smaller, changing the balance between the $Nd$ and $N^2 d$ terms in the IO complexity. FlashAttention's block size selection $B_c = \lfloor M / (4d) \rfloor$ assumes the standard head dimension 64–128 for K and V — it might be suboptimal for shared KV heads.

3. Attention variants with different scoring functions. FlashAttention's tiling strategy depends on decomposing softmax using the running maximum and sum statistics ($m$, $\ell$). If the attention function uses a different normalization (e.g., linear attention's $\phi(Q)\phi(K)^T$, or sigmoid attention, or learned similarity functions), the decomposition technique may not apply — or would require a different set of running statistics. The paper's algorithm is specifically tied to the exponential softmax form.

4. Vision Transformers (ViT) and other modalities. While LRA includes an Image task (sequence length 1024, presumably flattened image patches), the paper does not test FlashAttention on standard ViT benchmarks (ImageNet classification, object detection, segmentation) where the sequence length is often shorter (196–577 patches for typical 224×224–384×384 images) and the computational bottleneck may be the feed-forward layers rather than attention. The paper's speedup claims may not extrapolate to these settings.

5. Sparse mixture-of-experts (MoE) Transformers, where different tokens are routed to different feed-forward experts. The attention computation is the same, but the overall training time bottleneck shifts — if feed-forward expert computation dominates, FlashAttention's speedup on the attention sub-component translates to a smaller end-to-end speedup. The paper's GPT-2 experiments suggest 1.7–3.5× end-to-end speedup, but these use dense models where attention is a large fraction of total compute. For MoE models with many experts, the relative benefit would be smaller.

What evidence exists in the paper. The paper provides no experiments on cross-attention, MQA/GQA, non-softmax attention variants, or Vision Transformers. The closest evidence of generality is the diversity of tasks: MLM (BERT), autoregressive LM (GPT-2), classification (LRA, MIMIC-III, ECtHR), and synthetic reasoning (Path-X, Path-256). This covers multiple output types (token-level, sequence-level) and multiple sequence lengths (128 to 64K), but all within the standard Transformer architecture with standard dot-product softmax self-attention. The LRA benchmark includes an Image task, but it treats images as sequences of flattened pixels (not patches as in ViT), so it tests sequence length scaling rather than vision-specific attention patterns.

Mitigation status. The paper does not claim generality to other attention variants — the scope is explicitly "exact attention" as defined in Section 2.2. The title and abstract make no mention of cross-attention, MQA, or vision models. This is a reasonable scope limitation for an initial paper introducing a new primitive. However, the paper's framing in Section 5 — "We believe that the IO-aware approach can extend beyond attention... every layer in a deep network touches GPU HBM. We hope our work inspires IO-aware implementations of additional modules" — suggests broader ambitions. A practitioner using a non-standard attention variant or a non-Transformer architecture cannot rely on FlashAttention as a drop-in solution; they would need to adapt the tiling and recomputation strategy to their specific computation graph, which may or may not admit the same decomposition.

The block-sparse extension provides one data point on extensibility: it shows that the FlashAttention framework can accommodate a different computation pattern (skipping blocks) while maintaining the IO-awareness benefits. But block-sparse attention is still dot-product softmax self-attention — it changes which scores are computed, not how scores are computed. Extending to cross-attention or non-softmax scoring functions would test the framework's generality more severely, and the paper provides no evidence either way.

7. Implications and Future Directions

How This Work Changes the Landscape

FlashAttention represents a methodological reframing rather than a paradigm shift — it does not change what attention computes or how Transformers work, but it fundamentally changes how the field should think about optimizing deep learning operations. The paper's lasting contribution is establishing IO-awareness as a first-class design criterion for neural network primitives, displacing the field's default assumption that FLOP reduction is the primary path to wall-clock speed.

This reframing has several concrete consequences for how research is conducted and evaluated:

1. FLOP-counting is exposed as an unreliable proxy for runtime. The paper provides the clearest empirical demonstration to date in the attention literature that FLOPs and wall-clock time can move in opposite directions. Figure 2 (left) shows FlashAttention's forward+backward pass executing 13% more FLOPs (75.2 vs. 66.6 GFLOPs) while running 5.7× faster (7.3 ms vs. 41.7 ms). This single datapoint invalidates the common practice of reporting FLOP reductions as evidence of efficiency — a practice ubiquitous in the efficient Transformers literature. The implication is that future papers proposing attention variants must report wall-clock runtime and memory usage on standard hardware, not just asymptotic FLOP counts. Conferences and reviewers now have a clear precedent for requiring these metrics.

2. It reconciles the contradictory results between approximate attention methods that reduce FLOPs and their failure to achieve proportional speedups. The paper identifies the missing variable: memory access overhead. Prior work like Reformer, Linformer, Performer, and BigBird all achieved asymptotic FLOP reductions (typically from $O(N^2)$ to $O(N)$ or $O(N \log N)$), but their wall-clock speedups were inconsistent and often disappointing. FlashAttention's IO-complexity analysis (Theorem 2) provides the explanation: these methods reduced computation but not data movement, or introduced irregular memory access patterns (hashing, gather/scatter, sparse indexing) that themselves became the bottleneck. The proof that standard attention's IO complexity is $\Theta(Nd + N^2)$ — dominated by reading and writing the $N \times N$ attention matrix — reveals that FLOP-focused methods were optimizing the wrong cost model. This reconciliation is valuable because it converts a confusing empirical landscape ("why doesn't this $O(N)$ method run $N\times$ faster?") into a principled diagnosis.

3. It elevates recomputation from a memory-saving technique to a speed-enhancing one. Gradient checkpointing (Chen et al., 2016) and the memory-efficient attention of Rabe and Staats (2021) both used recomputation to reduce peak memory, but traded speed for memory — the backward pass became slower because it recomputed values it could have read from memory. FlashAttention inverts this logic: on memory-bound operations, recomputation can be faster than storage because the cost of recomputing from on-chip data is lower than the cost of reading from off-chip memory. The paper's backward pass (Algorithm 4) recomputes $P_{ij} = \text{diag}(\ell_i)^{-1} \exp(S_{ij} - m_i)$ on-chip, converting what would be $O(N^2)$ HBM reads into $O(N^2 d^2 M^{-1})$ HBM reads plus extra FLOPs — and comes out ahead. This insight generalizes: any memory-bound layer where intermediate activations can be reconstructed from smaller stored statistics and cheap on-chip compute is a candidate for the same treatment (layer norm statistics, activation function inputs, certain pooling operations). The paper's explicit call for IO-aware implementations of "additional modules" (Section 5) points toward this generalization.

4. It establishes a theoretical ceiling for exact attention optimization. Proposition 3's lower bound — proving that no exact attention algorithm can achieve $o(N^2 d^2 M^{-1})$ HBM accesses for all SRAM sizes — provides a principled stopping criterion for future engineering. FlashAttention achieves this bound within constants, meaning further asymptotic improvements to exact attention are impossible without hardware changes (larger SRAM, higher bandwidth, different memory hierarchy). This channels research effort into two productive directions: (a) hardware-algorithm co-design (designing GPUs or accelerators with SRAM sizes chosen to minimize attention's IO complexity), and (b) acceptable approximation (block-sparse FlashAttention already demonstrates that sparsification breaks through the exact-attention lower bound, trading accuracy for IO complexity proportional to the sparsity ratio). This is analogous to how Chinchilla scaling laws (Hoffmann et al., 2022) provided a compute-optimal pretraining recipe that structured subsequent research — FlashAttention provides an IO-optimal attention recipe that structures subsequent optimization work.

5. It changes the narrative around hardware-dependent research. FlashAttention is unapologetically hardware-specific — the algorithm's block sizes, loop ordering, and recomputation strategy are all parameterized by GPU SRAM size, HBM bandwidth, and compute throughput. The paper demonstrates that hardware-specific optimization yields gains that hardware-agnostic methods cannot match (the block-sparse FlashAttention results in Figure 2 right and Table 3 show speedups proportional to sparsity, something no prior sparse attention method achieved). This challenges the field's preference for hardware-agnostic algorithmic contributions and suggests that the "hardware lottery" (Hooker, 2021) can be engaged with productively — by designing algorithms that explicitly account for specific hardware characteristics — rather than only lamented. The paper's call for compilers that automate this hardware specialization (Section 5, Halide analogy) points toward a middle ground where researchers express algorithms at a high level and compilers handle IO-aware code generation for each target platform. This reframes the relationship between ML research and systems engineering: IO-awareness is not an implementation detail but a core algorithmic property that should influence how operations are designed.

Follow-Up Research This Work Enables

Compiling attention to IO-aware CUDA kernels from high-level specifications. The paper's most explicit future-work call is Section 5's suggestion of "a method that supports writing attention algorithms in a high-level language (e.g., PyTorch), and compiling to IO-aware implementations in CUDA—similar to efforts such as Halide in image processing." This is a systems research problem: given a declarative specification of an attention variant (e.g., "compute $O = \text{softmax}(\text{mask}(QK^T/\sqrt{d}))V$ with this sparsity pattern"), automatically generate a tiled CUDA kernel with optimal block sizes, loop ordering, and recomputation strategy for the target GPU. A strong result would demonstrate that auto-generated kernels for a range of attention variants (sliding window, causal, block-sparse, cross-attention, linear attention) achieve ≥90% of the speed of hand-optimized FlashAttention, on at least two GPU architectures (A100 and H100, or Nvidia and AMD). The existence proof is FlashAttention itself — it shows the target performance level is achievable — and the research contribution would be the compiler infrastructure that makes it accessible without CUDA expertise. This would directly address the limitation that "implementations may also not be transferrable across GPU architectures" (Section 5).

Rigorous characterization of the compute-bound to memory-bound transition for FlashAttention across hardware generations. Figure 2 (middle) shows that as block size increases, FlashAttention runtime decreases until it plateaus at block size ~256, where the operation becomes compute-bound. This crossover point depends on SRAM size, HBM bandwidth, and compute throughput — all of which change across GPU generations (A100 → H100 → future architectures). A systematic study measuring the IO complexity and runtime of FlashAttention on at least three GPU generations (V100, A100, H100), while varying head dimension (32, 64, 128, 256), sequence length (512–32K), and precision (FP16, BF16, FP32), would produce an empirical scaling law for attention's memory behavior. The specific question: how does the optimal block size $B_c = \lfloor M / (4d) \rfloor$ from the paper's analysis evolve as SRAM grows ($M$ increases) and as head dimensions increase (larger $d$ reduces $B_c$)? Extending this to model the backward pass separately — where recomputation changes the arithmetic intensity — would provide practitioners with a lookup table for configuring FlashAttention optimally on their specific hardware.

IO-aware cross-attention and encoder-decoder models. The paper evaluates only self-attention (BERT, GPT-2) and does not benchmark cross-attention, where Q comes from a decoder and K, V from an encoder. In encoder-decoder models (T5, BART, Whisper, translation models), cross-attention has a different memory access pattern: K and V are encoder outputs that are typically of different sequence length than Q, and are reused across all decoder layers. The IO-optimal tiling strategy may differ — perhaps the outer loop should be over Q blocks (decoder) and the inner loop over K-V blocks (encoder), reversing the paper's Algorithm 1 ordering, because K and V are accessed by every decoder layer and maximizing their reuse in SRAM could be more important than maximizing Q reuse. A strong follow-up would implement an IO-aware cross-attention kernel (or verify that the existing FlashAttention works well for cross-attention), benchmark it on a standard seq2seq task (e.g., WMT translation with T5-base), and measure both training speedup and the effect on inference latency for beam search decoding, where the KV cache management differs from self-attention.

IO-aware attention for multi-query and grouped-query attention (MQA/GQA). Modern deployed LLMs (LLaMA 2, Mistral, Gemini, PaLM) increasingly use multi-query attention (multiple query heads share a single key-value head) or grouped-query attention (a small number of KV heads shared across multiple query heads). These variants dramatically reduce the KV cache size for inference, but they also change the attention computation's arithmetic intensity: the K and V matrices have fewer heads, so reading them from HBM costs less bandwidth relative to the computation they drive. FlashAttention's block size selection $B_c = \lfloor M / (4d) \rfloor$ assumes all heads have independent K, V — with shared KV heads, the effective head dimension for K/V loading is smaller, potentially allowing larger $B_c$ and fewer outer loop iterations. A concrete experiment: benchmark FlashAttention on LLaMA-2-7B's attention (which uses GQA with 32 query heads and 8 KV heads), measure whether modifying the block sizes to account for KV head sharing yields additional speedup, and characterize how the speedup scales with the query-to-KV head ratio (from 1:1 for standard MHA to 32:1 for extreme MQA). This would determine whether FlashAttention's IO-optimal strategy generalizes to the attention variant used in most deployed LLMs.

Stress-testing the IO lower bound with near-SRAM-capacity sequence lengths. Proposition 3 proves that no exact attention algorithm can beat $\Theta(N^2 d^2 M^{-1})$ HBM accesses for all $M$, but the proof uses the $M = \Theta(Nd)$ regime (SRAM large enough to hold a full row/column). What about intermediate regimes — e.g., $M = \Theta(N\sqrt{d})$ or $M = \Theta(N)$ — that might arise in practice with very long sequences or small SRAM? A more refined lower bound analysis as a function of the $M/Nd$ ratio could reveal whether FlashAttention's specific tiling strategy (outer loop over K-V, inner loop over Q) is optimal for all ratios, or whether alternative loop orderings or block shapes (non-square attention blocks) could improve the constant factors. This is the "parameterized complexity" question the paper explicitly leaves for future work (Section 3.2). A strong negative result — finding an alternative exact attention algorithm that beats FlashAttention's constant factors for some $(M, N, d)$ combinations — would refine the optimality claims and potentially inspire new implementations. A strong positive result — proving FlashAttention is optimal for all parameter regimes — would close the exact-attention optimization problem definitively.

Systematic evaluation of IO-awareness for non-attention deep learning primitives. The paper's central claim is that IO-awareness generalizes beyond attention ("every layer in a deep network touches GPU HBM," Section 5), but it provides no evidence. A high-impact follow-up would identify the next most IO-intensive operations in Transformer training — layer normalization, cross-entropy loss, activation functions (GELU, SiLU), positional embeddings, feed-forward layers — and quantify their arithmetic intensity and memory access patterns using GPU profilers. For operations that prove to be memory-bound (the paper hypothesizes elementwise and reduction operations are), implement IO-aware fused kernels following the FlashAttention template (tiling to fit in SRAM, recomputation instead of storage for intermediate values during backprop). A rigorous study would report: (a) the fraction of total training time spent in each operation for a representative model (e.g., LLaMA-7B at sequence length 4K), (b) which operations are memory-bound vs. compute-bound, and (c) the speedup achieved by IO-aware implementations of the top memory-bound operations. This would convert the paper's speculative claim ("we hope our work inspires") into an actionable roadmap.

Practical Applications and Downstream Use Cases

Training LLMs on long documents without model parallelism for attention. The paper's headline result — training GPT-2 with context length 4K on 8×A100 GPUs in 3.6 days, 30% faster than Megatron with context length 1K — implies that many long-context training workloads that previously required tensor parallelism (splitting attention across multiple GPUs) can now fit on a single GPU or data-parallel setup. For organizations fine-tuning LLMs on long legal documents (as in the ECtHR experiment, where sequence length 8K yielded 8.5 F1 points of lift), medical records (MIMIC-III, 4.3 points of lift at 16K), or code repositories (where context windows of 8K–32K tokens capture entire files or modules), FlashAttention's linear memory footprint and 2–4× training speedup directly reduce the GPU-hours required. The concrete deployment scenario: a team fine-tuning LLaMA-2-7B on a corpus of 50K–100K scientific papers (average length 8K tokens) can use FlashAttention to train at sequence length 8K on 8×A100 GPUs without model parallelism, avoiding the communication overhead and engineering complexity of tensor-parallel training. The paper's memory benchmarking (Table 21: 836 MB at 4K, 3344 MB at 16K, 13376 MB at 64K for 8 heads) provides the memory budget for planning such deployments.

Enabling large-batch inference for retrieval-augmented generation (RAG) and document processing pipelines. Inference workloads that process many documents simultaneously — RAG retrieval scoring, document clustering, large-scale text classification — batch multiple sequences through the same attention operation. FlashAttention's throughput advantages are largest at moderate batch sizes (the paper uses batch size 16–64 for benchmarking, and training batches of 448–512). For a document processing pipeline scoring 10,000 passages against a query using cross-attention, batching 64 passages at a time with FlashAttention's 2–3× attention speedup directly reduces latency and cost. The Path-X and Path-256 results demonstrate that FlashAttention makes 16K–64K sequence lengths feasible — for applications like processing entire legal contracts, scientific papers, or code files in a single forward pass (rather than chunking them and losing cross-chunk context), FlashAttention's memory efficiency at extreme lengths (13.4 GB at 64K for the tested configuration) is what makes single-pass processing possible on a 40 GB A100.

Democratizing long-context research for academic and small-industry labs. Before FlashAttention, training a Transformer on sequence length 4K required either (a) model parallelism across multiple GPUs (engineering complexity), (b) approximate attention methods that traded accuracy for memory (Linformer, Performer), or (c) gradient accumulation with tiny per-GPU batch sizes (extremely slow). FlashAttention's open-source release and drop-in compatibility with PyTorch mean a researcher with access to a single 8×A100 node — the standard academic GPU cluster configuration — can train models at sequence lengths that were previously accessible only to industry labs with large GPU fleets. The paper's GPT-2 experiments are a direct demonstration: FlashAttention trains GPT-2 small at context 4K on 8×A100s in 3.6 days, producing a model with 0.7 better perplexity than a 1K-context model. This is a turnaround time that an academic lab can afford for a single experiment, enabling research on long-context language modeling, in-context learning with many examples, and retrieval-free document understanding without requiring industry-scale compute budgets.

Real-time or near-real-time applications requiring long context. The T4 forward-pass benchmarking (Figure 8 bottom, Appendix E.5) shows FlashAttention achieving speedups over PyTorch attention at inference time, though the magnitude is smaller than for training. For applications like real-time speech recognition with long audio context (e.g., transcribing hour-long meetings with Whisper-style encoder-decoder models) or interactive document Q&A (where a user's question is attended against a 20K-token document), FlashAttention's forward-pass speedup reduces the latency of each attention operation. The paper does not provide A100 inference-only benchmarks, but the T4 data suggests speedups of 1.5–3× for forward-pass at sequence lengths 1K–8K (depending on masking and dropout configuration). Combined with the memory savings (which allow fitting the entire long sequence in GPU memory without chunking), this enables lower-latency responses for applications where the user waits for the model to process a long context — a growing use case as LLMs are deployed for document-grounded conversation and codebase understanding.