FlashAttention tiled forward

Background

Standard attention materializes the full $N\times N$ score matrix, costing $O(N^2)$ memory. FlashAttention computes the exact same result without ever storing that matrix: it streams over blocks of keys/values and maintains an online softmax — a running max $m$, a running normalizer $\ell$, and a running output $O$ — rescaling the accumulators each time a new block shifts the max. This is the trick behind training long-context transformers efficiently.

Problem statement

Implement flash_attention(Q, K, V, block_size) returning the attention output, computed by tiling over key/value blocks with online softmax. For each key block $S = \frac{QK_b^\top}{\sqrt d}$:

$$m^{\text{new}} = \max(m,\ \text{rowmax}(S)), \quad \alpha = e^{\,m - m^{\text{new}}}, \quad P = e^{\,S - m^{\text{new}}}$$ $$\ell \leftarrow \alpha\,\ell + \text{rowsum}(P), \qquad O \leftarrow \alpha\,O + P V_b$$

After all blocks, divide: $O \leftarrow O/\ell$. Use scale $1/\sqrt{d}$.

Input

• Q, K, V — np.ndarray of shape (N, d) (single head, no batch).
• block_size — int, number of keys processed per tile.

Output

An np.ndarray of shape (N, d): identical to standard softmax attention $\operatorname{softmax}(QK^\top/\sqrt d)V$.

Examples

Example 1

Input:  Q = [[1, 0]], K = [[1, 0], [0, 1]], V = [[1, 2], [3, 4]], block_size = 1
Output: [[1.6604, 2.6604]]

Explanation: scores are $[1,0]/\sqrt2 = [0.707, 0]$; softmax gives weights $[0.670, 0.330]$; the output is $0.670\,[1,2] + 0.330\,[3,4] = [1.660, 2.660]$ — the same whether computed in one block or tiled.

Constraints

• Initialize the running max to $-\infty$ and the normalizer/output to 0; the first block's $\alpha = 0$ correctly discards the empty initial state.
• Rescale both $\ell$ and $O$ by $\alpha = e^{m - m^{\text{new}}}$ before adding the new block's contribution.
• The result must match standard attention for any block_size in $1..N$.

Notes

• The online-softmax rescaling is what keeps the result exact: subtracting the running max prevents overflow, and $\alpha$ corrects earlier terms once a larger max appears.
• Real FlashAttention also tiles the queries and fuses everything into one GPU kernel; the math here is the same, only the loop nest differs.