Home / Fundamentals / Transformer Architecture

Transformer ArchitectureSeniorQuestion 3 of 9

Compare MHA, MQA, GQA and MLA. What problem do they solve?

Try answering out loud before reading on. The walkthrough below is paced for a 40-minute interview round.

0–5 min Clarify and frame

Framing: "These variants differ in how many key/value heads exist and what is cached. That number controls KV-cache size, which bounds batch size and context length at inference. I'll define each, quantify memory, and discuss quality and systems trade-offs, including MLA."

5–15 min Core explanation

MQA and MHA are the extremes of GQA (G = 1 and G = H). In GQA, each group of H/G query heads reads the same K and V.

15–25 min Deep dive

Memory. With L = 80 layers, H = 64, d_h = 128, bf16:

MHA: 2 · 80 · 64 · 128 · 2 bytes = 2.62 MB per token  → 32k tokens ≈ 86 GB per sequence
GQA (G = 8):                        0.33 MB per token  → 32k tokens ≈ 10.7 GB
MQA:                                0.04 MB per token  → 32k tokens ≈ 1.3 GB

With MHA a single long sequence wouldn't fit beside the weights on one GPU; with GQA it does, and batch size rises accordingly.

Bandwidth. Decode reads the whole KV cache every step. Fewer KV heads means fewer bytes per token generated, so lower latency and higher throughput at long contexts.

Quality. MQA forces all heads to use the same keys and values, reducing representational capacity and sometimes stability. GQA with 4–8 groups recovers nearly all of MHA's quality.

Converting checkpoints. An MHA model can be "uptrained" to GQA by mean-pooling the K and V projection weights within each group, then training briefly on a small fraction of the original compute.

MLA (multi-head latent attention). Compress each token's keys and values into a low-dimensional latent c (e.g. a few hundred dims) that is cached, and up-project it to per-head K and V when attending. Up-projections can be absorbed into the query and output projections, so you don't materialize full K/V at decode time. Because RoPE doesn't commute with that absorption, MLA keeps a small separate positional key component. The result is a cache smaller than GQA's with quality comparable to MHA, at the cost of extra compute and implementation complexity.

25–32 min Trade-offs

• Tensor parallelism: with very few KV heads (MQA), splitting heads across GPUs means replicating K/V on each device; GQA's group count is often chosen to divide evenly by the tensor-parallel degree.
• Training cost: fewer KV heads slightly reduce training FLOPs and parameters; the main benefit is inference.
• Alternatives for the same problem: KV-cache quantization (FP8/INT8), sliding-window layers, cross-layer KV sharing, paged memory management. These combine with GQA.

32–37 min Production example

A team serving a 70B-class model found throughput capped by KV memory, not compute. Moving to a similar-quality GQA model (8 KV heads) and FP8 KV cache allowed roughly 4–8× more concurrent sequences at the same latency target, which cut cost per token more than any kernel optimization.

37–40 min Likely follow-ups

• Why not reduce query heads too? Queries aren't cached; reducing them saves little memory and costs quality.
• How is GQA implemented efficiently? Kernels index the shared K/V per group rather than copying them for each query head.
• Does GQA help prefill? Slightly; prefill is compute-bound and dominated by matmuls, so the benefit is mainly in decode.

What interviewers listen for

• Correct definitions with MHA and MQA as GQA's extremes.
• KV-cache arithmetic with realistic numbers, and the link to batch size and bandwidth.
• Knowledge of MLA's latent caching and the RoPE complication.

Red flags

• "GQA reduces the number of query heads."
• Not connecting the change to inference memory.