Deep Dive - FlashAttention

In full dense self-attention, each query compares with every key. Causal decoders mask future positions, but the number of allowed query-key pairs still grows quadratically. The arithmetic is familiar, but the GPU has a second problem: where each intermediate lives while that arithmetic runs. FlashAttention keeps exact attention while changing the order and location of those reads and writes.

This chapter reads the Dao-AILab implementation as an engineering project. It connects original papers to the source tree, works through a two-tile softmax by hand, then identifies when a production system should use FlashAttention, a framework backend, or a fallback. The earlier scaled dot-product attention lesson supplies the formula. Here the focus is the memory traffic around it.

What stays mathematically unchanged when FlashAttention is used?

Answer

The operator is still softmax of scaled $QK^T$ multiplied by $V$. FlashAttention changes scheduling and memory movement, not the attention definition.

The memory wall behind exact attention

Scaled dot-product attention forms a score for every query-key pair. For sequence length $n$, that is $n^2$ scores per head before multiplying by values. A baseline implementation can write scores, read them for softmax, write probabilities, read them for the value multiply, and finally write the output. Each round trip crosses HBM (high-bandwidth memory), the large GPU memory pool.

HBM has high capacity and high aggregate bandwidth, but an on-chip SRAM (static random-access memory) workspace is much smaller and faster to revisit. GPU arithmetic units can perform many fused multiply-adds while a kernel waits for HBM transactions. Kernel analysis therefore needs both FLOPs and bytes crossing the chip boundary.

The original FlashAttention paper describes an A100-era gap between aggregate on-chip SRAM and HBM capacity and bandwidth.^{[1]Reference 1FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness.https://arxiv.org/abs/2205.14135} Exact values depend on the GPU and the available per-kernel budget. Treat those paper-era values as an explanation of the design pressure, not as a current hardware specification.

The asymptotic score calculation stays $O(n^2d)$ FLOPs. The memory that must remain live for intermediate scores changes from quadratic to linear in sequence length. That distinction is why "memory efficient" doesn't mean "less attention math."

A size check before blaming weights

Suppose a training batch has 8 sequences, 32 heads, sequence length 8,192, and FP16 scores. One dense score tensor would need:

$8 \times 32 \times 8192^2 \times 2\;\text{bytes} = 34{,}359{,}738{,}368\;\text{bytes} \approx 32\;\text{GiB}.$

That is only scores. Probabilities, model weights, activations, gradients, and allocator fragmentation still need space. The number is a sizing exercise, not a benchmark, and it assumes an implementation actually materializes that tensor.

Expected output is score values: 17,179,869,184 and binary size: 32.00 GiB. A fused kernel can avoid this allocation, but the pair count still matters for compute and data movement.

Why can a fused attention kernel still be compute-heavy at long context?

Answer

It avoids storing the full score matrix, but it still evaluates the allowed query-key pairs. IO savings remove a large memory cost; they don't make the quadratic dot-product work disappear.

Tiling keeps the result exact

FlashAttention chooses a query tile $Q_i$ and walks over key/value tiles $(K_j, V_j)$. A tile is small enough for on-chip SRAM. The kernel computes $S_{ij}=Q_iK_j^T$ there, applies mask and scale, updates a row-wise softmax state, and discards $S_{ij}$ before loading the next tile.

The control flow is easier to see than the CUDA code. Each row of $Q_i$ carries three pieces of state:

• $m_i$: largest score seen so far, used for numerical stability.
• $\ell_i$: sum of exponentials after shifting by that largest score.
• $O_i$: normalized output accumulated from the tiles seen so far.

m_i and \ell_i are running softmax statistics from which $\operatorname{LSE}_i=m_i+\log \ell_i$ is formed. They are small vectors, one entry per query row, not $n \times n$ matrices.

The diagram's last edge is a loop in the kernel: the next tile reuses the same query tile and softmax state. Only after the final key tile does the kernel write the output row to HBM.

Online softmax recurrence

Normal softmax appears to need every score before it can compute a denominator. The online algorithm instead merges two partial summaries. Let a previous tile have state $(m_{old}, \ell_{old}, O_{old})$ and the new score tile have row scores $s$. Compute:

$m_{new}=\max(m_{old},\max(s)).$

Then shift both old and new contributions to the same maximum:

$\ell_{new}=e^{m_{old}-m_{new}}\ell_{old}+\sum_k e^{s_k-m_{new}}.$

For values $V_j$, update the numerator and normalize at the end:

$O_{new}=\frac{e^{m_{old}-m_{new}}\ell_{old}O_{old}+\sum_k e^{s_k-m_{new}}V_{j,k}}{\ell_{new}}.$

The exponential shift keeps exponents at or below one. When a later tile contains a larger score, the old accumulator is rescaled exactly instead of being thrown away. Milakov and Gimelshein describe this streaming normalizer, which FlashAttention uses as part of its tiled exact computation.^{[2]Reference 2Online normalizer calculation for softmax.https://arxiv.org/abs/1805.02867}

Worked numbers across two tiles

Use one query row and four keys split into two tiles. The first tile has scores $[1,2]$ and values $[[10,0],[0,10]]$. The second has scores $[3,0]$ and values $[[5,5],[9,1]]$. We omit the common scale in this arithmetic so the state update stays readable.

For tile one, $m_1=2$ and $\ell_1=e^{-1}+1=1.3679$. Its numerator is:

$e^{-1}[10,0]+[0,10]=[3.6788,10].$

So $O_1=[2.689,7.311]$. That output is provisional because key tile two hasn't been seen.

Tile two contains a larger score, so $m_2=3$. The old state must be rescaled by $e^{2-3}=0.3679$:

$\ell_2=0.3679\times1.3679 + (1+e^{-3}) = 1.5530.$

The new numerator is:

$0.3679[3.6788,10] + [5,5] + e^{-3}[9,1] = [6.8014,8.7286].$

The final result is $O_2=[4.379,5.620]$. Computing ordinary softmax over all four scores gives the same values up to rounding. No $1\times4$ score row had to survive between tiles.

What error would appear if a later tile had a larger maximum but the old numerator wasn't rescaled?

Answer

The old contributions would be compared in the wrong exponential scale. Their relative probability would be too large, so the result wouldn't equal dense softmax. Rescaling by $e^{m_{old}-m_{new}}$ is the invariant that makes tile order exact.

Read the implementation boundary

The official repository keeps Python wrappers and GPU kernels in separate layers.^{[3]Reference 3FlashAttentionhttps://github.com/Dao-AILab/flash-attention} The public flash_attn_func interface accepts tensors shaped like (batch, sequence, heads, head_dim), normalizes optional arguments, and calls a compiled forward operator. The wrapper also exposes variable-length and key/value-cache paths.

The CUDA path registers custom PyTorch operators, returning output, row-wise softmax LSE, an optional probability or dropout-mask buffer, and RNG state. The LSE tensor is compact: it has one value per query row and head. It gives backward a stable summary without saving every probability.

The source has a ROCm branch selected by FLASH_ATTENTION_TRITON_AMD_ENABLE. On CUDA, the wrapper imports the compiled flash_attn_2_cuda extension. On ROCm, it can route to an AIter Triton implementation. This dispatch is part of the product: a fast algorithm still needs a kernel matching the device, dtype, head dimension, and layout.

The expected output is flash-attn-cuda followed by framework-fallback. The snippet models a guard, not the full runtime selector. Real dispatch also checks GPU capability, dtype, sequence lengths, masks, dropout, and compiled extension availability.

Backward pass: recompute instead of save

Training needs gradients for $Q$, $K$, and $V$. A materializing implementation can save the full probability matrix for backward, but that repeats the quadratic memory problem. FlashAttention saves $Q$, $K$, $V$, output $O$, row-wise LSE, and any required random-number state, then recomputes score and probability tiles during backward.

For each tile, backward reconstructs the same shifted probabilities in SRAM, uses incoming gradient $dO$, and accumulates $dQ$, $dK$, and $dV$. Recomputing scores costs additional FLOPs and HBM reads, but it avoids storing and loading a full $P$ matrix. This is a classic memory-compute trade-off, not a free speedup.

FlashAttention-2 adds work partitioning changes so more thread blocks can process long sequences and less serial work sits in each block.^{[4]Reference 4FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning.https://arxiv.org/abs/2307.08691} The repository exposes a deterministic option for backward. Deterministic backward uses more memory and can be slower, while forward remains deterministic according to the project interface.^{[3]Reference 3FlashAttentionhttps://github.com/Dao-AILab/flash-attention}

Why does backward recomputation often make sense for training?

Answer

Saving a quadratic probability matrix can exceed activation memory at long context. Recomputing score tiles spends extra arithmetic and reads, but keeps only linear-size output and LSE state, often enabling a larger batch or sequence length.

Masks and structured layouts

FlashAttention isn't limited to unmasked full-sequence attention. The same tile loop can skip or alter score elements before softmax.

Causal attention

Autoregressive decoders can't read future tokens. A causal tile masks positions above the diagonal, treating them as $-\infty$ before softmax. The wrapper aligns causal masks for unequal query and key lengths, which matters during decode when a short query attends to a long prefix.

Local or sliding-window attention

The window_size=(left,right) argument limits each query to nearby keys. Local attention reduces the number of allowed scores for long contexts, but it isn't automatically faster for every shape. Tile occupancy, mask shape, and GPU generation decide whether the work reduction outweighs branch and scheduling overhead.

Variable length and packed batches

Training batches often contain sequences with different lengths. The varlen interface uses cumulative sequence lengths (cu_seqlens_q and cu_seqlens_k) so padding doesn't force every row to the longest sequence. This changes indexing and launch metadata, not the online-softmax invariant.

Paged key/value cache

Decode systems may store K/V in fixed-size physical pages and map each request through a block table. The implementation accepts a page table for this mode. FlashAttention handles the attention math and page lookup; a serving engine still owns allocation, eviction, and request scheduling. Don't confuse the kernel's paged path with a complete KV-cache manager.

MQA and GQA

Multi-query attention (MQA) shares one K/V head across query heads. Grouped-query attention (GQA) shares K/V heads in groups. Passing fewer K/V heads than Q heads reduces cache traffic, provided the query-head count is divisible by the K/V-head count. This is an architectural choice with quality and bandwidth trade-offs; FlashAttention's interface supports it but doesn't choose the model's head layout.

FlashAttention 1 through 4

Each release combines an algorithmic idea with a hardware target. Paper benchmark numbers are snapshots from specific GPU, dtype, sequence, and baseline conditions. They shouldn't be copied into a service SLO.

The repository README describes FA3 as a beta path with H100/H800 and CUDA requirements, while FA4 is exposed through a CuTeDSL package for Hopper and Blackwell.^{[3]Reference 3FlashAttentionhttps://github.com/Dao-AILab/flash-attention} Read those release notes as compatibility guidance. A paper name in a citation doesn't mean every installed wheel has that generation enabled.

Project identity

The team is best understood as an open research-and-engineering project. The Dao-AILab repository carries the implementation, tests, build scripts, and issue history; the papers identify Tri Dao and collaborators as authors. Downstream framework maintainers contribute integrations and backend work. That split explains why paper ideas, compiled kernels, and framework dispatch evolve on different schedules.

Framework dispatch and application surface

PyTorch exposes scaled dot-product attention (SDPA), which can select among flash, memory-efficient, and math implementations based on input and device constraints.^{[9]Reference 9torch.nn.functional.scaled_dot_product_attentionhttps://docs.pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html} A model library should call the framework API when it wants portability and let the dispatcher pick a valid backend. It should call the Dao-AILab package when it needs a feature or kernel version that the framework path doesn't provide.

The repository usage page lists integrations or adoption examples in PyTorch, Transformers, DeepSpeed, Megatron-LM, diffusion systems, and protein-structure models.^{[3]Reference 3FlashAttentionhttps://github.com/Dao-AILab/flash-attention} Those examples cover two broad workloads:

1. Training: reduce activation memory and improve attention throughput, allowing longer sequences or larger batches under a fixed memory budget.
2. Inference: accelerate prompt processing and some decode attention, especially when the shape reaches the kernel's efficient regime. KV-cache layout and scheduler behavior still dominate many decode workloads.

Use the paper or a local benchmark for any speed claim. The measured gain depends on sequence length, batch shape, head dimension, dtype, dropout, mask, GPU generation, compiler version, and baseline backend. "Flash" is not a guarantee that a tiny or unusual tensor wins.

Strengths and weaknesses

The most important weakness is an operational one: a kernel can be correct but still be the wrong choice for a workload. Small batches, unsupported head dimensions, CPU or Apple MPS runs, unusual masks, and missing compiled extensions all need explicit fallback behavior.

A small correctness exercise

The following pure-Python check compares a full stable softmax with a two-tile streaming implementation. It uses one query row and scalar values so every intermediate can be inspected. This is a teaching model, not a GPU benchmark.

Both lines print full: 4.379542 and tiled: 4.379542 to six decimals. Try moving score 3.0 into the first tile. The result should stay the same, which checks that state merging is independent of the tile boundary.

Production checklist

Before enabling a FlashAttention backend, record the exact workload and hardware:

• Shape: batch, query length, key length, head count, K/V head count, and head dimension.
• Numerics: FP16 or BF16 input, FP32 accumulation where supported, scale, dropout, and tolerance against a stable reference.
• Mask: causal alignment, local window, padding, varlen offsets, and page-table indexing.
• Device: GPU compute capability, CUDA or ROCm version, driver, PyTorch version, and compiled extension.
• Dispatch: which backend was selected, whether a fallback occurred, and why.
• Evidence: p50 and tail latency, memory peak, throughput, and correctness error on representative prompts.

PyTorch's SDPA documentation is a good first check for backend eligibility.^{[9]Reference 9torch.nn.functional.scaled_dot_product_attentionhttps://docs.pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html} The repository tests are the next check for a direct package integration. A benchmark should compare the backend that production would otherwise choose, not a deliberately slow baseline.

If a release upgrades CUDA, PyTorch, GPU generation, or FlashAttention, rerun correctness and performance checks. Keep a fallback path in the service. A request should fail over to a framework implementation or a safe math path with an observable metric, rather than silently return wrong output or crash because a niche shape missed a kernel specialization.

⚠️ Common mistake: Treating "no quadratic allocation" as "constant-time attention." FlashAttention still computes dense query-key pairs, and long-context decode can remain dominated by KV-cache reads and scheduler effects.

Key takeaways

• FlashAttention computes exact scaled dot-product attention while controlling HBM traffic.
• Tiling keeps active Q, K, V, score, and value work in SRAM; online softmax preserves exact normalization with $(m,\ell,O)$ state.
• Backward recomputes score tiles to trade extra FLOPs for much smaller saved activation memory.
• Causal, local, variable-length, paged, MQA, and GQA modes widen the application surface, but every mode has shape and dispatch constraints.
• FA1 through FA4 pair the same core IO idea with different parallelism and hardware pipelines. Paper-era speedups are conditional measurements.
• Production confidence requires correctness checks, dispatch observability, representative benchmarks, and a tested fallback.