FlashAttention from scratch: how it works and how to code it

In this article

1. Before you start
2. Why interviewers ask this
3. FlashAttention from first principles
4. Worked example
5. Trade-offs and edge cases
6. How to answer it in an interview
7. Common mistakes
8. Follow-up questions
9. Sources

Before you start

This page is one lesson in the PyTorch interview questions and the machine learning (ML) coding round guide.

Why interviewers ask this

FlashAttention appears in coding rounds because it tests three things at once. The first is whether you can write attention without a library, the second is whether you understand why softmax can be computed in pieces, and the third is whether you know that speed on a GPU is often about memory traffic, not arithmetic. Candidates report being asked to implement single-headed flash attention in NumPy, and one report describes a three-part round: cross attention, then flash attention, then the backward pass. These are candidate reports from our question bank, not confirmed by a public source.

A weak answer says "it splits attention into blocks to save memory" and cannot write the merge. A strong answer writes the running maximum and running sum, explains the rescale, and checks the result against ordinary attention. A Staff-level answer adds why it is faster despite extra arithmetic and how the backward pass recomputes score tiles.

FlashAttention from first principles

FlashAttention

FlashAttention is an exact attention algorithm that processes the keys and values in blocks held in fast on-chip memory and merges the blocks with running softmax statistics, so the full score matrix is never written to the GPU's main memory.

Recomputation

Recomputation is a training technique that discards intermediate values in the forward pass and rebuilds them when the backward pass needs them, trading extra arithmetic for lower memory use and, in FlashAttention, less memory traffic.

In plain words

Picture a cook with a tiny cutting board and a huge pantry down the hall. Walking to the pantry takes far longer than chopping. A slow cook fetches everything, chops it, writes the whole result on a giant sheet, walks it back, then fetches it again. A fast cook carries a small armful, chops it, and keeps a short running tally. The analogy breaks in one place: the tally must be corrected whenever a later armful changes the scale of the total. That correction is the maths of this page.

In one precise sentence, FlashAttention gives the same output as standard attention without storing the score matrix, by merging key and value blocks into running statistics.

Step 1: Split Q, K and V into blocks that fit in on-chip memory

Standard attention on a GPU is slow because it writes a huge intermediate matrix to the GPU's slow main memory, and its softmax and elementwise steps are limited by that memory traffic. A GPU has two memories that matter: high bandwidth memory (HBM), which is large but slow, and static random-access memory (SRAM), which is small, on-chip and fast.

For the A100 GPU, Dao et al. (2022) list 40 to 80 gigabytes of HBM at 1.5 to 2.0 terabytes per second against 192 kilobytes of SRAM per streaming multiprocessor at about 19 terabytes per second. Standard attention writes the N-by-N score matrix S and weight matrix P to HBM, which needs O(N²) memory (Dao et al., 2022). For N of 1,024 in 32-bit floats, one head's S is 4 MiB; at 16,384 it is 1,024 MiB.¹

FlashAttention splits Q, K and V into blocks and loads one block of K and V at a time. The paper sets the block sizes from the SRAM size M and head size d: Bc = ⌈M/4d⌉ and Br = min(⌈M/4d⌉, d). For an illustrative M of 25,600 elements and d of 64, Bc is 100 and Br is 64.¹

Step 2: Rescale the old partial result when the maximum grows

A softmax needs its whole row's maximum and sum, but a block sees only part of the row. FlashAttention therefore keeps two running numbers per query row: m, the largest score so far, and l, the sum of exp(score − m) so far. Subtracting the maximum is the trick of a numerically stable softmax.

When a later block raises the maximum, every earlier exponential used a stale maximum. The fix is to multiply the old sum and the old output by e^(m_old − m_new), a factor between zero and one. In Algorithm 1 of Dao et al. (2022) this is the update l_new = e^(m_old − m_new) · l_old + e^(m_tile − m_new) · l_tile. The tile's own exponentials are scaled by e^(m_tile − m_new) the same way. This page keeps the running output unnormalized, which is a FlashAttention-2-style choice. Computing a softmax in one pass with running statistics predates FlashAttention: Milakov and Gimelshein (2018) proposed computing softmax with fewer memory accesses.

Step 3: Loop over the blocks and divide once at the end

For one query row, the loop visits every block of keys and values. Each pass computes the score tile, takes its local maximum, forms local exponentials, merges them into m, l and the running output o with the Step 2 rule, and moves on. This page keeps o unnormalized and divides the output by the running sum once, after the last block. The original Algorithm 1 instead divides on every pass and multiplies the old output back by l_old before the next merge. Both are exact. Deferring the division, as FlashAttention-2 does, saves non-matrix-multiply work. What is wrong is dividing early without undoing it.

Dao et al. (2022, Theorem 1) show the result equals softmax(Q @ K.T) @ V with O(N²d) floating-point operations and only O(N) extra memory beyond inputs and output. The paper's loops run over K and V blocks on the outside and Q blocks on the inside, so output rows are rewritten to HBM on each pass; the arithmetic is the same. Fusing all of this into one GPU kernel keeps score tiles in SRAM; see GPU kernel questions in ML coding rounds for kernel fusion.

Step 4: Save only the statistics and recompute the tiles in the backward pass

Training needs gradients, and the usual way is to keep S and P from the forward pass. FlashAttention keeps only the output plus the per-row max and sum. In the backward pass it reloads Q, K and V blocks and rebuilds the score tiles on-chip (Dao et al., 2022). That is more arithmetic, yet the paper reports it is faster than reading the stored matrix back from HBM. Ordinary gradient checkpointing saves memory by spending time; here the recomputation also saves time, because the bottleneck is memory traffic.

1 · Load blocks2 · Merge with rescale3 · Loop, then divide4 · BackwardReplay

Step-by-step description

1. A block of K and V fits in fast memory.
2. Old state is scaled when the max grows.
3. Divide by the running sum once.
4. Score tiles are recomputed, not stored.

• Blocks of K and V are loaded into fast on-chip memory.
• The old output is scaled by e to the old maximum minus the new maximum.
• This page divides the output by the running sum once, after the last block.
• The backward pass recomputes score tiles from the saved maximum and sum.

Worked example

Take one head with four queries, four keys and two-dimensional vectors, all illustrative: Q rows (1,0), (0,1), (1,1), (2,0); K rows (1,0), (0,2), (1,1), (3,0); V rows (1,0), (0,1), (2,2), (1,3). Scores are plain dot products with no 1/√d scaling, as in the paper's formula. Use blocks of two keys and follow query 3, whose scores against the four keys are 1, 2, 2 and 3.¹

Block 1 has scores 1 and 2, so m is 2 and the weights are 0.3679 and 1, giving l of 1.3679 and an unnormalized output of (0.3679, 1.0). Block 2 has scores 2 and 3. Its local maximum 3 beats the old 2, so the old state is scaled by e^(2 − 3) = 0.3679, which gives l of 1.8711 and an output of (1.8711, 4.1036). Dividing by l gives (1.0, 2.1932), identical to the ordinary softmax over all four scores. The merged sum is 1.8711 and the final output is (1.0, 2.1932).¹

If the rescale is skipped, the code adds block 1's exponentials (taken against maximum 2) to block 2's (taken against maximum 3) as if they shared one scale, and returns (0.769, 1.731) instead of (1.0, 2.1932).¹ On random inputs with 64 queries and keys, head size 16 and blocks of 8, the largest difference from ordinary attention is below 1e-14, which is rounding error.¹

Text description of the diagram

• After block 1 the max is 2, the sum is 1.3679 and the output is (0.3679, 1.0).
• Block 2 raises the max to 3, so the old state is scaled by 0.3679.
• The merged sum is 1.8711 and the final output is (1.0, 2.1932).

The same loop in NumPy

The function below handles all query rows together: m, l and O hold one value per row, and each pass of the loop merges one block of keys and values. It is the form candidates are asked to write; the line to explain is the one that multiplies O by a.

import numpy as np

def flash_attention(Q, K, V, block=2):
    N, d = Q.shape
    O = np.zeros((N, d))        # running unnormalized output
    m = np.full(N, -np.inf)     # running max per row
    l = np.zeros(N)             # running sum of exponentials per row
    for j in range(0, N, block):
        Kj, Vj = K[j:j+block], V[j:j+block]
        S = Q @ Kj.T                      # score tile
        m_tile = S.max(axis=1)
        P = np.exp(S - m_tile[:, None])
        l_tile = P.sum(axis=1)
        m_new = np.maximum(m, m_tile)
        a = np.exp(m - m_new)             # rescales the old state
        b = np.exp(m_tile - m_new)        # rescales the new tile
        l = a * l + b * l_tile
        O = a[:, None] * O + b[:, None] * (P @ Vj)
        m = m_new
    return O / l[:, None]

Run on the toy inputs above, this exact listing returns the same query-3 row as the reference, and on 64 random queries with blocks of 8 and of 3 (3 does not divide 64) it differs from ordinary attention only at rounding level.¹ NumPy runs this on the central processing unit (CPU), so it demonstrates the arithmetic, not the speed.

Measured on a real GPU

Dao et al. (2022) measured forward plus backward passes for a Generative Pre-trained Transformer 2 (generative pre-trained transformer (GPT)-2) medium model on an A100: sequence length 1,024, head size 64, 16 heads, batch 64. FlashAttention did more arithmetic and still finished far sooner.

Standard attention moved 40.3 gigabytes and took 41.7 milliseconds. FlashAttention moved 4.4 gigabytes and took 7.3 milliseconds. The ratios are computed from the paper's figures.¹ Memory traffic fell about nine times and runtime about 5.7 times, so runtime tracks data movement, not arithmetic.

Text description of the diagram

• Standard attention moved 40.3 gigabytes and took 41.7 milliseconds.
• FlashAttention moved 4.4 gigabytes and took 7.3 milliseconds.

Trade-offs and edge cases

FlashAttention changes where the work happens, not how much attention costs. Memory becomes linear in sequence length, but the arithmetic stays O(N²d) (Dao et al., 2022, Theorem 1), so it is not a sub-quadratic method. The paper's abstract contrasts it with approximate attention methods, which trade quality for compute.

Edge cases to know. A causal mask is applied to each score tile before the row maximum is taken; see what the causal mask does. A tile that is fully masked has a row maximum of −∞, and exp(−∞ − (−∞)) is NaN, so guard that case (treat the tile as contributing nothing); blocks entirely above the diagonal can simply be skipped. Real attention divides scores by √d_k first, as in scaled dot-product attention. With several heads, each head runs the same loop; see multi-head attention without nn.

Later work keeps the idea and changes the scheduling: Dao (2023) reports that FlashAttention reached only 25 to 40 percent of the theoretical maximum floating-point throughput and tweaks the algorithm to cut non-matrix-multiply operations (one tweak is to defer the division to the end) and to parallelize across thread blocks. As of this page's check, Shah et al. (2024) carry the idea to newer GPUs with asynchrony and low-precision arithmetic in FlashAttention-3. Where attention sits in a full model is covered in build GPT from scratch.

How to answer it in an interview

The spoken structure: name the running statistics, write the rescale, divide at the end, then give the memory-traffic reason and the test.

Common mistakes

• "I normalize each block by its own sum." (wrong) Each block's sum misses the other blocks. (right) "I merge into a running sum and divide once at the end."
• "I just add the block outputs." (wrong) Blocks with different maxima have different scales. (right) "I rescale the old state by e to the old max minus the new max."
• "FlashAttention approximates attention to go faster." (wrong) The output equals standard attention up to rounding. (right) "It is exact; the speed comes from less memory traffic."
• "It is faster because it does fewer operations." (wrong) The paper measured more arithmetic. (right) "It does more arithmetic and moves less data, and data movement dominates."
• "Linear memory means linear time." (wrong) The arithmetic is still quadratic. (right) "Memory is linear in N; the arithmetic is still O(N²d)."

Follow-up questions

How do you implement the backward pass of flash attention?

Save the output and the per-row max and sum from the forward pass. In the backward pass, reload Q, K and V blocks, rebuild each score tile and its weights from the saved statistics, and accumulate the gradients block by block (Dao et al., 2022). With the saved log-sum-exp L = m + log l, each tile gives P = exp(S − L), then dV += Pᵀ dO and dP = dO Vᵀ. With D = rowsum(dO ∘ O), dS = P ∘ (dP − D), then dQ += dS K and dK += dSᵀ Q; if S includes the 1/√d_k scale, dQ and dK carry the same factor (Appendix B and Algorithm 4 of the paper).

Is FlashAttention an approximation?

No. It computes the same output as softmax(Q @ K.T) @ V, and the toy check above differs only by floating-point rounding. Approximate attention methods change the result to save compute; FlashAttention changes only the order of the work.

Why is FlashAttention faster if it does more arithmetic?

Standard attention on a GPU is limited by main-memory traffic. In the paper's measured run, reads and writes fell from 40.3 to 4.4 gigabytes while arithmetic rose from 66.6 to 75.2 billion operations, and runtime fell from 41.7 to 7.3 milliseconds (Dao et al., 2022).

How does FlashAttention-2 differ from FlashAttention?

It keeps the tiling and the running statistics but reorganizes the work. Dao (2023) reduces non-matrix-multiply operations (it defers the division to the end), parallelizes across thread blocks even for one head, and distributes work between warps (groups of GPU threads) to cut communication through shared memory.

How do you test your implementation?

Compare against a plain reference on random inputs, with block sizes that do and do not divide the sequence length, and check the maximum absolute difference is at rounding level. Include a case where a later block has the largest score, since that exercises the rescale.

Sources

1. FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness (arXiv 2205.14135) — arXiv (Dao, Fu, Ermon, Rudra, Ré), 24 Jun 2022
2. Online normalizer calculation for softmax (arXiv 1805.02867) — arXiv (Milakov and Gimelshein), 8 May 2018
3. FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning (arXiv 2307.08691) — arXiv (Dao), 18 Jul 2023
4. FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision (arXiv 2407.08608) — arXiv (Shah et al.), 11 Jul 2024