Mini FlashAttention

1What attention is

Attention is how a transformer token looks at the other tokens and decides what to copy into its own representation. Three matrices per head:

• Q (query) — “what am I looking for?” (your token’s projection)
• K (keys) — “what do I contain?” (every token’s projection)
• V (values) — “what can I hand over?” (every token’s payload)

The math, in three lines:

S = Q @ K^T * scale        # scores: how much each query wants each key   (N×N)
P = softmax(S, rows)        # probabilities: lesson 3's dance, per row      (N×N)
O = P @ V                   # outputs: weighted mix of the values           (N×d)

With scale = 1/√d keeping the scores’ variance sane. In autoregressive decode, the causal mask keeps S[i, j] = −∞ for j > i — a token never attends to the future.

2Why fuse it

The unfused path (torch) writes the N×N scores matrix to HBM, reads it back for the softmax, writes P, reads P back for the PV product. For a 4096-token sequence that’s a 4096² × 4 bytes = 64 MB round trips per head, times 32 heads, times 32 layers — and the scores matrix is exactly as expensive as the tokens it describes. FlashAttention’s entire thesis: never write S to memory. Compute a tile of scores, softmax it in registers, multiply by V, keep only the output. The kernel below does exactly that — the full sequence in one tile (Part A), then the real flash version that chunks K/V and streams (Part B).

3One-shot fused attention

Everything you know, in one kernel. Every lesson is in here — spot them all:

@triton.jit
def attention_kernel(q_ptr, k_ptr, v_ptr, o_ptr, NSEQ, HEAD, scale,
                     BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_D: tl.constexpr,
                     CAUSAL: tl.constexpr):
    pid_m = tl.program_id(0)
    offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)   # rows of Q I own (L2)
    offs_n = tl.arange(0, BLOCK_N)                     # all sequence positions
    offs_d = tl.arange(0, BLOCK_D)                     # head dim

    q = tl.load(q_ptr + offs_m[:, None] * HEAD + offs_d[None, :],
                mask=offs_m[:, None] < NSEQ, other=0.0)
    k_t = tl.load(k_ptr + offs_n[:, None] * HEAD + offs_d[None, :],
                  mask=offs_n[:, None] < NSEQ, other=0.0)   # (BLOCK_N, HEAD)
    v = tl.load(v_ptr + offs_n[:, None] * HEAD + offs_d[None, :],
                mask=offs_n[:, None] < NSEQ, other=0.0)     # (BLOCK_N, HEAD)

    s = tl.dot(q, tl.trans(k_t)) * scale               # (BLOCK_M, BLOCK_N) scores
    if CAUSAL:
        s = tl.where(offs_m[:, None] >= offs_n[None, :], s, -float("inf"))

    m = tl.max(s, axis=1)                              # lesson 3, in miniature
    m = tl.where(offs_m < NSEQ, m, 0.0)                # ghost-row pin
    p = tl.exp(s - m[:, None])
    l = tl.sum(p, axis=1)
    o = tl.dot(p.to(tl.float32), v) / l[:, None]       # P @ V, normalized
    tl.store(o_ptr + offs_m[:, None] * HEAD + offs_d[None, :], o,
             mask=offs_m[:, None] < NSEQ)

4The flash trick: online softmax

Part A holds the whole sequence in one tile — fine for N ≤ 128, impossible for N = 4096. The real FlashAttention streams K/V in BLOCK_N-column chunks, and that creates a problem: softmax needs the whole row’s max and sum, but each chunk arrives separately. The fix is the paper’s key idea — keep a running state per row and rescale when the max moves:

# running state per row:
m_i = running max score        # starts at -inf
l_i = running sum of exp()     # starts at 0
acc = running P @ V            # starts at 0

for each K/V chunk:
    s = Q @ K_chunk^T * scale
    m_new = max(m_i, max(s))               # maybe a bigger score arrived
    alpha = exp(m_i - m_new)               # rescale factor for OLD state
    p = exp(s - m_new)                     # new chunk, stable
    l_i = l_i * alpha + sum(p)             # old sum re-weighted
    acc = acc * alpha + p @ V_chunk        # old output re-weighted
    m_i = m_new

O = acc / l_i                              # after the last chunk

The companion’s Part B implements exactly this loop (plus the ghost-row guard: m_new = tl.where(m_new == -inf, 0.0, m_new), so fully-masked rows can’t NaN). Verify it against torch — it matches to 1e-7 on ragged sequences too.

5Attention playground

Watch the three-line math happen on real numbers. Click a score cell to trace it through the softmax row; flip the causal mask and watch the future disappear (× = −∞):

6The surgery map, complete

You started with “why GPUs exist”. Here is where the course lands — decode-lab’s surgery map, with the lessons that unlock each row:

7Homework — the last set

Run lesson7_flash_attention.py on Colab, then try each of these. Hints are collapsible — try it first, peek only if stuck.

8The whole journey