Single-Head Attention

Start from brute force

The textbook matrix version materializes everything: build the full score matrix S = Q·Kᵀ / √d, run exp on every entry and divide each row by its sum, then multiply the weights by V. It is the formula written out, so on small, tame inputs it is right, at O(n · m · (d + d_v)) time for n queries and m keys and O(n · m) space for S. Two things break it. Memory: for self-attention over 100,000 tokens, S alone holds 10¹⁰ scores. And plain exp overflows: math.exp(1000) raises OverflowError in Python, while JavaScript returns Infinity and the weights come out NaN.

The approach

Work one query at a time and stabilize the softmax. The reference attention sets d = len(Q[0]) and scale = math.sqrt(d), then for each query row q computes scores, one dot product per key divided by scale. It subtracts the row maximum m before exponentiating: softmax is unchanged when every score shifts by the same amount, the largest term becomes exp(0) = 1 so nothing overflows, and the sum z is at least 1 so the division is safe. The weights w = [e / z for e in exps] are non-negative and sum to 1.

Each output entry is sum(w[j] * V[j][c] for j in range(len(V))), a convex blend of the value rows, so every output column lies between that column’s smallest and largest value. The output width is len(V[0]), which need not equal d. Dividing by √d matters because the dot product of two vectors with independent unit-variance entries has variance d; scaling brings it back to 1 and keeps the softmax out of saturation. The TypeScript version runs the same three steps per row.

Walk through an example

The third example is the one where the weights do real work: Q = [[10, 0]], K = [[1, 0], [0, 1]], V = [[1, 1], [2, 2]], and the query matches key 0 far better than key 1.

Complexity

Time O(n · m · (d + d_v)) for n queries, m keys, key width d and value width d_v: each query scores m keys at O(d) each and blends m value rows into d_v columns. Working space is O(m) per query for scores, exps and w, plus O(n · d_v) for the output. The full-matrix version has the same time but O(n · m) space.

Edge cases to test

• Q = [[1000]], K = [[1], [0]], V = [[1], [2]] → [[1.0]]. The max shift turns the exponents into 0 and -1000; without it, math.exp(1000) raises OverflowError.
• Q = [[1, 5]], K = [[0, 0], [0, 0]], V = [[1, 2], [3, 4]] → [[2.0, 3.0]]. Identical keys give uniform weights and the mean of V, whatever the query.
• Q = [[1, 0]], K = [[0, 1]], V = [[7, 8, 9]] → [[7.0, 8.0, 9.0]]. The output width follows V, not d, and a single key always gets weight 1.
• Q = [[1, 1, 1, 1]], K = [[1, 1, 1, 1], [0, 0, 0, 0]], V = [[1], [0]] → [[0.8807970779778823]]. Without dividing by √4 the weight on the first key rises to 0.982: the scale softens the distribution.
• Q = [[1, 0], [0, 1]] against three keys [[1, 0], [0, 1], [1, 1]] with V = [[1], [2], [3]] → [[2.0], [2.203336278039358]]. There is one output row per query; the number of keys only sets the length of each weight vector.

Common mistakes

• Dropping the √d scaling — the dot products grow with dimension, saturate the softmax, and gradients vanish.
• Softmax over the wrong axis — it must normalize over the keys, per query.
• Forgetting the causal mask when the model is autoregressive (a token must not attend to the future).

Follow-ups an interviewer asks next

Add a causal mask for a decoder.

Query i may only attend to keys 0 through i. Set the scores of later keys to negative infinity before the softmax, so their weights are exactly 0, or in this loop simply score K[:i + 1]. The max shift still works, because every row keeps at least its own key, so m is finite and z is at least 1.

Turn it into multi-head attention.

Project the inputs with h separate learned matrices into h heads, each with key width d / h, run this function on every head independently, concatenate the h outputs, and apply one more learned output projection. The total work is about the same as one head of width d, but each head can learn a different pattern of what to attend to.

The sequence is 100,000 tokens long. What do production kernels do?

The full score matrix no longer fits in fast on-chip memory. FlashAttention tiles K and V into blocks and keeps a running maximum and running sum for every query, the online softmax, rescaling the partial output as each new block arrives. It never stores the whole matrix and still computes exact attention, not an approximation.

You generate text one token at a time.

Earlier tokens’ keys and values never change, so cache them. Each new token computes one query row against the cached K and V, costing O(n · d) per step instead of recomputing attention for the whole prefix. That is the KV cache, and because it grows with every token, its size is the main memory cost of serving long contexts.

Where this shows up

This is one head of the multi-head attention inside every Transformer, so implementing it from scratch is the canonical “do you actually understand attention” check. The O(n²) cost is also the direct reason long context is expensive and why sparse and linear-attention variants exist.

Reference solution

import math

def attention(Q, K, V):
    d = len(Q[0])
    scale = math.sqrt(d)
    out = []
    for q in Q:
        # 1. scaled dot-product scores against every key
        scores = [sum(q[t] * k[t] for t in range(d)) / scale for k in K]
        # 2. numerically stable softmax → weights that sum to 1
        m = max(scores)
        exps = [math.exp(s - m) for s in scores]
        z = sum(exps)
        w = [e / z for e in exps]
        # 3. weighted sum of the value rows
        dv = len(V[0])
        row = [sum(w[j] * V[j][c] for j in range(len(V))) for c in range(dv)]
        out.append(row)
    return out


print(attention([[1]], [[1]], [[5, 6]]))                 # -> [[5.0, 6.0]]
print(attention([[1, 5]], [[0, 0], [0, 0]], [[1, 2], [3, 4]]))  # -> [[2.0, 3.0]]

function attention(Q: number[][], K: number[][], V: number[][]): number[][] {
  const d = Q[0].length;
  const scale = Math.sqrt(d);
  const out: number[][] = [];
  for (const q of Q) {
    // 1. scaled dot-product scores against every key
    const scores = K.map((k) => {
      let s = 0;
      for (let t = 0; t < d; t++) s += q[t] * k[t];
      return s / scale;
    });
    // 2. numerically stable softmax
    const m = Math.max(...scores);
    const exps = scores.map((s) => Math.exp(s - m));
    const z = exps.reduce((a, b) => a + b, 0);
    const w = exps.map((e) => e / z);
    // 3. weighted sum of the value rows
    const dv = V[0].length;
    const row: number[] = [];
    for (let c = 0; c < dv; c++) {
      let acc = 0;
      for (let j = 0; j < V.length; j++) acc += w[j] * V[j][c];
      row.push(acc);
    }
    out.push(row);
  }
  return out;
}

console.log(attention([[1]], [[1]], [[5, 6]]));                       // -> [[5, 6]]
console.log(attention([[1, 5]], [[0, 0], [0, 0]], [[1, 2], [3, 4]])); // -> [[2, 3]]