Agent MoE Experiment: MHA (no GQA)

[ClassicLarry]

TL;DR

Test full multi-head attention (num_kv_heads = num_heads) vs baseline 4:1 GQA. Measures the quality cost of GQA at current compute budgets with the current recipe.

User prompt

follow agent.md and implement multi-head-attention, (no gqa)

Scope

• Parent: MoE Scaling up to April goal #4281
• Branch: moe_mha
• Sweep file: experiments/grug/moe/mha_sweep.py
• Prior work: Experiment: Test Loss impact of GQA #4371 found ~+0.004 BPB per halving of KV heads on an older recipe

Head counts

d	GQA (baseline) heads/kv	MHA heads/kv
512	4/1	4/4
768	6/1	6/6

Gate 1 runs (2 total)

Dim	Budget	Status
512	2.19e17	pending
768	1.70e18	pending

Decision log

empty

Conclusion

pending

[ClassicLarry]

🤖 Gate 1 Results: MHA — no GQA (macro_loss)

Both runs finished. Strong improvement from removing GQA.

Dim	Run	macro_loss	tps	tps ratio	Eff Speedup
512	baseline (kv=1)	3.8104	405,630	—	—
512	MHA (kv=4)	3.7868 (−0.024)	396,878	0.978	1.097
768	baseline (kv=1)	3.4339	273,532	—	—
768	MHA (kv=6)	3.3999 (−0.034)	263,894	0.965	1.177

MHA costs ~2-4% throughput (more KV params) but the quality gain is substantial. Prior GQA ablation (#4371) found +0.008 BPB for 6:1 GQA; this confirms the macro_loss cost is ~0.024-0.034 at current compute budgets.

Verdict: PASS gate 1 (1.10–1.18×).

[ClassicLarry]

🤖 Gate 2 Analysis: MHA — no GQA (all 4 scales, macro_loss)

All 4 runs finished.

Results

Dim	Run	macro_loss	tps	tps ratio	Eff Speedup
512	baseline (kv=1)	3.8104	405,630	—	—
512	MHA (kv=4)	3.7868 (−0.024)	396,878	0.978	1.097
768	baseline (kv=1)	3.4339	273,532	—	—
768	MHA (kv=6)	3.3999 (−0.034)	263,894	0.965	1.177
1024	baseline (kv=2)	3.1605	175,162	—	—
1024	MHA (kv=8)	3.1412 (−0.019)	171,003	0.976	1.114
1280	baseline (kv=2)	3.0065	128,278	—	—
1280	MHA (kv=10)	2.9811 (−0.025)	123,312	0.961	1.167

Scaling law projections (pinned L∞=1.6, α=0.0941)

Fitted A: baseline=94.98, variant=93.56

Budget	BL projected	Var projected	Δ
1e21	2.604	2.589	−0.015
1e23	2.251	2.241	−0.010

Verdict

PASS gate 2. Effective speedup 1.10–1.18× at all 4 scales with only 2–4% throughput cost. Removing GQA gives consistent quality improvement that scales well. Projections −0.015 at 1e21, −0.010 at 1e23.

Note: MHA runs used the same step count as the GQA baseline. Since MHA has more FLOPs per token (more KV projections), these runs used slightly more total compute than baseline. The effective speedup formula accounts for this through throughput.

[ClassicLarry]

🤖 Closing — MHA passed gate 2 (1.10–1.18×).

[ClassicLarry]

We are intentionally giving up this ~15% speedup in order to have a 4x smaller KV cache. Later we can explore techniques like multihead latent attention to reclaim some portion of this 15%.