Click any words you think contain an error. Click again to unmark.
GQA requires num_attention_heads % num_key_value_heads == 0. 32 mod 7 = 4, so the reshape fails at load time and 7 was almost certainly a typo for 8.
Picture 32 students in a classroom that need to be split into equal teams sharing one shared textbook per team. If you ask for 8 teams, every team gets 4 students; clean and fair. If you ask for 7 teams, the math fails: you cannot hand the same 32 students to 7 teams of equal size, because 32 does not divide cleanly into 7 parts. GQA works the same way. The 32 query heads are the students, the KV heads are the shared textbooks, and each KV head must serve the same number of query heads. Pick a divisor of 32 (1, 2, 4, 8, 16, 32) and life is easy. Pick 7 and the reshape blows up the moment the model tries to run.
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.
This bug is one of the most common configuration mistakes in production LLM stacks, and the framing matters as much as the answer. The teammate did not pick 7 out of malice; they picked it because they were tuning the KV memory budget by hand and 7 KV heads at the right d_head landed close to the target cache size. The divisibility constraint is the architectural detail that vetoes the otherwise reasonable choice.
This deep dive walks the constraint itself, explains where exactly the error surfaces (and where it does not), surveys the standard production GQA configurations to anchor what 'looks normal', and ends with the practical CI test that prevents this class of bug from recurring.
Mental model: GQA is constrained by how the kernels reshape tensors, not by what the algorithm could express in principle. A custom kernel could allow unequal group sizes; production kernels do not, so you live with divisibility.
The divisibility constraint and why it exists
GQA partitions n_heads query heads into n_kv_groups groups, each sharing one K and one V projection. The shared projection lives at shape (d_model, n_kv_heads * d_head), and at attention time the kernel reshapes the query tensor from (B, T, n_heads, d_head) into (B, T, n_kv_heads, group_size, d_head) so the reduced K and V can broadcast across the group.
Why production kernels enforce equal group sizes
The reshape view(B, T, n_kv_heads, group_size, d_head) requires n_heads == n_kv_heads * group_size exactly. Both sides are integers, so n_kv_heads must divide n_heads. Every mainstream attention implementation, FlashAttention 2, FlashAttention 3, xFormers, vLLM's paged kernels, TensorRT-LLM, PyTorch's scaled_dot_product_attention, takes this layout for granted.
A custom kernel could in principle support unequal groups by scattering query heads into uneven buckets, padding to the next multiple, or implementing per-group masking. None of the production stacks do this because the engineering cost is real and the divisibility constraint is a non-issue for every standard configuration.
The 32 mod 7 = 4 problem
For n_heads = 32, the legal n_kv_heads values are exactly the divisors of 32: {1, 2, 4, 8, 16, 32}. 7 is not among them. The reshape from a 32-element axis into 7 equal sub-groups fails with a shape error.
The constraint is a property of the production kernel layout, not of the GQA algorithm. Knowing which is which is what distinguishes a senior debug from a junior one.
Situations where this technique stops working.
2–4 min · Everything important, quickly.
# Validation that should live in your model config loader.
def validate_gqa_config(n_heads: int, n_kv_heads: int) -> None:
if n_heads % n_kv_heads != 0:
raise ValueError(
f"GQA requires n_heads ({n_heads}) divisible by "
f"n_kv_heads ({n_kv_heads}). "
f"For {n_heads} query heads, legal n_kv_heads values are "
f"divisors of {n_heads}: "
f"{[d for d in range(1, n_heads + 1) if n_heads % d == 0]}"
)
# This catches the bug at PR review, not at the first forward pass.
validate_gqa_config(n_heads=32, n_kv_heads=7) # raises ValueError
validate_gqa_config(n_heads=32, n_kv_heads=8) # passes (group_size=4)Real products, models, and research that use this idea.
- Llama-3 70B (2024): 64 query heads, 8 KV heads, group size 8.
- Mistral 7B: 32 query heads, 8 KV heads, group size 4.
What an interviewer would ask next. Try answering before peeking at the approach.
QWhy did Llama-2 70B pick exactly 8 KV heads for 64 query heads rather than 4 or 16?
Empirical sweet spot from the GQA paper. Group size 8 (64/8) preserves over 99% of MHA quality on common evals while shrinking the KV cache by 8x. Group size 16 (64/4) loses meaningful quality at long context. Group size 4 (64/16) saves less KV memory without measurable quality gain. The choice is a Pareto point, not a theoretical optimum.
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.
Assuming the config loaded successfully because the model file opened, missing that the shape error only fires when the attention reshape executes during the first real forward pass.
60 second bullets to scan on the way to the call.
Primary sources. Browse if you want the original framing.
Same topic, related formats. Practice these next.