Learning 11 · every intermediate visible

Attention, worked all the way through

Context: M2, one head → GQA → full verified block · Status: living

Attention is not vague “focus.” It is a content-addressed weighted read: Q asks what this token seeks, K says how each token can be matched, and V is what that token carries if selected.

Q, K, and V are different learned projections of the residual stream—not raw token vectors. Q and K determine weights; those weights combine V rows.
📖 Inference Engineering · §2.2.3🔧 ds4 · flash_attn / softmax🧭 Raschka · MHA / GQA / QK-norm

One head · explicit axes

Score, mask, normalize, read

stageshapeaxes
Q, K, V[seq,d]token position × head feature
S = QKᵀ/√d[seq,seq]query position × key position
causal mask, then A = row_softmax(S)[seq,seq]one distribution over visible keys per query row
O = AV[seq,d]query position × carried feature

Output row O[t,:] mixes value rows using A[t,:]. It does not shuffle token order: each query position still owns exactly one output row.

scaled scoresS = QKᵀ/√2 · [2,2][[0.7071, 0], [0, 0.7071]]Q = K = [[1,0],[0,1]]d = 2 → 1/√d = 0.7071 causal maskfuture j > t → −∞ · [2,2][[0.7071, −∞], [0, 0.7071]]t=0 cannot read j=1 row softmaxA · each row sums to 1 · [2,2][[1, 0], [0.33023846, 0.66976154]]subtract row max, then exp weighted VO = AV · [2,2][[10, 0], [3.3023846, 13.395231]]V = [[10,0],[0,20]] masksoftmax× V t=0 → v₀ only · t=1 → 0.33023846·v₀ + 0.66976154·v₁
matching scores legal prefix read weights carried output
The exact known-answer test in src/forward.rs: a complete score → mask → stable row-softmax → weighted-value trace, with no hidden vectors.

Why scaling and stable softmax matter

A dot product sums d terms, so its typical magnitude grows with head width. Dividing by √d keeps scores from driving softmax into saturation. Softmax subtracts each row's maximum before exponentiation: exp(xᵢ−max(x))/Σexp(xⱼ−max(x)). The probabilities are unchanged, but exponentials cannot overflow. A masked −∞ becomes weight zero.

Multi-head · grouped-query attention

Independent queries, shared K/V where specified

Projections pack heads along the feature axis. For Qwen3-0.6B: H=1024, 16 query heads, 8 KV heads, d=128. Thus Q is [seq,2048]; K and V are each [seq,1024]. Each query head runs attention independently. Its 16 outputs concatenate to [seq,2048], then o_proj → [seq,1024], back to H.

query headsshared K/V headsdistinct outputs q0q1q2q3 kv0shared K and Vkv1shared K and V o0o1o2o3 toy: group = 4/2 = 2 · Qwen: group = 16/8 = 2 · kv(h) = floor(h/2)
independent Q shared K/V independent result
Sharing K/V does not average query heads or their outputs. Different Q vectors score the shared keys differently, so every query head keeps its own weighted read.

Architecture fidelity

Qwen3's exact order

residual h
→ q/k/v projections → split heads
→ RMSNorm q and k per d (not V)
→ RoPE q and k (not V)
→ per-query-head causal attention with GQA sharing
→ concatenate → o_proj → H

QK-norm ≠ scaling

QK-norm is learned per-feature RMS normalization of Q/K vectors. 1/√d is fixed scaling of their dot products. They control different quantities, so both happen.

RoPE ≠ causal mask

RoPE adds relative-position information to Q/K matching. The mask forbids illegal future reads. Position-aware scores still need a visibility rule.

Bridge to code

From one visible head to a verified real block

attention_one_head computes the causal prefix and stable softmax shown above. multi_head_attention composes projections, QK-norm, RoPE, GQA, concatenation, and o_proj. The assembled real block 0 now matches all 5,120 official fp32 values, and the completed 28-layer pass matches every final logit. See M2 · forward pass → logits.

Common misconceptions