Implement FlashAttention-style tiled attention

The key identity: softmax can be computed incrementally. For each query row, sweep K/V tiles while maintaining a running max m, running normalizer l, and running output o, rescaling previous partial results when a new max appears.

import numpy as np

def flash_attention(Q, K, V, tile=128):
    N, d = Q.shape
    O = np.zeros((N, d))
    m = np.full(N, -np.inf)   # running row max
    l = np.zeros(N)           # running softmax denominator
    for j in range(0, N, tile):
        Kj, Vj = K[j:j + tile], V[j:j + tile]
        S = Q @ Kj.T / np.sqrt(d)              # (N, tile) — only a tile, never N×N
        m_new = np.maximum(m, S.max(axis=1))
        scale = np.exp(m - m_new)              # rescale old accumulators
        P = np.exp(S - m_new[:, None])
        l = l * scale + P.sum(axis=1)
        O = O * scale[:, None] + P @ Vj
        m = m_new
    return O / l[:, None]

Explain the systems framing: standard attention is O(N²) HBM traffic; tiling keeps the working set in SRAM, so FlashAttention is an IO-aware algorithm — same FLOPs, far less memory movement, exact (not approximate) result.

Follow-ups: Also tile the queries (the real algorithm is a 2D loop) and handle the causal mask per tile. Why is attention memory-bound rather than compute-bound on modern GPUs? Backward pass: why does FlashAttention recompute the attention matrix instead of storing it? How does this interact with KV caching at inference — what's different in decode (one query row) vs prefill?