Problem 106 · Attention & Inference

Scaled Dot-Product Attention with a Causal Mask

MediumAI engineering

01Problem

Using only Python lists, implement `attention(Q, K, V, causal=False)` where `Q` is `Tq x d`, `K` is `Tk x d` and `V` is `Tk x dv`. Compute `softmax(Q K^T / sqrt(d)) V` row by row. With `causal=True`, query i (aligned to the **end** of the key sequence, so `Tq <= Tk` works for incremental decoding) may only attend to keys at positions `<= i + (Tk - Tq)`.

02Key insight

The 1/sqrt(d) scale keeps the dot products' variance near 1 as the head dimension grows; without it softmax saturates to one-hot and gradients vanish. The causal mask is applied as -inf **before** softmax, never as a zero after it, or the remaining weights do not sum to 1. Aligning queries to the end of the keys is the detail that lets the same function serve both full prefill (Tq = Tk) and one-token decode against a KV cache (Tq = 1).

03Solution

solution.pypython27 lines

import math

Matrix = list[list[float]]


def _softmax(row: list[float]) -> list[float]:
    m = max(row)
    exps = [math.exp(x - m) if x != -math.inf else 0.0 for x in row]
    total = sum(exps)
    return [e / total for e in exps]


def attention(Q: Matrix, K: Matrix, V: Matrix, causal: bool = False) -> Matrix:
    if not Q or len(K) != len(V) or len(Q[0]) != len(K[0]):
        raise ValueError("shape mismatch")
    scale = 1.0 / math.sqrt(len(Q[0]))
    offset = len(K) - len(Q)
    if causal and offset < 0:
        raise ValueError("causal attention needs Tq <= Tk")
    out: Matrix = []
    for i, q in enumerate(Q):
        scores = [sum(a * b for a, b in zip(q, k)) * scale for k in K]
        if causal:
            scores = [s if j <= i + offset else -math.inf for j, s in enumerate(scores)]
        w = _softmax(scores)
        out.append([sum(w[j] * V[j][c] for j in range(len(V))) for c in range(len(V[0]))])
    return out

Try it first.

TimeO(Tq * Tk * (d + dv))

SpaceO(Tk) per row plus O(Tq * dv) output

tests.py · 6 passing

tests.pypython18 lines

close = lambda A, B: all(abs(x - y) < 1e-12 for ra, rb in zip(A, B) for x, y in zip(ra, rb))
Q = [[1.0, 0.0], [0.0, 1.0], [1.0, 1.0]]
K = [[1.0, 0.0], [0.0, 1.0], [1.0, 1.0]]
V = [[1.0, 0.0], [0.0, 1.0], [5.0, 5.0]]
out = attention(Q, K, V, causal=True)
assert close([out[0]], [V[0]])
s = 1 / math.sqrt(2)
w = [math.exp(0.0), math.exp(s)]
assert close([out[1]], [[w[0] / sum(w), w[1] / sum(w)]])
full = attention(Q, K, V)
assert close([full[2]], [out[2]])
assert close(attention([[3.0, 4.0]], [[1.0, 1.0]] * 3, V), [[2.0, 2.0]])
assert close(attention([Q[2]], K, V, causal=True), [out[2]])
try:
    attention(Q, K[:2], V[:2], causal=True)
    assert False
except ValueError:
    pass

04Staff-level follow-up

Memory, not FLOPs, is the bottleneck at long context: the naive version materialises a Tq x Tk score matrix. FlashAttention tiles Q and K, uses the online softmax (running max and sum) and never writes the full matrix to HBM. Also be ready to explain sliding-window and ALiBi or RoPE as changes to the score matrix, not to this skeleton.