Vibe Engines Learn
CODING CHALLENGE · N°08

Single-Head Attention

Medium AI EngineeringTransformersLLM

The one operation at the heart of every Transformer: scaled dot-product attention. Given queries, keys, and values, let each query pull a weighted blend of the values — weighted by how much its query matches each key. No numpy, just the math.

The problem

Implement scaled dot-product attention. You are given three matrices as lists of rows: Q (queries), K (keys) — each row length d — and V (values). For each query row, compute a score against every key as their dot product divided by √d, softmax those scores into weights that sum to 1, then return the weighted sum of the value rows. Return a matrix with one row per query and one column per value dimension: softmax(Q·Kᵀ / √d) · V.

EXAMPLE 1
Input Q=[[1]], K=[[1]], V=[[5, 6]]
Output [[5.0, 6.0]]
one key → weight 1 → the value itself
EXAMPLE 2
Input Q=[[1,5]], K=[[0,0],[0,0]], V=[[1,2],[3,4]]
Output [[2.0, 3.0]]
equal keys → uniform weights → the mean of V
EXAMPLE 3
Input Q=[[10,0]], K=[[1,0],[0,1]], V=[[1,1],[2,2]]
Output ≈[[1.0, 1.0]]
query matches key 0 → mostly V[0]
CONSTRAINTS
  • Scale the scores by 1 / √d, where d = len(Q[0]) is the key/query dimension.
  • Use a numerically stable softmax (subtract the row max before exp).
  • Each output row is a convex combination of the value rows — the attention weights sum to 1.
  • Pyodide has no numpy — use the math stdlib and plain lists only.
SOLVE IT YOURSELF

Your turn — write it

Edit the stub, hit Run (or ⌘/Ctrl + Enter), and watch the hidden tests. Stuck? the hints are right above and Reveal solution is one click away.

YOUR TASK

Implement attention(Q, K, V). For each query: score it against every key (dot product ÷ √d), softmax the scores, and return the weighted sum of the value rows.

HINTS — 4 IDEAS
  1. For query row q and key row k, the raw score is their dot product: sum(q[t] * k[t]).
  2. Divide every score by √d with d = len(Q[0]) — that is the "scaled" in scaled dot-product attention.
  3. Softmax each query row of scores into weights (subtract the max first for stability).
  4. Output row = Σ_j weight[j] * V[j] — a weighted blend of the value vectors, one entry per value column.
CPython · WebAssembly

Explore the topic

See this challenge alongside everything else on the same subject — handbooks, system designs, algorithms and tools, in one place.

LLM Engineering