Agent MoE Experiment: Wide attention (1.5x heads, head_dim=128)

[ClassicLarry]

TL;DR

Increase number of attention heads by 1.5x while keeping head_dim at 128. Tests whether more attention capacity improves quality.

User prompt

follow agent.md and implement increased attn width. that is, instead of having number of attention heads be hidden_dim/128, make it 1.5*hidden_dim/128. Keep the head_dim at 128. So d768 will have 9 heads instead of 6. make sure GQA still works.

Scope

• Parent: MoE Scaling up to April goal #4281
• Branch: moe_wide_attn
• Sweep file: experiments/grug/moe/wide_attn_sweep.py

Head counts

d	Current heads/kv	1.5x heads/kv
512	4/1	6/1
768	6/1	9/1
1024	8/2	12/3
1280	10/2	15/3

GQA 4:1 ratio preserved at all sizes.

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: Wide Attention 1.5x (macro_loss)

Both runs finished. Marginal pass — quality improves but ~6% throughput cost nearly offsets it.

Dim	Run	macro_loss	tps	tps ratio	Eff Speedup
512	baseline (4 heads)	3.8104	405,630	—	—
512	wide-attn (6 heads)	3.7884 (−0.022)	384,948	0.949	1.055
768	baseline (6 heads)	3.4339	273,532	—	—
768	wide-attn (9 heads)	3.4160 (−0.018)	256,924	0.939	1.043

Note: The FLOPs-per-token increase from wider attention (~33% more attention FLOPs) is reflected in the throughput drop. The effective speedup accounts for this — the quality gain just barely outweighs the extra compute.

Verdict: Marginal PASS gate 1 (1.04–1.06x). Quality improvement is real but the throughput cost is significant. Compare to partial rope (1.09–1.12x) or PKO+prope (1.09–1.23x) which achieve similar or better quality gains with no throughput cost.

[ClassicLarry]

🤖 Gate 2 Analysis: Wide Attention 1.5x (all 4 scales, macro_loss)

All 4 runs finished.

Results

Dim	Run	macro_loss	tps	tps ratio	Eff Speedup
512	baseline (4 heads)	3.8104	405,630	—	—
512	wide-attn (6 heads)	3.7884 (−0.022)	384,948	0.949	1.055
768	baseline (6 heads)	3.4339	273,532	—	—
768	wide-attn (9 heads)	3.4160 (−0.018)	256,924	0.939	1.043
1024	baseline (8 heads)	3.1605	175,162	—	—
1024	wide-attn (12 heads)	3.1405 (−0.020)	163,586	0.934	1.071
1280	baseline (10 heads)	3.0065	128,278	—	—
1280	wide-attn (15 heads)	2.9888 (−0.018)	117,532	0.916	1.048

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

Fitted A: baseline=94.98, variant=93.91

Budget	BL projected	Var projected	Δ
1e21	2.604	2.592	−0.011
1e23	2.251	2.243	−0.007

Note on FLOPs

Wide attention increases attention FLOPs ~33% (more Q/K/V projections + attention compute at 1.5x heads with head_dim=128). The compute_flops_per_token function underestimates this because it infers head_dim = hidden_dim / num_heads instead of using the actual 128. The effective speedup accounts for this via throughput.

Verdict

Marginal PASS gate 2. Effective speedup 1.04–1.07× across all scales. Consistent quality improvement (−0.018 to −0.022) but 5–8% throughput cost. Projections show −0.011 at 1e21 and −0.007 at 1e23.

[ClassicLarry]

🤖 Closing — wide attention marginal pass gate 2 (1.04–1.07×).