Self-attention is the mechanism at the center of every transformer: for each position, ask "how relevant is every other (earlier) position to me?", turn those relevance scores into weights, and produce an output as a weighted blend of every position's value vector.

Q, K, V are each a list of n vectors (n = sequence length). dot and softmax are done for you.

Implement attention(Q, K, V, causal=True, pad_mask=None):

• For query position i, its score against key position j is dot(Q[i], K[j]) / sqrt(d_k) (d_k = the vector dimension).
• If causal is True, position i may only attend to j <= i -- mask later positions to float("-inf") before softmax (softmax of -inf correctly contributes 0 weight).
• Turn the (masked) scores into weights with softmax.
• Position i's output is the weighted sum of V, using those weights.

The missing link: masking isn't just the causal rule

A single boolean (causal) is enough for a toy example, but real inference almost always batches multiple sequences of DIFFERENT lengths together, padded to a common length so they fit in one tensor. Those padding slots aren't real tokens -- they must never receive attention weight, from ANY query position, regardless of the causal rule. That's pad_mask: a list of bool (True = real token, False = padding), applied as an INDEPENDENT masking condition alongside the causal one. A position is masked if it fails causality OR if it's padding -- both conditions matter, and an implementation that only checks one of them will silently attend to garbage padding vectors the moment a real batch has mixed-length sequences.

Example

With an all-zero Q, every dot product is 0 regardless of K, so weights are uniform over whatever's visible:

Q = [[0], [0], [0]]
V = [[10.0], [20.0], [30.0]]
attention(Q, K, V, causal=True)[1][0]  # -> 15.0  (average of V[0], V[1])

Constraints

• The scaling by sqrt(d_k) happens BEFORE softmax, not after.
• Masked-out positions (causal OR padding) must contribute exactly 0 weight, not just a small one.
• pad_mask=None (the default) must behave exactly as if no padding masking applies at all -- only causal masking, same as before this parameter existed.

Real-world grounding

• GPT-2's attention implementation (a clean, readable reference): src/transformers/models/gpt2/modeling_gpt2.py
• nanoGPT's attention (even more minimal): model.py

Real implementations operate on batches of matrices via a tensor library (PyTorch, GGML) instead of Python loops over lists, and split Q/K/V into multiple "heads" computed in parallel -- but the core per-position computation is exactly what you just implemented.