Derive online softmax, prove it is exact, and explain why it makes FlashAttention fast without reducing FLOPs.

TL;DR: Keep a running max $m$, a running normalizer $\ell$ and an unnormalized output, and rescale all three by $e^{m_{\text{old}} - m_{\text{new}}}$ whenever a new block raises the max: the final $\text{acc}/\ell$ equals softmax attention exactly. That lets attention stream K and V through on-chip SRAM once, so the $N \times N$ score matrix (128 MiB per head at 8K tokens in BF16) never touches HBM. FLOPs are unchanged; bytes moved fall, and attention is bound by bytes.

How to approach it

Lead with the obstacle: softmax needs the maximum and the sum over the whole row before any output can be written, which seems to force materializing every score. Then say the plan, which is to make the statistics mergeable so the row can be processed block by block. Clarify whether the interviewer wants the forward pass only, or also the backward pass, where the same trick reappears as recomputation.

A strong answer

Softmax attention for one query is $o = \sum_i \frac{e^{s_i}}{\sum_j e^{s_j}} v_i$ with $s_i = q \cdot k_i / \sqrt{d}$. Subtracting any constant $m$ from every $s_i$ leaves it unchanged, and choosing $m = \max_i s_i$ keeps every exponent at or below zero. That is not optional: $e^{x}$ overflows FP16 for $x > 11.09$, so real kernels keep these statistics in FP32 with the max subtracted.

Process the keys in blocks. After block $j$, hold three things: the max so far $m^{(j)}$, the normalizer $\ell^{(j)} = \sum_{i \le \text{block } j} e^{s_i - m^{(j)}}$, and the unnormalized output $a^{(j)} = \sum_{i \le \text{block } j} e^{s_i - m^{(j)}} v_i$. When block $j+1$ arrives with block max $\tilde m$:

$$m^{(j+1)} = \max(m^{(j)}, \tilde m), \quad \ell^{(j+1)} = e^{m^{(j)} - m^{(j+1)}}\ell^{(j)} + \sum_{i \in j+1} e^{s_i - m^{(j+1)}}$$

and $a^{(j+1)}$ updates with the same scale factor. The proof is one line: every term already in $\ell^{(j)}$ is $e^{s_i - m^{(j)}}$, and multiplying it by $e^{m^{(j)} - m^{(j+1)}}$ turns it into $e^{s_i - m^{(j+1)}}$, exactly the term it would have been had the new max been known from the start. By induction the final state is the full-row sum referenced to the true max, and $o = a / \ell$ is exact, not approximate.

import numpy as np

def online_attention(q, K, V, block=128):
    """q: [d], K, V: [N, d]. One pass over K/V, O(d) state."""
    m, l = -np.inf, 0.0
    acc = np.zeros(V.shape[1])
    for s in range(0, K.shape[0], block):
        scores = K[s:s + block] @ q / np.sqrt(q.shape[0])
        m_new = max(m, scores.max())
        scale = np.exp(m - m_new)        # 0 on the first block: exp(-inf)
        p = np.exp(scores - m_new)
        l = l * scale + p.sum()
        acc = acc * scale + p @ V[s:s + block]
        m = m_new
    return acc / l

Against a direct softmax on 1,000 keys (deliberately not a multiple of the block), the maximum error is $10^{-16}$ at ordinary score magnitudes and zero at magnitudes of $10^4$, where a softmax without max subtraction returns NaN.

Now the cost. Naive attention writes the score matrix $S = QK^\top$ to HBM, reads it for the softmax, writes $P$, and reads $P$ again for $PV$. At $N = 8{,}192$:

S per head          = 8,192 x 8,192 x 2 B   = 128 MiB
x 64 heads          = 8 GiB per layer, written and read at least twice
H100 at 3.35 TB/s   -> about 10 ms per layer in score traffic alone
softmax pass intensity: a few FLOPs per byte, far below the ridge of 295

FlashAttention tiles Q, K and V so a block of queries and a block of keys fit in an SM's shared memory (up to 228 KB per SM on H100), computes the block's scores, applies the recurrence in registers, multiplies by the V block, and moves on. Only Q, K, V and O cross HBM. The paper's bound is $\Theta(N^2 d^2 / M)$ HBM accesses for SRAM of size $M$, against $\Theta(Nd + N^2)$ for the naive kernel. The FLOPs are identical; for long sequences the bytes fall by a factor of roughly $M/d^2$, and because the naive version was memory-bound, wall time follows bytes.

The reversal condition: at very short sequences, say under a few hundred tokens, the score matrix fits in cache anyway and launch overhead dominates, so fused attention gains little. FlashAttention matters most exactly where context grows.

What interviewers probe next

"What does the backward pass store?" Only the per-row log-sum-exp $m + \log \ell$, one scalar per query. The backward kernel recomputes $P$ block by block from Q, K and that scalar. More FLOPs, far fewer bytes, and it wins for the same reason the forward pass does.

"What did FlashAttention-2 and -3 change?" FA2 parallelized over sequence length as well as batch and heads, and moved non-matmul work out of the inner loop, because FA1 left SMs idle at small batch. FA3 targets Hopper: warp-specialized producers issuing TMA loads, WGMMA tensor-core instructions, softmax overlapped with the next block's matmuls, and FP8 support. FlashAttention-4 (arXiv 2603.05451, March 2026) targets Blackwell, where tensor-core throughput grew faster than the exponential units, and computes part of the exponentials in software on FMA units.

"Batch 1, 128K context, decoding: does this kernel still parallelize?" Not over queries, since there is one. Split the KV length across thread blocks, run the recurrence on each split, and merge the partial $(m, \ell, a)$ triples with the same rescaling rule. That is split-KV, or FlashDecoding.

How it's scored

Weak: says FlashAttention "reduces attention complexity" or "approximates softmax", or cannot say why the max is subtracted.

Hire: writes the recurrence with the rescale factor, argues exactness, and explains that the saving is HBM traffic for the score matrix.

Strong Hire: gives the IO count, connects the overflow threshold to FP32 statistics, explains backward recomputation from the log-sum-exp, and generalizes the merge rule to split-KV decoding.

Common mistakes

• Rescaling $\ell$ but not the accumulator, which leaves the output wrong by a factor that depends on block order and only shows on inputs whose maxima arrive late.
• Initializing $m = 0$ instead of $-\infty$, which silently breaks rows whose scores are all negative.
• Claiming fewer FLOPs. FlashAttention does the same multiplications, plus recomputation in the backward pass.
• Treating causal masking as a separate pass. Fully masked blocks are skipped before any loads, which is where causal attention gets close to 2x.