Nightly 2026-08-03: first MoE architecture (OLMoE-1B-7B), parity-verified

[physics515]

Nightly run 2026-08-03. Four increments; the headline is that the zoo is no longer dense-only — OLMoE-1B-7B is ported from scratch and passed its llama.cpp parity gate on the first run.

1. Dependency freshness (b327514)

cargo update took 8 packages to their latest compatible versions (clap 4.6.4→4.6.5, clap_builder 4.6.2→4.6.5, cubecl-hip-sys 7.2→7.14, data-encoding, hybrid-array, ipnet, libredox, time 0.3.54→0.3.55). cargo upgrade --incompatible offered exactly one major — wgpu 29 → 30 — which stays held, per the standing pin rationale: wgpu is not ours to pick, burn 0.21 resolves 29 transitively, so it unblocks with a burn bump rather than a cargo upgrade. tokenizers 0.23.1 remains current. Green + Cargo.lock committed as a standalone increment.

2. First MoE architecture — OLMoE-1B-7B (a5208a2)

Closes the P2 first MoE architecture item from the 2026-07-30 colibri parity scan.

nn::SparseMoe — a softmax top-k router over a fused 3-D expert bank in [experts, out, in] layout, which is the row-major twin of ggml's ffn_{gate,up,down}_exps, so 64 experts load as three tensors per layer instead of 192. Compute is dense-mask: every expert processes every token and the router's weight row (exactly zero off the top-k) scales the rest away — numerically identical to the sparse formulation and it keeps the whole forward on-device with no data-dependent gather. The routing scatter is an on-device arange-compare rather than burn's one_hot, which reads indices back to the host and would sync every layer.

GqaAttention learned OLMoE's q/k-norm placement — RMSNorm over the whole projection before the head split, versus the per-head placement Qwen3/LFM2 use. Which one applies is inferred from the loaded gamma's own width (head_dim → per-head, n·head_dim → projection), so the module shape is unchanged and every existing checkpoint loads identically. The cache-equivalence proof now runs all three norm modes.

models::olmoe — config-driven off olmoe.* GGUF metadata (expert counts, either feed_forward_length spelling, expert_weights_norm), untied head, CausalLm impl. GGUF-only by design: HF ships the experts unfused as mlp.experts.{i}.*, so a safetensors path needs a 64-way concat-on-import — split out as its own roadmap item. Registry gains the allenai Q4_K_M catalog entry; the tokenizer registry gains the olmo pre (stock GPT-2 regex + NFC, per the checkpoint's own tokenizer.json).

Parity gate — PASSED, first run

parity_gguf.rs gains an olmoe leg; the harness is now generic over the backend so this tier can compare on Cpu (the ~28 GB f32 build does not fit a 16 GB card). llama.cpp runs the same Q4_K_M file:

• top-5 first-forward ids match exactly in order: [1992, 17833, 11202, 4943, 1394]
• 24-token greedy sequence byte-identical: "To solve this problem, we need to identify the first five prime numbers. A prime number is a natural number greater than"
• max |Δlogprob| 3.687691131310693e-1, inside the shared 7.5e-1 tolerance and right beside Qwen3's 4.02e-1

Real-model verification

real_olmoe.rs: the registry spec fetched the 4.21 GB GGUF, the one file loaded 16 layers × 64 experts in 92 s and greedy-decoded "2 + 2 equals 4." on the CPU backend (sanity spread 36.9). Its GGUF-built tokenizer is byte-identical to the checkpoint's tokenizer.json across an 8-prompt battery (specials, CJK/emoji, whitespace runs, contractions, empty).

Every prior gate re-verified after the attention change

Gate	Result
Qwen3 GGUF vs llama.cpp	bit-identical — top-5 exact in order, max |Δlogprob| 4.015608155114805e-1, greedy byte-identical
Qwen2.5 both legs	Candle max |Δlogit| 1.9e-5; Ollama greedy byte-identical
GPU budget	104.4 ms TTFT / 13.2 tok/s (budgets 150 ms / 10 tok/s)
CPU budget	15.3 tok/s (budget 6)
Template byte gate	10/10 byte-identical
Qwen3 real-GPU tool call	clean parseable <tool_call> emitted

202 unit tests.

3. Docs + a recorded negative result (a01d829)

The P2 item ticks [x] with its gate numbers. Two follow-ups split out: HF safetensors import, and making routed-expert compute pay.

That second item records a measured negative result instead of a shipped regression. The obvious optimization — gather only the 8 routed expert slices per decode step via a device-side select off the router's own index tensor — is exact-same-math (a unit test confirms the two paths agree to 1e-6) but measured 1.58 s/token against the dense path's 1.15 on burn-flex: the ~200 MB per-layer-per-token gather copy costs more than the dense matmul it removes, i.e. select materializes where the dense path streams. It was reverted rather than shipped, and the roadmap now names the three routes worth trying (a fused kernel that never materializes the gathered bank; measuring on GPU, blocked until a MoE fits VRAM; llama.cpp's --n-cpu-moe-style placement that moves whole banks instead of slicing them) and what each is blocked on.

Research folded in: burn 0.22 is still pre-release (0.22.0-pre.1 remains newest — the P0 migration item stays gated), CubeCL tagged 0.11.0-pre.1 the same day (Metal backend, new CPU runtime, frontend refactor — all arriving with the burn bump), and llama.cpp's --n-cpu-moe plus its two-tier GPU+RAM expert-cache RFC (#20757) as prior art for the P6 streaming tier.

4. MoE budget row + regression gate (aa88a00)

budget_moe.rs defends the new tier like every other path: load 81.9 s (ceiling 300) and decode 0.76 s/token warm (ceiling 2.0, ~2.6× headroom for host-load noise). Budgeted in seconds/token rather than tok/s because the dense-mask forward touches all 7B params per token — s/token is the honest unit and the number the routed-compute work has to beat. bench/BASELINE.md records both the 0.76 warm figure and the 1.15 end-to-end one with what each measures, so the gap reads as methodology rather than drift.

What's next

• Routed-expert compute — the 0.76 s/token baseline is now the target; the three candidate routes are scoped in P2.
• OLMoE from HF safetensors — mechanical concat-on-import, wants its own ~14 GB fixture.
• burn 0.22 migration — still gated on a stable release; the pre-release telegraphs breaking changes to exactly the B::FloatElem seam Mummu's dtype pinning rides.
• P9 keep-quantized VRAM — the lever that would put a MoE on the 16 GB card at all.

🤖 Generated with Claude Code

[Copilot]

Pull request overview

This nightly update expands Mummu’s model zoo beyond dense-only architectures by adding an OLMoE Mixture-of-Experts (MoE) implementation (GGUF-first), updating attention q/k-norm handling to support OLMoE’s projection-wide norm placement, and introducing parity + real-weights + budget gates for the new tier.

Changes:

• Add OLMoE MoE architecture support (nn::SparseMoe, models::olmoe) with GGUF import, plus new parity/real-weights validation.
• Extend GqaAttention to support (and infer) per-head vs projection-wide q/k RMSNorm placement.
• Add MoE CPU perf budget gate + baseline documentation; update dependency lockfile and roadmap/readme narrative.

Reviewed changes

Copilot reviewed 17 out of 18 changed files in this pull request and generated 5 comments.

Show a summary per file

File	Description
ROADMAP.md	Marks first MoE architecture as shipped; records MoE follow-ups and related research.
README.md	Updates public-facing architecture list and adds an MoE overview section.
crates/mummu/tests/real_olmoe.rs	Adds ignored real-weights validation for OLMoE GGUF load + decode + tokenizer parity.
crates/mummu/tests/parity_gguf.rs	Generalizes GGUF parity harness over backend and adds an OLMoE parity leg.
crates/mummu/src/tokenizer.rs	Adds olmo tokenizer pre-spec for OLMo/OLMoE family.
crates/mummu/src/registry.rs	Adds Architecture::Olmoe and registers the OLMoE Q4_K_M GGUF catalog entry.
crates/mummu/src/nn/moe.rs	Introduces nn::SparseMoe and fused expert-bank implementation + unit tests.
crates/mummu/src/nn/mod.rs	Exposes the new MoE module/types from nn.
crates/mummu/src/nn/attention.rs	Adds projection-wide q/k norm placement support + updated tests.
crates/mummu/src/models/qwen3.rs	Plumbs the new attention config field for Qwen3 (per-head norm).
crates/mummu/src/models/qwen2.rs	Plumbs the new attention config field for Qwen2 (no q/k norm).
crates/mummu/src/models/olmoe.rs	Adds the OLMoE model implementation and GGUF import path.
crates/mummu/src/models/mod.rs	Exposes the new models::olmoe module.
crates/mummu/src/models/lfm2.rs	Plumbs the new attention config field for LFM2 (per-head norm).
crates/mummu/src/gguf.rs	Adds GgufValue::as_bool helper used by OLMoE metadata parsing.
crates/mummu-bench/tests/budget_moe.rs	Adds ignored MoE CPU decode/load budget gate for OLMoE.
Cargo.lock	Updates several dependencies to latest compatible versions.
bench/BASELINE.md	Records MoE baseline numbers and documents the routed-gather negative result.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.