Learning 11 · every intermediate visible
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.
One head · explicit axes
| stage | shape | axes |
|---|---|---|
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.
src/forward.rs: a complete score → mask → stable row-softmax → weighted-value trace, with no hidden vectors.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
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.
Architecture fidelity
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 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 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
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.
1/√d; RoPE does not replace the causal mask.src/forward.rs · ds4 flash_attn/softmax · HF modeling_qwen3.py.