Implement online softmax for a whole tile of queries, not one query, and write the rescaling step in matrix form.

TL;DR: For a query tile of Br rows, keep vectors m and l of length Br and an unnormalized accumulator A of shape Br by d. Each key block updates them as A <- diag(alpha) A + P V with alpha = exp(m_old - m_new) per row, and the division by diag(l) happens once at the end, which is the FlashAttention-2 change. For Br = d = 128 the FP32 accumulator is 64 KiB of registers per tile, and the rows with no visible key yet need a guard or alpha becomes NaN.

How to approach it

Say what changes from one query to many: every scalar in the recurrence becomes a per-row vector, and every multiply by a scalar becomes a left multiply by a diagonal matrix, which in code is a broadcast. Then ask about masking (causal, window, padding), because it decides whether a row can see no keys in a block. Finally commit to where the normalization happens, since that is the detail interviewers check.

A strong answer

Fix a query tile $Q_i \in \mathbb{R}^{B_r \times d}$ and stream key blocks $K_j, V_j \in \mathbb{R}^{B_c \times d}$. The state is $m \in \mathbb{R}^{B_r}$, $\ell \in \mathbb{R}^{B_r}$ and $A \in \mathbb{R}^{B_r \times d}$. For each block:

$$S_j = \tfrac{1}{\sqrt d} Q_i K_j^\top, \quad m' = \max\big(m, \operatorname{rowmax}(S_j)\big), \quad P_j = \exp\big(S_j - m' \mathbf{1}^\top\big)$$ $$\ell' = \operatorname{diag}(\alpha)\,\ell + P_j \mathbf{1}, \quad A' = \operatorname{diag}(\alpha)\,A + P_j V_j, \quad \alpha = \exp(m - m')$$

After the last block, $O_i = \operatorname{diag}(\ell)^{-1} A$. Each row is exactly the single-query recurrence, so exactness carries over row by row. The matrix form matters for the hardware: $P_j V_j$ is a tensor-core GEMM, while $\operatorname{diag}(\alpha) A$ is an elementwise scale of registers the thread already owns, so it costs no data movement.

import numpy as np

def flash_tile(Q, K, V, q0=0, causal=False, Bc=64):
    """Q: [Br, d], one query tile starting at position q0. K, V: [N, d]."""
    Br, d = Q.shape
    m = np.full(Br, -np.inf); l = np.zeros(Br); acc = np.zeros((Br, V.shape[1]))
    qpos = q0 + np.arange(Br)
    for s in range(0, K.shape[0], Bc):
        S = Q @ K[s:s + Bc].T / np.sqrt(d)                  # [Br, Bc]
        if causal:
            S = np.where(np.arange(s, s + S.shape[1]) <= qpos[:, None], S, -np.inf)
        m_new = np.maximum(m, S.max(axis=1))                 # [Br]
        m_ref = np.where(np.isneginf(m_new), 0.0, m_new)     # rows with nothing yet
        alpha = np.exp(m - m_ref)                            # diagonal of diag(alpha)
        P = np.exp(S - m_ref[:, None])                       # [Br, Bc]
        l = alpha * l + P.sum(axis=1)
        acc = alpha[:, None] * acc + P @ V[s:s + Bc]         # diag(alpha) @ acc + P V
        m = m_new
    return acc / l[:, None]                              # diag(l)^-1 once, at the end

Tested against a direct softmax with N = 1,000 keys (the last block is partial), Br = 128 and d = 64: the maximum error is under 5.2e-15 non-causal, causal at tile offsets 0, 300 and 872, and with the largest scores planted in the last five keys so the max arrives late.

The guard on m_ref is the edge case. Under a sliding window or padding, a row can see zero keys in a block while its running max is still negative infinity. Then m - m_new is negative infinity minus negative infinity, which is NaN, and the NaN propagates into the accumulator forever. Referencing such rows to 0 gives alpha = 0 and P = 0, which is the correct empty contribution. Plain causal attention from position 0 never triggers it, because every row sees key 0 in the first block, so tests that only cover causal miss it.

Where the division goes is the FlashAttention-2 point. FA1 kept the output normalized and rescaled it by $\operatorname{diag}(\ell')^{-1}$ on every block; FA2 keeps it unnormalized and divides once:

tile Br = 128, d = 128, 8K keys in blocks of Bc = 64:
  key blocks per tile         = 8,192 / 64 = 128
  FA1 divides per block       = 128 x 128 = 16,384, repeated 128 times
  FA2 divides per tile        = 16,384, once
state per tile (FP32): A = 128 x 128 x 4 B = 64 KiB, m and l = 2 x 128 x 4 B = 1 KiB
score tile S_j (FP32)         = 128 x 64 x 4 B = 32 KiB

Non-matmul FLOPs run far slower than tensor-core FLOPs on Hopper, so removing per-block divides is real time, not tidiness. The 64 KiB accumulator is why kernels split a 128-row tile across warpgroups and keep it in registers rather than shared memory.

The reversal: in the backward pass the accumulator form is not needed. The forward saves the per-row log-sum-exp, and the backward rebuilds normalized $P_j = \exp(S_j - \text{LSE})$ directly with no running max.

What interviewers probe next

"Can you skip the rescale when the max did not change?" Yes, alpha is exactly 1 for those rows. FlashAttention-4 (arXiv 2603.05451) goes further and rescales only when the max grows past a threshold, because the final division by the matching sum cancels any common reference.

"Why is l kept in FP32 even when P is cast to BF16 for the GEMM?" l sums thousands of values below 1. BF16 has 8 bits of significand, so once l passes 256 an increment of 1 rounds away.

"Which loop order: queries outer or keys outer?" Queries outer, so each tile's state lives in registers for the whole sweep and O is written once. Keys outer forces the state to round-trip through HBM.

How it's scored

Weak: writes the scalar recurrence and loops over rows, or rescales l and forgets the accumulator.

Hire: writes the vector state, the diag(alpha) update and the final diag(l) inverse, and tests against a reference with a partial last block.

Strong Hire: handles the fully masked row, explains the FA2 deferred division with a count, sizes the register state, and connects the matrix form to which ops land on tensor cores.

Common mistakes

• Using one scalar max for the whole tile. It is still exact in exact arithmetic but underflows rows whose scores sit far below the tile max.
• Dividing by l inside the loop and then rescaling the normalized output by the wrong ratio.
• Testing only with N a multiple of the block size, which never exercises the partial last block.
• Leaving NaN handling to the final division, after the NaN has already reached the accumulator.