EXPERT_BUDGET=N: miss-aware cap on distinct experts/layer (+75% decode tok/s, 6x prefill on low-RAM hosts)

[woolcoxm]

Summary

Adds EXPERT_BUDGET=N env var that caps the number of distinct experts loaded per layer across the batch-union. When the union exceeds the budget, keeps cache hits + the highest-aggregate-gate-weight misses, drops the rest — they're never loaded from disk. Targets the low-RAM design corner (cap=1-2, nearly every routed expert is a disk miss).

Closes #220.

The problem

Every expert miss costs ~19 MB of disk I/O. On low-RAM hosts where the LRU cache cap is tiny (e.g. cap=2 on 24 GB RAM), nearly every routed expert is a miss. With topk=8 and 75 sparse layers, a single-token decode reads ~8.5 GB of experts from disk. Under MTP the batch-union produces 20-32 distinct experts per layer — almost all misses.

TOPP trims experts within a single position's top-K. It cannot reduce the cross-position union — the deduplication across batch positions (prefill, MTP verification) that multiplies disk loads.

The fix

In moe(), after the batch-union uniq[] is built but before the cache-resolve/load loop:

if(EXPERT_BUDGET > 0 && nu > EXPERT_BUDGET):
    1. Count cache hits in uniq[] (these are always kept — miss-aware)
    2. miss_budget = EXPERT_BUDGET - nhits
    3. Sort MISS experts by descending aggregate gate weight
    4. Keep top miss_budget misses + all hits
    5. Remove dropped experts from each position's idxs[]/keff[]
       (decrement keff, renormalize remaining weights if norm_topk)

Dropped experts are removed from idxs[] entirely — never resolved, never loaded, never computed. The downstream code already tolerates keff[s] < K (the TOPP path produces the same state).

Miss-aware design: cache hits are always kept regardless of budget. Only misses compete for the remaining budget slots. This means on a warm cache, budget has no effect — it only constrains cold-cache disk reads.

Code change

c/glm.c — 68 lines added, single file:

• g_expert_budget global + g_budget_dropped counter
• Env var parsing: EXPERT_BUDGET=N (default 0 = off)
• Budget cap logic in moe() batch-union
• Stats line: EXPERT_BUDGET=N (dropped X experts, ~Y GB I/O saved)

A/B measurements

Setup: GLM-5.2 744B per-row int4, 24 GB RAM (cache cap=2), Core Ultra 9 185H (AVX-VNNI), same prompt: "Explain the concept of recursion in programming. Provide a simple example.", 32 tokens, greedy (temp=0).

MTP off — speed

Config	tok/s	vs baseline	hit rate	prefill	decode	experts dropped	I/O saved
Baseline (budget=0)	0.16	—	10.0%	38.3s	205.7s	0	0
EXPERT_BUDGET=8	0.16	+0%	13.5%	9.5s	196.8s	3,798	71.8 GB
EXPERT_BUDGET=6	0.20	+25%	19.4%	10.4s	162.2s	8,738	165.1 GB
EXPERT_BUDGET=4	0.28	+75%	30.7%	6.2s	113.6s	13,399	253.2 GB

MTP on — speed + acceptance rate

Config	tok/s	hit rate	MTP acceptance	tokens/forward	experts dropped
Baseline (MTP=1)	0.12	9.6%	38% (17/45)	2.13	0
EXPERT_BUDGET=8 (MTP=1)	0.23	33.4%	21% (12/57)	1.68	23,267

MTP interaction — the honest finding

The maintainer's concern (#220 comment) was exactly right: dropped experts on draft positions force MTP rejects. At budget=8 + MTP, acceptance drops from 38% → 21%. The draft tokens route to experts that were budgeted away, so verification rejects them.

Despite the lower acceptance, tok/s still improves (0.12 → 0.23, +92%) because the per-forward cost drops dramatically (fewer experts loaded). The net win is positive but the acceptance hit is real.

Recommendation: use a higher budget with MTP (e.g. EXPERT_BUDGET=16-24) so the draft's experts survive. Without MTP, budget=4-6 is the sweet spot.

Token agreement — the confound

Greedy-token agreement between budgeted and unbudgeted runs is 0% at every budget level on this model. But this is not meaningful: the per-row int4 model is already incoherent at baseline (output: "Hire Some To Take Object-Oriented Programming Assignment"), so all outputs are garbled differently. A proper token-agreement test requires the grouped-int4 model (#242, merged) on a warm cache where the baseline produces coherent text. That A/B is the next step once the grouped conversion completes.

Safety

• Default OFF (EXPERT_BUDGET=0) — zero behavior change unless explicitly set
• Opt-in quality trade-off — same design philosophy as TOPP
• No output corruption — dropped experts' contributions are omitted and remaining weights renormalized (identical math to TOPP)
• Compatible with all features — PILOT, MTP, CUDA, CACHE_ROUTE, PIPE, TOPP all work unchanged
• No crash risk — no new allocation patterns, purely a filter on existing data structures

Prior art

MoE-Spec (arXiv 2602.16052) — "Expert Budgeting for Efficient Speculative Decoding": "top 32 of 64 experts capture 93% of routing weight." Our approach applies the same principle at the batch-union level, making it effective for prefill and single-token decode too.