Note: FlashAttention, the same math computed without materializing the full attention matrix

The problem, from first principles

[03-attention-quadratic-complexity] showed that the fully materialized attention weight matrix, shape (seq_len, seq_len) per head, is itself a quadratic-memory cost, separate from the quadratic COMPUTE cost. For long sequences, simply HOLDING that matrix in memory (even briefly) can become the actual bottleneck, especially on GPUs, where moving data between slow, large "HBM" memory and fast, small on-chip "SRAM" is often slower than the arithmetic itself. Dao et al. (2022, "FlashAttention") observed that the FINAL attention output never actually needs the full matrix to exist all at once: it only needs, for every query, a running WEIGHTED SUM over keys/values seen so far, exactly the kind of computation that can be done incrementally, in CHUNKS, never holding more than one chunk's worth of scores in memory at a time.

The trick that makes this work correctly is "online softmax" (Milakov & Gimelshein, 2018): softmax needs to know the MAXIMUM score (for numerical stability, [01-scaled-dot-product-attention]'s - max trick) and the SUM of exponentiated scores, both computed over the ENTIRE row, before it can produce a single normalized weight. Processing keys in chunks means neither the true max nor the true sum is known until the LAST chunk. Online softmax solves this by maintaining a RUNNING max and running sum, and RESCALING every previously-accumulated partial result whenever a new chunk reveals a larger max than seen so far, mathematically producing the EXACT same final answer as computing the whole row at once, just incrementally.

From theory to code

Implement flash_attention(query, key, value, block_size, mask), computing EXACTLY [01-scaled-dot-product-attention]'s output, but processing key/value in chunks of block_size along the key axis: maintain a running max, running softmax denominator, and running (unnormalized) output, updating all three after each chunk via the online-softmax rescaling rule, and dividing by the final running denominator only at the very end.

Constraints

• Must produce EXACTLY the same output as computing full, unchunked attention, for ANY block_size from 1 up to seq_len (all inclusive).
• The running max must be updated with the TRUE maximum seen SO FAR (max(running_max, this_chunk's_max)), never a value from just the current chunk alone.
• Whenever the running max changes, every PREVIOUSLY accumulated quantity (the running sum, the running output) must be rescaled by exp(old_max - new_max) before adding the new chunk's contribution, or the result is mathematically wrong, not just imprecise.
• Never materialize a (seq_len_q, seq_len_k)-shaped score matrix; only ever a (seq_len_q, block_size)-shaped one, one chunk at a time.

Hints

block_max = scores.max(axis=-1, keepdims=True)
new_max = np.maximum(running_max, block_max)
correction = np.exp(running_max - new_max)     # rescales everything accumulated so far
probs = np.exp(scores - new_max)

running_sum = correction * running_sum + probs.sum(axis=-1, keepdims=True)
running_output = correction * running_output + probs @ value_chunk
running_max = new_max