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. Predict 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 sees a 128-dimensional slice of the residual stream.
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 divides those notes into 32 equal piles, one per head, with 128 notes each. Each head only looks at its own pile when computing Q, K, and V. The arithmetic is simply 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. The interesting part is why every modern config file you encounter follows this pattern and what the number actually controls inside multi-head attention.
This walkthrough derives head_dim from first principles, traces it through the projection shapes, shows where it enters the softmax scaling factor, and ends with how GQA and MLA partially break the formula while preserving its spirit.
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.