Add DiffusionGemma (Gemma-4 MoE block-diffusion) support: MoE imatrix variants, AWQ MoE groups, llama.cpp bump, diffusion eval backend

[pearsonkyle]

Summary

Adds end-to-end support for DiffusionGemma 26B-A4B (google/diffusiongemma-26B-A4B-it) — a 128-expert (8 active) Gemma-4 MoE block-diffusion LLM — covering the <3bpw quantization methods, bench, and full task-level eval. This is the repo's first MoE and first diffusion architecture.

Answers to the motivating questions

• Does imatrix work for this model? Yes. llama-imatrix runs a standard forward, and DiffusionGemma's causal encoder mode is exactly that, over the same weights the diffusion decoder uses (encoder/decoder are tied; only the attention mask differs). Two caveats are handled: (1) the UNIFIED graph treats the last canvas_length (256) positions of each chunk as a bidirectional canvas — recipes raise imatrix_ctx to 2048 to keep that fraction small, and quant-vs-quant KLD remains valid; (2) MoE imatrix entries carry per-expert counts and 3D weights, which the variant builders previously dropped silently (which would make llama-quantize refuse IQ2_*). They now normalize and rerank per expert, with an expert-coverage check wired into the pipeline.
• Can it be tuned with AWQ? Yes. New L{i}_moe scale groups fold into pre_feedforward_layernorm_2, whose only consumer is the fused 3D experts.gate_up_proj. The router reads the raw residual through its own internal norm, so routing logits are exactly invariant under the fold (unit-tested in f32). moe_expert_sample (default 8) caps the α-search cost; Gemma's (1+γ) RMSNorm convention rides the existing rmsnorm_plus_one path.
• GPTQ on the 128 experts is deferred (badly-conditioned per-expert Hessians × CPU Cholesky on 26B ≈ days); documented in CLAUDE.md.

Submodule bump

vendor/llama.cpp: 9e58d4d6 → c84e85af, the head of llama.cpp PR #24423 — current master has no DiffusionGemma support; the PR adds the diffusion-gemma arch, converter, llama-diffusion-cli canvas mode, and a persistent logits server. Setup instructions in CLAUDE.md show how to fetch the PR ref; re-pin to master once it merges. calibrate/_iq2_grids.py regenerated against the new pin (grids byte-identical; ksigns parity assertion passed). llama-perplexity still has no --parse-special (re-verified); the imatrix GGUF schema and converter CLI are unchanged.

Task eval without llama-server

llama-server cannot diffusion-decode, so eval/diffusion_server.py provides running_diffusion_server(...): an OpenAI-compatible shim that renders chat prompts (with tools) via the HF tokenizer, generates through a numpy port of llama.cpp's entropy-bound block-diffusion sampler driven by the persistent llama-diffusion-gemma-server (one model load per eval), and parses Gemma-4 <|tool_call>call:NAME{…}<tool_call|> markup into OpenAI tool_calls. The tool-call/MMLU-Pro evals and eval.reps are untouched apart from a server_factory hook; reps scripts gain --backend diffusion --hf-dir <snapshot>. scripts/bench_diffusion_speed.py times llama-diffusion-cli (llama-bench decode tok/s is meaningless for diffusion generation).

New recipes

• q2_k_diffusiongemma_imatrix — low-resource path (no HF load; llama-imatrix streams the GGUF).
• iq2_m_diffusiongemma_awq — AWQ + hybrid_custom imatrix stacked, IQ2_S codebook proxy auto-selected.

Testing

• 353 unit tests pass (33 new: MoE tensor mapping, per-expert imatrix rerank/normalization with real gguf-py round-trips, expert coverage, MoE group discovery on standard + DiffusionGemma trunks, f32 fold/routing invariance, deterministic expert subsampling, Gemma-4 markup parsing, canvas trimming, entropy-bound acceptance rule).
• ruff and mypy at or better than parity with main (mypy 36 → 28 pre-existing errors; the one ruff finding predates this PR).
• Not yet run (needs GPU/disk beyond this environment): real 26B end-to-end. The recipe headers carry the resource checklist (≥220 GB disk for the AWQ path, 80GB-class GPU or ≥96 GB RAM; imatrix-only path needs neither) and post-run checks (expert coverage ≥95%, llama-quantize applies imatrix to all ffn_*_exps, manual llama-diffusion-cli sanity generation at F16 vs IQ2_M).

https://claude.ai/code/session_011BaPPBxkVCggmj1McnSC32

---

Generated by Claude Code