Attention, Controlled · Paper 1 of the series

The Attention Variant Is the Only Variable:
A Controlled Study of MLA vs MHA, MQA, and GQA at 124M Parameters

1  Introduction

Autoregressive decoding stores one key and one value vector per attended token, per layer—the KV cache. Its size grows linearly with context length and, at long context, dominates both the memory budget and the memory-bandwidth cost of each decode step. A 124M MHA model in this study needs 38.6 GB of KV cache to decode a single 128k-token context at batch 8—more than twice the 16 GB of the GPU it was trained on. The attention mechanism is therefore the lever that decides whether long-context inference is feasible on commodity hardware at all.

MQA and GQA shrink the cache by sharing key/value heads, trading representational capacity for memory. MLA takes a different route: it caches a low-rank latent from which per-head keys and values are reconstructed on the fly, plus a small decoupled key that carries rotary position information. The open question this study targets: at a fixed compute and parameter budget, how much quality does each cache-reduction strategy actually cost, and what does each buy back in inference efficiency?

Public comparisons usually confound the attention mechanism with changes in depth, width, data, or training recipe. We remove that confound. Every variant here shares an identical 12-layer / 768-width / 12-head backbone, identical RoPE, identical SwiGLU MLP, identical optimizer and schedule, and the same token budget. The attention module is the sole independent variable.

2  Background: four ways to cache

All variants share the same per-head dimension $d_h = d_\text{model}/n_h = 64$ and the same RoPE applied to queries and keys. They differ only in what gets cached. With $L$ layers and 2-byte (bf16) elements, the KV cache per token is:

MHA caches a full key and value for each of $n_h$ heads. MQA collapses these to a single shared KV head, and GQA interpolates with $g$ groups. MLA instead caches a compressed latent $\mathbf{c}^{KV}\!\in\!\mathbb{R}^{d_c}$ ($d_c=256$ here) and a single decoupled RoPE key $\mathbf{k}^{R}\!\in\!\mathbb{R}^{d_\text{rope}}$ ($d_\text{rope}=16$), shared across heads—so its cache width is $d_c + d_\text{rope}$ regardless of head count.

3  Multi-head Latent Attention

MLA (DeepSeek-V2) factors the KV path through a low-rank bottleneck. From hidden state $\mathbf{h}_t$ it computes a down-projection to the cached latent and a decoupled rotary key:

$$\mathbf{c}^{KV}_t = W^{DKV}\mathbf{h}_t \in \mathbb{R}^{d_c}, \qquad \mathbf{k}^{R}_t = \mathrm{RoPE}(W^{KR}\mathbf{h}_t) \in \mathbb{R}^{d_\text{rope}}.$$

Only $(\mathbf{c}^{KV}_t, \mathbf{k}^{R}_t)$ are cached. At attention time the latent is up-projected to per-head content keys and values, $[\mathbf{k}^{C}; \mathbf{v}] = W^{UK,UV}\mathbf{c}^{KV}$, and each head's key is the concatenation $[\mathbf{k}^{C}_i; \mathbf{k}^{R}]$ of its content part and the shared rotary part. Queries are split the same way into a non-rotary "content" part ($d_\text{nope}=48$) and a rotary part ($d_\text{rope}=16$).

3.1  Weight absorption: a fair decode path

Naively, decoding MLA re-materializes per-head keys/values for the entire context at every step—which would erase its memory-bandwidth advantage. The fix is an algebraic re-association. For the content score of head $i$,

$$\mathbf{q}^{C\top}_i \mathbf{k}^{C}_i = \mathbf{q}^{C\top}_i \big(W^{UK}_i \mathbf{c}^{KV}\big) = \big(W^{UK\top}_i \mathbf{q}^{C}_i\big)^{\!\top} \mathbf{c}^{KV},$$

so the up-projection $W^{UK}_i$ can be "absorbed" into the query once per step and the score taken directly against the cached latent $\mathbf{c}^{KV}$—the per-head content key is never formed over the context. The value projection $W^{UV}_i$ absorbs symmetrically into the output. We verify numerically (float64, math SDPA backend) that this absorbed path is identical to the naive one to $<\!10^{-9}$, and use it for every MLA decode measurement below.

4  Method

• Backbone (fixed): 12 layers, $d_\text{model}=768$, $n_h=12$, $d_h=64$, SwiGLU MLP ($\tfrac{8}{3}$ expansion), RMSNorm, tied embeddings, RoPE $\theta{=}10^4$, vocab 50,304, context 1024.
• Variant knobs: GQA $g=4$; MLA $d_c=256$, $d_\text{nope}=48$, $d_\text{rope}=16$, $d_v=64$, no query compression.
• Honest parameter counts. We do not fake parameter-matching; the backbone is identical and we report exact counts (MHA 123.6M, GQA 114.2M, MLA 116.1M, MQA 110.6M—the spread comes only from the differing K/V/-projection widths) and FLOP-normalize in analysis.
• Training (fixed): AdamW, cosine schedule with warmup, bf16 autocast, gradient accumulation, identical token budget per variant. torch.compile on. Same seed.
• Correctness: KV-cache incremental decode is checked against a full forward pass to $10^{-9}$ for every variant; causality (no future leak) is asserted; MLA's absorbed path is checked against its naive path.

5  Efficiency results

All measurements: RTX 4080 Super, bf16, decode batch 8, single token/step against a pre-filled context. Decode uses MLA's weight-absorbed path.

5.1  KV-cache footprint

The cache is exact and grows linearly with context. At a 128k context (batch 8) the total KV footprint is MHA 38.65 GB, GQA 12.88 GB, MLA 6.85 GB, MQA 3.22 GB. MHA is physically impossible on this 16 GB card; MLA fits with room to spare while caching a single shared latent rather than four KV groups.

5.2  Decode latency, memory, and throughput

Because MLA attends against a narrow latent, its per-step cost stays nearly flat as context grows from 512 to 2048 (4.68 → 4.65 ms), where MHA already climbs (3.29 → 5.84 ms). At 16k context MLA decodes at 359 tok/s using 3.0 GB, versus MHA's 258 tok/s at 11.1 GB and GQA's 269 tok/s at 4.3 GB. There is a crossover: at short context MLA pays a small absorption overhead and trails the others; at long context it overtakes both MHA and GQA on throughput while using a fraction of their memory.

5.3  Prefill and training throughput

With no cache to amortize, prefill and training throughput are close across variants—within ~20% (prefill 250–308k tok/s; training 54–58k tok/s), with the lighter-KV variants slightly ahead. Cache reduction is an inference-time win, not a training-time cost.

6  Quality results

6.1  TinyStories (dev scale, ~50M)

At dev scale the four variants land within 0.024 val loss (0.10 perplexity, ~2.4%) of one another. MHA leads, as expected from the largest and most expressive cache; the three cache-reduced variants are separated by less than run-to-run noise (MQA 4.258, GQA 4.284, MLA 4.302). The practical reading is that cutting the KV cache 4–8× costs essentially nothing in quality at this token budget. MLA matches the GQA/MQA quality band while—unlike them—reconstructing full per-head keys and values; whether that retained head diversity converts into a quality edge is a question for the 124M/10B headline scale (§6.2), where 50M/4k-iter under-training no longer masks it. We do not over-read the ordering among the three cheap variants here.

6.2  FineWeb-Edu (headline, 124M)

Headline sweep: identical 12-layer / 768-wide backbone, 20,000 iterations (491,520 tokens/step ≈ 9.8B tokens, ~1 epoch of FineWeb-Edu 10B), identical AdamW / cosine schedule and seed; only --attn varies. Evaluated on the held-out split (47.8M tokens) from each variant's best checkpoint.

At headline scale the four variants span just 0.045 val loss (20.3–21.3 perplexity, ~4.6%), and the ordering sharpens into a clean result. MHA leads, as its full-width cache should. But among the three cache-reduced variants MLA wins (val loss 3.0309, vs GQA 3.0349 and MQA 3.0578) while caching less than GQA—6,528 vs 12,288 bytes/token. MLA therefore Pareto-dominates GQA: strictly smaller cache and lower perplexity. It also closes most of the gap to MHA, trailing the full-cache model by only ~1.8% perplexity while using 5.6× less KV memory. The head-diversity question left open at dev scale (§6.1)—whether MLA's reconstructed per-head keys/values buy quality the head-sharing variants cannot—resolves in MLA's favor once 50M / 4k-iter under-training is removed.

The four variants trained concurrently, one per GPU, on NVIDIA H100/H200 GPUs of Georgia Tech's PACE-ICE cluster. Tokens/step is held identical across variants, so the comparison stays controlled regardless of which GPU each variant landed on.

7  Discussion

The efficiency picture is unambiguous and hardware-independent in its ordering: for a fixed backbone, the cached width sets the long-context memory wall, and MLA buys a 5.6× reduction over MHA while—unlike MQA/GQA—reconstructing full per-head keys and values, so it need not give up head diversity to save memory. The weight-absorption identity is what turns that smaller cache into an actual decode-throughput win at long context rather than a wash. The decisive question is whether that reconstruction preserves quality at a fixed token budget. At headline scale (§6.2) it does: MLA is the best of the cache-reduced variants and the only one to Pareto-dominate GQA, recovering most of MHA's quality at a fraction of its cache. The retained per-head key/value diversity—not merely the smaller cache—is what separates MLA from MQA/GQA once the token budget is fixed and the under-training of dev scale is removed.

8  Reproducibility

Everything is open and runs on one consumer GPU. Code: github.com/bryanvine/mla-gpt.

uv sync
uv run pytest                                   # correctness: cache equivalence, causality, absorption
uv run python scripts/benchmark.py --max-context        # §5 efficiency + figures
bash scripts/sweep_tinystories.sh               # §6.1 dev-scale quality sweep
uv run python scripts/eval.py --glob 'runs/tinystories_*' --out runs/tinystories_eval
sbatch scripts/sweep_fineweb_pace.sbatch        # §6.2 headline sweep (4×GPU, PACE-ICE/Slurm)
uv run python scripts/eval.py --glob 'runs/fineweb_*' --out runs/eval_fineweb

A single base config trains every variant by changing only --attn, which is what keeps the attention mechanism the sole independent variable.

9  Acknowledgments

The 124M headline sweep (§6.2) was run on PACE-ICE, the instructional and research computing cluster at the Georgia Institute of Technology. We gratefully acknowledge Georgia Tech PACE-ICE for the research GPU time (NVIDIA H100/H200) that made headline-scale training feasible. All dev-scale work (§5 efficiency and the §6.1 sweep) runs on a single consumer RTX 4080 Super.

Built from scratch in PyTorch. MLA follows DeepSeek-V2 (Liu et al., 2024); GQA follows Ainslie et al. (2023); the byte/RoPE backbone follows the Llama/nanoGPT lineage. © 2026 Bryan Vine, MIT-licensed.