Attention in forty lines of NumPy

The whole operation, runnable

The previous lesson described attention without matrices. Here it is with them, in code that runs on any laptop with NumPy installed and nothing else. Fifteen minutes with this file is worth more than any diagram.

import numpy as np

def softmax(x, axis=-1):
    x = x - x.max(axis=axis, keepdims=True)      # subtract max for numerical safety
    e = np.exp(x)
    return e / e.sum(axis=axis, keepdims=True)

def attention(X, Wq, Wk, Wv, causal=True):
    """X: (T, d) — one row per token position."""
    Q = X @ Wq          # (T, dk)   what each position is looking for
    K = X @ Wk          # (T, dk)   what each position advertises
    V = X @ Wv          # (T, dv)   what each position will hand over

    scores = Q @ K.T / np.sqrt(Q.shape[-1])      # (T, T)

    if causal:
        mask = np.triu(np.ones_like(scores), k=1).astype(bool)
        scores = np.where(mask, -np.inf, scores)

    W = softmax(scores)                          # (T, T) rows sum to 1
    return W @ V, W                              # (T, dv)

T, d, dk = 6, 8, 4
rng = np.random.default_rng(0)
X = rng.normal(size=(T, d))
out, W = attention(X, rng.normal(size=(d, dk)),
                      rng.normal(size=(d, dk)),
                      rng.normal(size=(d, dk)))
print(np.round(W, 2))

Run it. Print W. You will see a lower-triangular matrix whose every row sums to 1. That single object is the entire “which positions matter to which” story, and it is 6 × 6 for six tokens.

Three details that matter

The division by √dk. Dot products of vectors with dimension 64 have a standard deviation roughly eight times larger than those of dimension 1. Without the scaling, the scores spread wide, softmax saturates, and gradients vanish. Delete that division and re-run: the weight matrix becomes almost one-hot, each position attending to exactly one other. That is a live demonstration of why the term is there.

Subtracting the max in softmax. Purely numerical. exp(1000) overflows; exp(1000 - 1000) does not. Every production implementation does this and it changes no result.

-inf before the softmax, not zero after. Masking with negative infinity makes those entries exactly zero after normalisation, so the remaining weights still sum to one. Zeroing afterwards would leave rows summing to less than one, which quietly rescales every value.

Multiple heads

A head is one run of the above with its own three matrices. A layer runs several in parallel and concatenates the results, then applies one more projection:

def multihead(X, heads, Wo):
    outs = [attention(X, *h)[0] for h in heads]
    return np.concatenate(outs, axis=-1) @ Wo

In real implementations the heads are packed into single large matrices and split by reshaping, purely for speed. Conceptually it is the loop above.

Now watch the cost

Look at Q @ K.T. Its shape is (T, T). Set T to 1,000 and that is a million entries per head per layer. Set T to 100,000 and it is ten billion, per head, per layer. Time it yourself:

for T in (256, 512, 1024, 2048):
    X = rng.normal(size=(T, 512))
    import time; t = time.time()
    attention(X, *(rng.normal(size=(512, 64)) for _ in range(3)))
    print(T, round(time.time() - t, 3), "s")

Each doubling roughly quadruples the time. You have just measured, on your own machine, the quadratic cost that shapes the entire economics of long context.

What FlashAttention changes, and what it does not

Production systems do not build that (T, T) matrix in memory. FlashAttention computes the same result in tiles, keeping partial softmax statistics as it goes, so memory use grows linearly rather than quadratically and far less data moves between GPU memory and the compute units.

Be precise about the claim: it is an exact algorithm, not an approximation. The output is identical to the naive version up to floating-point rounding. The arithmetic is still quadratic; what became linear is the memory traffic, and on modern hardware memory traffic is usually the thing that costs you time. This is one of the clearest examples in the field of a large practical speed-up from changing nothing about the mathematics.

The thing worth keeping

Attention is a weighted average of value vectors, where the weights come from comparing a query against keys and pushing the results through a softmax. Everything else (heads, masks, scaling, caching, tiling) is engineering around that sentence. When you next read a paper claiming a new attention variant, you will be able to see immediately which part of these forty lines it is replacing.