MQA and GQA: sharing keys and values

Sharing \(K\) and \(V\) changes both the learned projections and the state kept during decoding. Counting query heads as KV heads gives the wrong memory estimate.

derive
implement

You can now

Derive MHA, MQA and GQA projection shapes and bias-free parameter counts at matched model and head widths.
Implement causal grouped attention with an explicit query-to-KV-head mapping and differential tests.
Explain KV-memory savings without claiming an equal reduction in attention work or unchanged expressivity.

Covers

Multi-query and grouped-query attention shapes
MHA, MQA and GQA parameter and KV trade-offs
One small MHA versus GQA memory comparison

Study, in this order

2 steps

MHA, MQA and grouped-query implementation
Modern LLM Notebook - decoder construction path

Multi-Head Attention: MHA and GQA, the MHA vs GQA vs MQA KV Cache hand calculation comparison, and GQA implementation

Depth · Stop before QK-Norm. Skip MLA and the architecture survey. Read Exercise 3, How much KV Cache does GQA save, for practice.

Done
Contiguous grouped-head reference
PyTorch scaled dot product attention reference

Equivalent Python implementation: enable_gqa and repeat_interleave; GQA constraints note; Parameters and Shape legend

Depth · Use the head mapping and functional contract after your explicit grouped oracle passes. Skip backend and kernel tuning.

Done

Focus on

Let \(D=H_q \cdot d\) and \(H_{\mathrm{kv}}\) divide \(H_q\). \(Q\) is \([B, H_q, T, d]\); unexpanded \(K\) and \(V\) are \([B, H_{\mathrm{kv}}, T, d]\). MHA has \(H_{\mathrm{kv}}=H_q\); MQA has \(H_{\mathrm{kv}}=1\).
In PyTorch weight layout, \(W_q\) and \(W_o\) are \([D, D]\), while \(W_k\) and \(W_v\) are \([H_{\mathrm{kv}} \cdot d, D]\). The bias-free total is \(2D^2 + 2D \cdot H_{\mathrm{kv}} \cdot d\).
SDPA maps query head \(h\) to KV head \(\lfloor h/(H_q/H_{\mathrm{kv}}) \rfloor\). The notebook's expand-then-reshape instead interleaves KV heads; choose one convention and make weight ordering and the oracle agree.
Optional comparison: allocate unexpanded \(K\) and \(V\) at the same \(B\), \(T\), \(d\) and dtype for MHA and GQA, then compare numel()*element_size(). Do not count expanded temporary tensors as compact cache storage.

Reproduce closed-book

Implement a causal grouped-head loop and batched path for \(D=16\), \(H_q=4\) and \(H_{\mathrm{kv}}=2\), using contiguous groups. Derive 768 projection parameters versus MHA's 1024, and show that compact KV tensors use half the bytes at matched \(B\), \(T\) and dtype.

Traps

Using repeat rather than repeat_interleave while assuming contiguous groups. Shapes match, but queries read the wrong KV head.
Storing expanded K and V as the cache. That loses the compact representation's memory saving.
Expecting independently initialized MHA and GQA to produce equal outputs. Sharing restricts projections; equality needs a deliberately expanded shared-weight MHA reference.

Practice

Modern LLM Notebook: Exercise 3, How much KV Cache does GQA save

Exercises: Exercise 3, How much KV Cache does GQA save; complete the three quantities and pass its assertions

Extend the exercise in your own repository with the reproduce task. Compare outputs and shared-KV gradients against a head loop for \(H_{\mathrm{kv}}=H_q\), \(H_{\mathrm{kv}}=1\) and an intermediate group count. Reject invalid divisibility.

Self-test

With \(H_q=8\) and \(H_{\mathrm{kv}}=2\) under contiguous grouping, which query heads share KV head 1?
For \(D=1024\), \(H_q=16\) and \(H_{\mathrm{kv}}=4\), what are the \(K\) weight shape and total projection parameters?
Does four times fewer KV heads imply four times less \(QK\) attention work?
For \(B=1\), \(T=2048\), \(d=128\) and two-byte elements, what does changing 32 KV heads to 4 save per layer?

Completion criterion

Stop when

The named exercise passes, your grouped outputs and gradients match the explicit loop for all three head structures, and you derive compact KV bytes and projection counts without notes.

Mark passed, then open 2.5 Only mark this when you can do it without notes.

Leads to

2.7 Assembling and verifying a decoder Make the head structure an explicit model configuration choice.
5.4 KV memory arithmetic Extend the per-layer KV calculation across layers and sequences.

Comes after

2.3 Multi-head attention: split, merge and gradientsChange the number of KV heads without changing the query-head layout.

Resources for this unit

Modern LLM Notebook - decoder construction path
Walking Labsprimarycourse
Work through English notebooks 03-embedding, 04-position-encoding, 05-transformer-block, 06-mini-gpt and the selected normalization, RoPE, GQA and SwiGLU comparisons in 08-gpt2-to-modern-models. Re-derive shapes and rerun with changed dimensions.
PyTorch scaled dot product attention reference
PyTorchsupportreference
Use the functional contract only after the explicit attention oracle passes. Transfer identical weights and masks, compare outputs and gradients, and explain dtype, causal and dropout behavior from the documented semantics.

Also in this section

LLMs from Scratch companion - attention and GPT construction
Sebastian Raschkarequiredimplementation
Use the ch03 attention notebooks and ch04 model code. Rebuild the model without a high-level Transformer module and compare one forward and backward path with an independent small reference.
Stanford CS336 Assignment 1 - model contract
Stanford Universityrequiredassessment
Use the named linear, embedding, rmsnorm, positionwise_feedforward, rope, softmax, scaled_dot_product_attention, multihead_self_attention, transformer_block, transformer_lm and transformer_accounting specifications and public tests. Keep implementation independent and respect the course AI and honor-code boundary.
Deep-ML candidate - three transformer closing drills
Deep-MLoptionalpractice
After the platform quality sample passes, attempt at most problems 490, 53 and 904 as changed-prompt drills for attention shapes and multi-head composition. Preserve independent tests; a judge result does not replace the section checkpoint.
nanoGPT - residual projection initialization
Andrej Karpathysupportimplementation
Read model.py: MLP, Block, GPT.__init__'s special scaled initialization of c_proj.weight, and GPT._init_weights. Contrast the depth-dependent residual initialization with CS336's component initialization; do not adopt the training or generation code.