Flash Attention: A Complete Architectural Autopsy

Prologue: The Problem Is Not Compute

Before a single line of Flash Attention code is read, the reader must internalize one fact that the textbook never states: the attention mechanism as originally written is not bottlenecked by arithmetic. It is bottlenecked by memory bandwidth, and specifically by the compulsive materialization of a matrix that the mathematics never required to exist.

The H100’s FP16 tensor core throughput sits at 989 TFLOPS. Its HBM3 bandwidth is 3.35 TB/s. The roofline ridge — the arithmetic intensity at which a kernel transitions from memory-bound to compute-bound — sits at roughly 295 FLOPS per byte. Standard attention, as you will see, lands around 64 FLOPS per byte. It never gets close to the compute ridge. Every optimization in the Flash Attention lineage is, at its root, a strategy to raise that number by eliminating unnecessary HBM round-trips.

Part 1: Standard Attention and the O(N²) Memory Crime

The reader knows the math. Given query, key, and value matrices Q, K, V ∈ ℝ^{N×d}, scaled dot-product attention computes:

S = QKᵀ / √d          (scores)        [N×N]
P = softmax(S)         (probabilities) [N×N]
O = PV                 (output)        [N×d]

The crime is in the brackets. S and P are N×N matrices. The standard PyTorch implementation writes S to HBM, reads it back to compute softmax into P, writes P to HBM, reads it back to compute PV. Four HBM transactions on a matrix whose size scales quadratically with sequence length.

The arithmetic: for N = 8192 and d = 128 in FP16, S alone is 8192 × 8192 × 2 bytes = 128 MB per attention head. A 70B model with 64 heads at this sequence length is materializing 8 GB of intermediate attention matrices per forward pass — matrices that do not survive the backward pass and serve no purpose except as staging areas for the softmax computation.

Formal complexity: O(N²d) FLOPs, O(N²) HBM memory. The FLOPs are unavoidable. The memory is not.

Arithmetic intensity of standard attention:

• FLOPs: 4N²d (two N×d-by-d×N matmuls for QKᵀ and PV)
• HBM bytes: dominated by the N×N round-trip — writing S, reading S, writing P, reading P = 4 × 2N² bytes in FP16
• I = 4N²d / 8N² = d/2 FLOPS/byte

For d = 128: I = 64 FLOPS/byte. The H100 is overprovisioned for compute by a factor of ~4.6×. The kernel spends most of its time waiting for HBM, not computing.

Part 2: Flash Attention 1 — Tile, Pipeline, Never Materialize

Dao, Fu, Ermon, Rudra, and Ré published FlashAttention in May 2022 (arXiv:2205.14135). The contribution is a rearrangement of where the computation happens, not a change to what is computed. The output is bit-identical to standard attention. The HBM traffic is not.

The Core Insight: SRAM Is Fast Enough for the Entire Working Set

The H100’s per-SM SRAM budget is 228 KB (at maximum carveout). The question FA1 answered is: can we structure the attention computation so that the N×N attention matrix is never written to HBM — not even in tiles — and all intermediate values live in SRAM?

The answer is yes, with one non-trivial ingredient: a numerically stable algorithm for computing softmax in a single pass over the keys, without ever knowing the global maximum of the row in advance. That ingredient is the online softmax of Milakov and Gimelshein (2018, Online normalizer calculation for softmax).

The Online Softmax Recurrence

Recall that standard softmax for row i is:

P[i,j] = exp(S[i,j] - max_j S[i,j]) / Σ_j exp(S[i,j] - max_j S[i,j])

The division by the global sum requires two passes: one to find the max, one to accumulate the sum. Neither pass can proceed until S[i,:] is fully materialized.

Online softmax restructures this as a recurrence over key-blocks. Maintain two scalars per query token: a running maximum m and a running sum of exponentials ℓ. For each new block of key scores s_new:

m_new = max(m_old, max(s_new))
ℓ_new = exp(m_old - m_new) · ℓ_old + Σ_j exp(s_new[j] - m_new)

When the running maximum increases (m_new > m_old), all previously accumulated exponentials must be rescaled by exp(m_old - m_new) — a correction factor that is always ≤ 1. The output accumulator O must be similarly rescaled:

O_new = exp(m_old - m_new) · O_old + exp(s_new - m_new) · V_block

At the end of all key-blocks, the true softmax output is O / ℓ. The complete derivation:

Let block index t iterate over Bc-sized key chunks. Initialize m⁰ = -∞, ℓ⁰ = 0, O⁰ = 0.

For each block t, load K_t, V_t from HBM into SRAM. Compute S_t = Q · K_tᵀ ∈ ℝ^{Br×Bc}. Then:

m̃_t = rowmax(S_t)                           # local max, shape [Br]
m_t = max(m_{t-1}, m̃_t)                     # updated running max
P̃_t = exp(S_t - m_t)                        # softmax numerators, unnormalized
ℓ_t = exp(m_{t-1} - m_t) ⊙ ℓ_{t-1} + rowsum(P̃_t)
O_t = diag(exp(m_{t-1} - m_t)) · O_{t-1} + P̃_t · V_t

After all T_c blocks: O = diag(ℓ_{T_c})⁻¹ · O_{T_c}

This is the entire recurrence. S_t and P̃_t live in SRAM for the duration of a single block and are discarded. The N×N attention matrix is never written to HBM. The m and ℓ statistics occupy O(N) space — one scalar per query token — and are stored for the backward pass.

Tiling and SRAM Budget

For d = 128, FP16, with Br = Bc = 64:

• Q tile: 64 × 128 × 2 = 16 KB
• K tile: 64 × 128 × 2 = 16 KB
• V tile: 64 × 128 × 2 = 16 KB
• O accumulator (FP32): 64 × 128 × 4 = 32 KB
• S_t tile: 64 × 64 × 4 = 16 KB (FP32 for stability)

Total: 96 KB, comfortably inside the 228 KB SRAM budget.

FA1 Complexity and Measured Gains

HBM traffic drops from O(N²) to O(N · d) for the forward pass:

• Reads: Q, K, V each once = 3 × 2Nd bytes
• Writes: O once + statistics = ~2Nd + 2N bytes

The factor of improvement over standard attention at N = 4096, d = 128: the standard path moves ~4 × 2 × 4096² ≈ 134 MB per head for intermediates; FA1 moves ~3 × 2 × 4096 × 128 ≈ 3 MB. Roughly 44× less HBM traffic for the intermediates.

Wall-clock speedup on A100: 2–4× end-to-end over PyTorch standard attention at long sequences. The backward pass, which recomputes S and P from saved (m, ℓ) rather than reading them from HBM, runs at ~2.5× the FLOPs of the forward but still beats the standard backward because the HBM savings dominate.

Part 3: Flash Attention 2 — The Algebra of Wasted FLOPs

FA1 was memory-IO optimal in its access pattern, but it left compute efficiency on the table. Dao’s follow-up (arXiv:2307.08691, July 2023) identified three sources of waste and fixed all three.

Fix 1: Eliminate Redundant Rescaling FLOPs

In FA1’s recurrence, the output accumulator O is rescaled by exp(m_{t-1} - m_t) at every block, even when the running maximum does not change. FA2 delays the final normalization to the very end:

O_t = O_{t-1} + P̃_t · V_t          # no per-block rescaling
# Final step only:
O = diag(exp(m_{T_c} - running_max))⁻¹ · diag(ℓ_{T_c})⁻¹ · O_{T_c}

When the local max equals the global max (which it does for all but the block containing the true maximum), the correction factor is exp(0) = 1 and can be skipped. This eliminates a diag() matrix multiply per block — non-matmul FLOPs that do not map to tensor cores and burn scalar execution units.

The net effect: FA2 reduces non-matmul FLOPs by roughly 2× compared to FA1. Since non-matmul FLOPs execute at roughly 16× lower throughput than matmul FLOPs on A100/H100 (scalar units vs tensor cores), this has outsized impact on latency.

Fix 2: Outer Loop Over Query Blocks

FA1 parallelizes across batch and head dimensions. The inner loop iterates over Q rows; the outer loop over K/V blocks. This means all Q tiles for a given (batch, head) are processed by a single thread block, sequentially. At long sequences with few heads, this underutilizes the SM grid.

FA2 swaps the loop order: the outer loop is over Q blocks, the inner loop over K/V blocks. Now each Q tile can be assigned to an independent thread block and run concurrently on different SMs. For a sequence of length N = 8192 with Br = 64, there are 128 Q blocks per head — 128-way parallelism across the sequence dimension, orthogonal to the existing batch and head parallelism.

This restructuring also makes the causal masking case more efficient: when the outer loop is over Q, a Q block at position t only needs to attend to K/V blocks at positions ≤ t. FA2 skips the masked blocks entirely, recovering the factor-of-2 in compute that FA1 wasted on computing masked-out score values.

Fix 3: Warp-Level Work Partitioning

FA1 split the K/V dimension across warps and required a synchronization step (an all-reduce across warps) to combine partial softmax statistics. FA2 partitions differently: each warp takes ownership of a contiguous slice of the Q sequence dimension, not the K/V dimension. Since each Q row’s softmax is independent, warps need not communicate. The inter-warp synchronization is eliminated entirely.

FA2 measured gains: approximately 2× over FA1 in FLOP/s utilization on A100, and 2× in wall-clock at typical sequence lengths (2k–8k). Peak HBM bandwidth utilization rises from FA1’s ~30–40% to FA2’s ~50–70%.

Part 4: Flash Attention 3 — Hopper’s Hardware Is a Different Machine

FA1 and FA2 were written for a hardware abstraction that Hopper broke. The H100’s Hopper microarchitecture introduced capabilities that FA2 could not exploit: TMA (Tensor Memory Accelerator) units for asynchronous HBM-to-SRAM copies, wgmma (warpgroup async MMA) for double-buffered tensor core operations, and a warp specialization model that separates producer and consumer threads at the hardware level. FA3 (arXiv:2407.08608, July 2024) is a ground-up rewrite to exploit all three.

TMA: Decoupling Memory from Compute

In pre-Hopper PTX, the warp scheduler issued memory loads. Load instructions consumed scheduler slots, competed with arithmetic instructions for issue ports, and stalled the warp until the load was in-flight. The programmer could hide latency only through software pipelining — keeping multiple live tiles in SRAM simultaneously so that one tile’s load overlapped with another tile’s compute.

TMA is a dedicated copy engine that operates asynchronously from the warp scheduler. An H100 SM can issue a TMA descriptor (a pointer, shape, and stride specification) and then proceed to arithmetic instructions while the TMA hardware DMA’s data from HBM into SRAM. The warp scheduler never stalls on memory.

FA3 exploits TMA through warp specialization: some warps in a warpgroup are designated “producers” and issue TMA copy instructions; others are “consumers” and issue wgmma matmul instructions. The two sets of warps never compete for functional units. This is not a software trick — it maps to distinct H100 hardware paths.

3-Stage Pipelining

FA2’s pipeline has two stages: load tile into SRAM, compute attention on tile. The latency of the load is partially hidden behind the compute of the previous tile.

FA3 introduces a 3-stage pipeline: load Q/K/V tiles (TMA), compute scores S_t = QK_tᵀ via wgmma, compute output update O += P̃_t · V_t. The three stages overlap across three consecutive K/V blocks:

Block t-1:  [TMA load] [score compute] [output compute]
Block t:               [TMA load]      [score compute]  [output compute]
Block t+1:                             [TMA load]       [score compute]  ...

The output-update matmul for block t-1 overlaps with the score compute for block t and the TMA load for block t+1. Three tile’s worth of SRAM staging buffers are maintained simultaneously. At Bc = 64, d = 128, FP16: 3 × (64×128×2) × 2 (K and V) = 96 KB of staging — well inside the 228 KB budget.

The gain from the 3-stage pipeline over FA2’s 2-stage pipeline is roughly 15–20% in isolation, but the warp specialization compounds it.

FP8 Path and Quantization

Hopper’s tensor cores support FP8 (E4M3 and E5M2) natively, doubling throughput over FP16 at the same clock. FA3 implements an FP8 attention path with per-tile quantization: before passing a K tile to wgmma, the producer warps quantize K_t to FP8 in SRAM using a per-tile scale factor derived from the tile’s absolute maximum. The score S_t is computed in FP8, then dequantized to FP32 for the softmax numerics. V tiles are similarly quantized.

Per-tile quantization degrades accuracy relative to FP16 in proportion to the outlier concentration within a tile. For standard language model attention, the degradation is negligible (< 0.1% perplexity on standard benchmarks). For attention patterns with extreme outliers — long documents where a single key dominates — it can be measurable. The FP8 path doubles FLOPs available for attention at the same HBM bandwidth, effectively raising the arithmetic intensity of the score computation to ~128 FLOPS/byte and pushing the kernel toward the FP8 compute ridge.

FA3 Measured Performance

On H100 SXM5:

• FP16, N = 8192, d = 128, causal: FA3 achieves ~740 TFLOPS vs FA2’s ~370 TFLOPS — ~1.5–2× over FA2
• FP8 path: up to ~1.2 PFLOPS on H100 (approaching the FP8 tensor core ceiling)
• End-to-end transformer throughput improvement over FA2: ~1.5× on representative workloads

The gap between FA2 and FA3 is larger for longer sequences (where the pipelining efficiency dominates) and smaller for batch-size-1 decode (where the working set is small enough that HBM latency, not bandwidth, is the limit).

Part 5: When Flash Attention Breaks, and When to Choose Something Else

FA is not universally applicable. The tile-and-recompute strategy carries implicit assumptions about attention structure that, when violated, either produce incorrect outputs or negate the performance benefits.

Masking Patterns

Causal masking is natively supported. FA2 and FA3 both skip the upper-triangular tiles of the attention matrix when processing causal sequences, recovering the factor-of-2 over bidirectional attention at no accuracy cost.

Padding masks (variable-length sequences in a batch) are handled via the “varlen” API in FlashAttention, which packs sequences without padding into a single tensor and uses a cumulative-length array to identify boundaries. Without this, short sequences in a padded batch cause FA to compute and discard quadratically many masked tokens.

Arbitrary custom bias (ALiBi, position-dependent slopes, document-boundary masks) requires that the bias be addable to S_t in-register before the softmax. FA supports bias addition per tile, but the bias tensor must be loaded from HBM, partially negating the IO savings. For bias tensors that are structured (ALiBi’s linear slope is computable on the fly from position indices), this cost is zero. For dense N×N bias matrices, the IO cost is O(N²) regardless — same as standard attention — and FA provides no benefit.

Sliding window attention (each token attends only to a local window of W positions) is supported by FA3 and by the xformers implementation. For window W << N, only ~W/N fraction of K/V tiles are loaded per Q block, giving IO complexity O(NW) instead of O(N²). This is the regime where FA’s benefits are largest in absolute terms: long documents with local attention patterns.

When FA Is Suboptimal

Decode-time single-token generation (batch size 1, one new query token per step): the Q tensor has shape [1, d]. The K/V cache has shape [context_length, d]. The matmul QKᵀ is [1, context_length] — a vector-matrix product, not a matrix-matrix product. Tensor cores are designed for matrix-matrix products; vector-matrix products expose the bandwidth bottleneck of loading K/V cache weights with essentially zero arithmetic intensity. FA’s tiling overhead (per-tile softmax state management) is pure overhead for this workload. PagedAttention (vLLM) addresses the KV cache management problem here, but the fundamental compute pattern is beyond FA’s remit.

Very short sequences (N < 128): the tile overhead dominates. Standard torch.nn.functional.scaled_dot_product_attention with enable_flash_sdp=True degrades to a non-FA kernel for short sequences automatically.

Non-standard data types or accumulation requirements: FA3’s FP8 path assumes the application can tolerate per-tile quantization. Mixture-of-experts routing attention or cross-attention patterns with non-contiguous K/V layouts may require custom CUDA kernels beyond FA’s default configuration.

Part 6: The Alternatives — xFormers and FlexAttention

xFormers Memory-Efficient Attention

Meta’s xFormers library (available via xformers.ops.memory_efficient_attention) predates FA1’s public release and implements a CUTLASS-based tiled attention kernel with the same fundamental optimization: tiled SRAM computation, no materialized N×N intermediates. The implementation is functionally equivalent to FA1 for standard attention patterns.

xFormers’ advantage is flexibility. Because it uses CUTLASS templates rather than hand-written PTX, it supports a wider range of head dimensions and arbitrary attention biases without kernel recompilation. Its disadvantage is throughput: on H100, xFormers’ memory-efficient attention achieves roughly 60–75% of FA3’s TFLOPS at the same configuration. For research code where custom bias tensors or non-standard head dimensions are common, xFormers remains the pragmatic choice.

FlexAttention (PyTorch 2.5+)

FlexAttention, introduced in PyTorch 2.5, takes a different approach to the flexibility-performance tradeoff. The user supplies a Python callable score_mod(score, batch, head, q_idx, k_idx) that transforms the raw attention score before the softmax. FlexAttention then torch.compiles this callable into a fused FA-style kernel that preserves the IO-optimal tiling structure while incorporating arbitrary score modifications.

The practical significance: any attention variant expressible as a pointwise transformation of scores — ALiBi, Rotary PE correction, document masking, soft-capping (as in Gemma-2) — can be implemented in FlexAttention with zero CUDA coding and close to FA3 throughput. The compiler handles tiling, SRAM staging, and the online softmax integration.

The limitation: the score_mod function must be a pointwise operation on the scalar score. Operations that require cross-position information (e.g., sparse attention patterns that depend on content rather than position) cannot be expressed in the current API. For these, custom CUDA kernels remain necessary.

Performance comparison on H100 (N = 4096, d = 128, causal, FP16):

Epilogue: The Lesson That Generalizes

Flash Attention’s trajectory — from a 2022 paper that the ML community initially found “too implementation-focused” to the de facto standard attention kernel on every major training cluster — illustrates a principle that extends beyond attention.

The principle: mathematically equivalent programs are not performance-equivalent on real hardware. The gap between standard attention and FA3 is not a gap in the mathematics. Both compute the same function. The gap is entirely in how many times bytes cross the HBM↔SRAM boundary, and in how well the computation respects the five-order-of-magnitude latency hierarchy between HBM (~400 cycles) and registers (~1 cycle).

FA1 moved the N×N matrix out of HBM. FA2 eliminated the algebraic waste in the recurrence. FA3 mapped the algorithm to Hopper’s warp-specialized hardware. Each step was not a new algorithm — it was a more faithful translation of the existing algorithm onto the actual capabilities of the memory hierarchy.

The implication for systems engineers is uncomfortable: no benchmark that runs at small N with float32 on a single GPU tells you whether your attention implementation is correct. The correctness you care about is not numerical — it is IO-correctness. Does the implementation respect the hierarchy? Does it avoid materializing tensors in HBM that the math does not require? The roofline model, not the loss curve, is the ultimate test.

References

Dao, T., Fu, D. Y., Ermon, S., Rudra, A., & Ré, C. (2022). FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. arXiv:2205.14135.

Dao, T. (2023). FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning. arXiv:2307.08691.

Shah, J., Bikshandi, G., Zhang, Y., Thakkar, V., Ramani, P., & Dao, T. (2024). FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision. arXiv:2407.08608.

Milakov, M., & Gimelshein, N. (2018). Online normalizer calculation for softmax. arXiv:1805.02867.

NVIDIA. (2022). NVIDIA H100 Tensor Core GPU Architecture Whitepaper.

He, H. (2022). Making Deep Learning Go Brrr From First Principles. https://horace.io/brrr_intro.html

@article{fp4-2606005,
  title   = {Flash Attention: A Complete Architectural Autopsy},
  author  = {fp4 editorial desk},
  year    = {2026},
  url     = {https://fp4.dev/algorithm/flash-attention-evolution/},
  journal = {fp4}
}