Compute the head dimension a config file implies for d_model 4096 and 32 heads
You open a model config: {"hidden_size": 4096, "num_attention_heads": 32}. The loader splits the hidden dimension evenly across heads. Compute head_dim, the per head dimension each Q, K, and V vector will have.head_dim = d_model / n_heads = 4096 / 32 = 128. Each attention head works in its own 128-dimensional slice of the projected Q, K, and V.
Imagine the model's hidden state at one position is a 4096-element vector, like 4096 sticky notes laid out in a row. Multi-head attention gives each of 32 heads its own pile of 128 notes. Each head reads the whole row and writes its own Q, K, and V into its pile. The arithmetic is 4096 divided by 32, which is 128. That number, 128, is the per-head dimension. It is also the dimension of every Q, K, and V vector inside the head. Many modern open models like Llama 3 chose this exact split because 128 is wide enough to be expressive and narrow enough that the dot-product attention scales nicely.
Concept explanation~2 min read
Everything you need to truly understand this topic: intuition, mechanics, step by step explanation, code, formulas, and worked example. Click to expand.
Concept explanation~2 min read
Everything you need to truly understand this topic: intuition, mechanics, step by step explanation, code, formulas, and worked example. Click to expand.
The arithmetic is one line: 4096 divided by 32 is 128. Every modern config file follows this pattern, and the number controls the shapes inside multi-head attention.
The contract: d_model equals n_heads times head_dim
Multi-head attention slices the residual stream into n_heads parallel subspaces. The standard contract is:
With d_model 4096 and n_heads 32, head_dim is exactly 128. The contract exists because the per-head attention outputs are concatenated back into a single d_model-shaped vector before the output projection W_O. If n_heads x head_dim did not equal d_model, the concat would not produce a d_model vector and the residual add at the end of the sublayer would fail.
Mechanically: the Q, K, V projection matrices have shape [d_model, n_heads x head_dim] = [4096, 4096]. Each projection produces a tensor of shape [batch, seq, n_heads, head_dim] after reshape. Attention runs per head in the 128-dim subspace. The concatenated output is [batch, seq, n_heads x head_dim] = [batch, seq, 4096], which feeds W_O of shape [4096, 4096] and lands back on the residual stream.
Most frameworks (PyTorch, JAX, vLLM) compute head_dim as hidden_size // num_attention_heads automatically when loading a config. Some configs ship an explicit head_dim field that may override this, but for vanilla LLM configs the formula is what gets used.
Situations where this technique stops working.
2–4 min · Everything important, quickly.
Real products, models, and research that use this idea.
- Llama 3 8B: d_model 4096, n_heads 32, head_dim 128.
- Llama 3 70B: d_model 8192, n_heads 64, head_dim 128.
What an interviewer would ask next. Try answering before peeking at the approach.
QWhy does head_dim appear inside a square root in scaled dot-product attention?
The dot product of two d-dim vectors with unit-variance components has variance d. Dividing by sqrt(d) restores unit variance for the pre-softmax logits, keeping the softmax in its non-saturated regime. Without the scaling, deep stacks with bf16 see one-hot attention by layer 4 or so.
Red flags & common mistakes
The phrases that signal junior thinking. Click to expand.
Red flags & common mistakes
The phrases that signal junior thinking. Click to expand.
Confusing head_dim with d_model itself or with the FFN intermediate size. head_dim is strictly the per-head slice, equal to d_model divided by n_heads when the split is uniform.
60 second bullets to scan on the way to the call.
The formula d_model = n_heads x head_dim and why it must hold
How head_dim shows up in the Q, K, and V projection shapes
Primary sources. Browse if you want the original framing.
Same topic, related formats. Practice these next.