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

CLAUDE.md

@@ -17,8 +17,16 @@ the canonical entry point for reproducing the published table.
 ## Setup and common commands
 
 ```bash
-# One-time: fetch + build vendored llama.cpp (pinned commit 9e58d4d6)
-git submodule update --init --recursive
+# One-time: fetch + build vendored llama.cpp.
+# Pinned commit c84e85af = head of llama.cpp PR #24423 (DiffusionGemma:
+# `diffusion-gemma` arch + converter + llama-diffusion-cli canvas mode +
+# llama-diffusion-gemma-server). NOT on master yet — if `submodule update`
+# can't resolve the gitlink, fetch the PR ref inside the submodule and retry.
+# Re-pin to master once the PR merges.
+git submodule update --init --recursive || (
+  git -C vendor/llama.cpp fetch origin pull/24423/head &&
+  git submodule update --recursive
+)
 cmake -S vendor/llama.cpp -B vendor/llama.cpp/build -DGGML_METAL=ON   # Linux+CUDA: -DGGML_CUDA=ON
 cmake --build vendor/llama.cpp/build -j
 
@@ -123,6 +131,62 @@ Pipeline behaviors worth knowing:
 ### Tensor naming
 `models/hf_gguf_map.py` maps HF parameter names to GGUF tensor names. Anything that
 crosses the HF↔GGUF boundary (imatrix variants, AWQ apply) goes through this mapping.
+MoE expert stacks map to fused 3D tensors (`ffn_gate_up_exps`/`ffn_down_exps`,
+`[n_expert, n_out, n_in]`); `is_moe_expert` gates per-expert handling the way
+`is_ssm` gates SSM passthrough. DiffusionGemma module paths (`model.decoder.*`,
+`model.encoder.language_model.*`) are prefix-normalized to `model.*` before
+matching — same rewrite llama.cpp's converter applies.
+
+### DiffusionGemma / Gemma-4 MoE (`diffusion-gemma` arch)
+`google/diffusiongemma-26B-A4B-it`: 128-expert (8 active) Gemma-4 MoE whose
+causal **encoder** (prompt prefill) and bidirectional diffusion **decoder**
+(256-token canvas denoising) share all backbone weights — only the attention
+mask differs. Requires the PR-#24423 llama.cpp pin (see Setup). Key facts the
+code depends on:
+- **Checkpoint shape**: arch `DiffusionGemmaForBlockDiffusion`, backbone under
+  `model.decoder.layers.N.*` (encoder tied via `_tied_weights_keys`), vision
+  tower on the encoder (the converter drops it), decoder-only
+  `self_conditioning` gated MLP, root-level `canvas_length`.
+  `extract_text_lm` passes the checkpoint through untouched (`ForBlockDiffusion`
+  predicate) — never rebuild a text-only CausalLM from it.
+- **Each layer has BOTH a dense MLP and an MoE block**, each behind its own
+  pre-norm (`pre_feedforward_layernorm` / `pre_feedforward_layernorm_2`).
+  Experts are ONE fused 3D param per projection (`experts.gate_up_proj`,
+  `experts.down_proj`); the router (`router.proj`) reads the **raw residual
+  through its own internal norm** — not the experts' pre-norm output.
+- **imatrix works** (encoder mode is a standard causal forward over the shared
+  weights) with two caveats: (1) llama-imatrix/llama-perplexity run the
+  UNIFIED graph, which treats the last `canvas_length` positions of every
+  chunk as a bidirectional canvas — quant-vs-quant KLD stays valid, absolute
+  PPL is not deployment-faithful; raise `imatrix_ctx` (recipes use 2048) to
+  keep the canvas fraction small. (2) MoE imatrix entries carry **per-expert
+  counts**; `_load_base_imatrix` normalizes per expert (zero-count experts →
+  neutral ones, mirroring llama-quantize), the analytic variants rerank **per
+  expert** (L1 normalization must not bleed across experts), and
+  `imatrix.warn_moe_coverage` flags expert tensors below 95% coverage — the
+  fix is more calibration tokens.
+- **AWQ works**: `discover_groups` emits `L{i}_moe` groups (the fused gate_up
+  stack under `pre_feedforward_layernorm_2`, anchored on the `experts` module
+  which sees every token pre-routing). Folding there leaves routing logits
+  exactly invariant because the router never consumes that norm's output —
+  do NOT add the router to the group. `moe_expert_sample` (default 8) caps the
+  α-search cost; `rmsnorm_plus_one: true` (Gemma `(1+γ)` convention). The HF
+  load goes through `_hf.load_calibration_model` (falls back to the concrete
+  arch class when `AutoModelForCausalLM` refuses) and sanity logits through
+  `_hf.causal_logits` (encoder + tied lm_head).
+- **GPTQ on the experts is unsupported** (deferred): 128 per-expert Hessians,
+  each seeing ~1/16 of tokens, is days of badly-conditioned CPU Cholesky.
+- **Task eval cannot use llama-server** (no diffusion decoding). Use
+  `eval/diffusion_server.py`: `running_diffusion_server(model_path, hf_dir=…)`
+  yields a `base_url` whose `/v1/chat/completions` runs the entropy-bound
+  block-diffusion sampler (numpy port of llama.cpp's reference, driven through
+  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`. Reps scripts take `--backend diffusion --hf-dir <snapshot>`.
+- **Speed**: llama-bench decode tok/s is meaningless for diffusion generation;
+  use `scripts/bench_diffusion_speed.py` (wraps llama-diffusion-cli).
+- Recipes: `q2_k_diffusiongemma_imatrix` (low-resource, no HF load),
+  `iq2_m_diffusiongemma_awq` (AWQ + hybrid_custom stacked).
 
 ### Bench
 `bench/runner.py` defines `BenchRow` and `CSV_COLUMNS`. Sub-modules:
@@ -233,7 +297,9 @@ creates `model_extracted/`, `corpus/`, `calibration/`, `gguf/`, `eval/` and rese
   `hybrid_custom` imatrix stacked, codebook proxies auto-selected),
   `q2_k_gptq`, `iq3_s_gptq` (asym grid + relaxed guardrails auto-derived).
 - Model-specific: `q4_k_m_qwen3_5_4b`, `{q4_k_m,q5_k_s}_qwen3_6_mtp{,_awq,_none}`,
-  `iq3_s_9b_mtp` (MTP heads kept via `extract.keep_mtp`).
+  `iq3_s_9b_mtp` (MTP heads kept via `extract.keep_mtp`),
+  `q2_k_diffusiongemma_imatrix` + `iq2_m_diffusiongemma_awq` (DiffusionGemma
+  26B-A4B MoE — see the DiffusionGemma section for the caveats).
 A unit test (`test_all_packaged_recipes_parse`) requires every shipped recipe to
 validate, and IQ1/IQ2 recipes to carry a calibration method — keep it green when
 adding recipes.

scripts/bench_diffusion_speed.py

@@ -0,0 +1,103 @@
+#!/usr/bin/env python3
+"""Denoising-throughput bench for diffusion GGUFs (DiffusionGemma et al.).
+
+llama-bench measures autoregressive prefill/decode — decode tok/s is NOT how
+a block-diffusion model generates. This script times `llama-diffusion-cli`
+end-to-end instead: wall time per generated token across N repetitions
+(mean ± stdev), the diffusion analog of bench/speed.py.
+
+Usage:
+  uv run python scripts/bench_diffusion_speed.py \
+      --model out/run/gguf/IQ2_M-awq-hybrid_custom.gguf \
+      --prompt "Explain the Fourier transform." \
+      --n-predict 256 --reps 5 --out results.csv
+"""
+
+from __future__ import annotations
+
+import argparse
+import csv
+import re
+import statistics
+import subprocess
+import sys
+import time
+from pathlib import Path
+
+sys.path.insert(0, str(Path(__file__).parent.parent / "src"))
+
+from quant_tuner.paths import llama_bin  # noqa: E402
+
+# entropy-bound CLI timing line: "total time: 1234.56ms, time per step: ..."
+_TOTAL_RE = re.compile(r"total time:\s*([0-9.]+)ms.*\((\d+) steps over (\d+) blocks", re.S)
+
+
+def run_once(args, rep: int) -> dict:
+    cmd = [
+        str(llama_bin("llama-diffusion-cli")),
+        "-m", str(args.model),
+        "-p", args.prompt,
+        "-n", str(args.n_predict),
+        "-ngl", str(args.ngl),
+        "--seed", str(args.seed + rep),
+    ]
+    if args.diffusion_steps:
+        cmd += ["--diffusion-steps", str(args.diffusion_steps)]
+    t0 = time.time()
+    proc = subprocess.run(cmd, capture_output=True, text=True, timeout=args.timeout)
+    wall_s = time.time() - t0
+    if proc.returncode != 0:
+        raise RuntimeError(f"llama-diffusion-cli failed (rc={proc.returncode}):\n{proc.stderr[-2000:]}")
+    combined = proc.stdout + proc.stderr
+    row = {"rep": rep, "wall_s": round(wall_s, 3)}
+    m = _TOTAL_RE.search(combined)
+    if m:
+        row["gen_ms"] = float(m.group(1))
+        row["steps"] = int(m.group(2))
+        row["blocks"] = int(m.group(3))
+        row["ms_per_step"] = round(row["gen_ms"] / max(row["steps"], 1), 2)
+    return row
+
+
+def main() -> int:
+    p = argparse.ArgumentParser(description=__doc__,
+                                formatter_class=argparse.RawDescriptionHelpFormatter)
+    p.add_argument("--model", type=Path, required=True)
+    p.add_argument("--prompt", default="Write a Python function that merges two sorted lists.")
+    p.add_argument("--n-predict", type=int, default=256)
+    p.add_argument("--reps", type=int, default=5)
+    p.add_argument("--seed", type=int, default=42)
+    p.add_argument("--ngl", type=int, default=99)
+    p.add_argument("--diffusion-steps", type=int, default=0,
+                   help="override step count (0 = model/GGUF default)")
+    p.add_argument("--timeout", type=int, default=1800)
+    p.add_argument("--out", type=Path, default=None, help="CSV output path")
+    args = p.parse_args()
+
+    rows = []
+    for rep in range(args.reps):
+        row = run_once(args, rep)
+        rows.append(row)
+        print(f"rep {rep}: {row}", flush=True)
+
+    walls = [r["wall_s"] for r in rows]
+    mean = statistics.mean(walls)
+    stdev = statistics.stdev(walls) if len(walls) > 1 else 0.0
+    tok_s = args.n_predict / mean if mean > 0 else 0.0
+    print(f"\n{args.model.name}: wall {mean:.2f}s ± {stdev:.2f} "
+          f"(~{tok_s:.1f} tok/s incl. model load, n={args.reps})")
+
+    if args.out:
+        args.out.parent.mkdir(parents=True, exist_ok=True)
+        cols = sorted({k for r in rows for k in r})
+        with args.out.open("w", newline="") as f:
+            w = csv.DictWriter(f, fieldnames=["model", *cols])
+            w.writeheader()
+            for r in rows:
+                w.writerow({"model": args.model.name, **r})
+        print(f"wrote {args.out}")
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

scripts/run_mmlu_pro_reps.py

@@ -58,6 +58,12 @@ def _build_arg_parser() -> argparse.ArgumentParser:
     p.add_argument("--ctx", type=int, default=8192)
     p.add_argument("--ngl", type=int, default=99)
     p.add_argument("--max-tokens", type=int, default=2048)
+    p.add_argument("--backend", choices=("server", "diffusion"), default="server",
+                   help="generation backend: llama-server (default) or the "
+                        "DiffusionGemma block-diffusion shim (requires --hf-dir)")
+    p.add_argument("--hf-dir", type=Path, default=None,
+                   help="HF snapshot dir with the tokenizer/chat template "
+                        "(required for --backend diffusion)")
     p.add_argument("--temperature", type=float, default=0.6,
                    help="Default 0.6 — reps only diverge at T>0. Set 0.0 for greedy.")
     p.add_argument("--top-p", type=float, default=0.95)
@@ -124,6 +130,16 @@ def on_model(mr) -> None:
     print(f"holdout:   {args.holdout}", flush=True)
     print(f"results:   {args.results} / {args.aggregated}", flush=True)
 
+    server_factory = None
+    if args.backend == "diffusion":
+        if args.hf_dir is None:
+            raise SystemExit("--backend diffusion requires --hf-dir (tokenizer source)")
+        from functools import partial
+
+        from quant_tuner.eval.diffusion_server import running_diffusion_server
+
+        server_factory = partial(running_diffusion_server, hf_dir=args.hf_dir)
+
     t0 = time.time()
     by_model = run_reps_for_models(
         models=list(args.models),
@@ -137,6 +153,7 @@ def on_model(mr) -> None:
         chat_template_kwargs=(args.chat_template_kwargs or None),
         per_rep_callback=on_rep,
         per_model_callback=on_model,
+        server_factory=server_factory,
     )
 
     write_csvs(

scripts/run_toolcall_reps.py

@@ -91,6 +91,12 @@ def _build_arg_parser() -> argparse.ArgumentParser:
                    help="Override sampling.top_p")
     p.add_argument("--top-k", type=int, default=None,
                    help="Override sampling.top_k")
+    p.add_argument("--backend", choices=("server", "diffusion"), default="server",
+                   help="generation backend: llama-server (default) or the "
+                        "DiffusionGemma block-diffusion shim (requires --hf-dir)")
+    p.add_argument("--hf-dir", type=Path, default=None,
+                   help="HF snapshot dir with the tokenizer/chat template "
+                        "(required for --backend diffusion)")
     p.add_argument("--min-p", type=float, default=None,
                    help="Override sampling.min_p")
     return p
@@ -158,6 +164,16 @@ def on_model(mr) -> None:
     print(f"results:     {args.results}  (per-rep) / {args.aggregated}  (aggregated)",
           flush=True)
 
+    server_factory = None
+    if args.backend == "diffusion":
+        if args.hf_dir is None:
+            raise SystemExit("--backend diffusion requires --hf-dir (tokenizer source)")
+        from functools import partial
+
+        from quant_tuner.eval.diffusion_server import running_diffusion_server
+
+        server_factory = partial(running_diffusion_server, hf_dir=args.hf_dir)
+
     t0 = time.time()
     by_model = run_reps_for_models(
         models=list(args.models),
@@ -170,6 +186,7 @@ def on_model(mr) -> None:
         chat_template_kwargs=(args.chat_template_kwargs or None),
         per_rep_callback=on_rep,
         per_model_callback=on_model,
+        server_factory=server_factory,
     )
 
     write_csvs(

src/quant_tuner/calibrate/_hf.py

@@ -2,6 +2,75 @@
 
 from __future__ import annotations
 
+import json
+from pathlib import Path
+from typing import Any
+
+# Candidate dotted paths (relative to the top-level HF model) for the decoder
+# layer stack. DiffusionGemma's causal *encoder* stack is listed rather than
+# its diffusion decoder: calibration forwards run in encoder (causal) mode,
+# and the two stacks share every backbone weight (tied parameters), so weight
+# mutations applied through encoder paths land in the decoder too.
+_TRUNK_CANDIDATES = (
+    "model",
+    "model.encoder.language_model",
+    "encoder.language_model",
+)
+
+
+def _get_attr_path(obj: Any, dotted: str) -> Any | None:
+    for part in dotted.split("."):
+        obj = getattr(obj, part, None)
+        if obj is None:
+            return None
+    return obj
+
+
+def resolve_text_trunk(model) -> tuple[str, Any]:
+    """Locate the text decoder trunk (the module owning ``.layers``).
+
+    Returns ``(dotted_path, module)`` where ``dotted_path`` is relative to the
+    top-level model (e.g. ``"model"`` for standard CausalLMs,
+    ``"model.encoder.language_model"`` for DiffusionGemma). Raises if no
+    candidate matches.
+    """
+    for path in _TRUNK_CANDIDATES:
+        trunk = _get_attr_path(model, path)
+        if trunk is not None and hasattr(trunk, "layers"):
+            return path, trunk
+    raise RuntimeError(
+        "no text trunk with .layers found (tried: "
+        + ", ".join(f"model.{p}" for p in _TRUNK_CANDIDATES)
+        + "); unsupported architecture"
+    )
+
+
+def load_calibration_model(model_dir: Path, *, torch_dtype, device: str):
+    """Load an HF checkpoint for calibration, handling non-CausalLM classes.
+
+    ``AutoModelForCausalLM`` covers the standard architectures; checkpoints
+    whose class isn't registered there (e.g. ``DiffusionGemmaForBlockDiffusion``)
+    are loaded via their concrete ``architectures[0]`` class from
+    ``transformers`` instead.
+    """
+    import transformers
+    from transformers import AutoModelForCausalLM
+
+    try:
+        model = AutoModelForCausalLM.from_pretrained(
+            model_dir, torch_dtype=torch_dtype, trust_remote_code=True
+        )
+    except (ValueError, KeyError):
+        config = json.loads((Path(model_dir) / "config.json").read_text())
+        archs = config.get("architectures") or []
+        cls = getattr(transformers, archs[0], None) if archs else None
+        if cls is None:
+            raise
+        model = cls.from_pretrained(
+            model_dir, torch_dtype=torch_dtype, trust_remote_code=True
+        )
+    return model.to(device)
+
 
 def forward_no_logits(model, input_ids) -> None:
     """Run a forward pass for hook-based stat collection, skipping the LM head.
@@ -10,11 +79,27 @@ def forward_no_logits(model, input_ids) -> None:
     logits are discarded, yet ``model(ids)`` materializes a
     ``[1, ctx, vocab]`` float tensor (several GB at ctx 4096 on a 150k-vocab
     model) and pays the lm_head matmul for every chunk. Calling the decoder
-    trunk (``model.model``) directly avoids both. Falls back to the full
-    model for architectures without the standard ``.model.layers`` shape.
+    trunk (``model.model``, or DiffusionGemma's causal encoder stack) directly
+    avoids both. Falls back to the full model for architectures without a
+    recognizable trunk.
     """
-    base = getattr(model, "model", None)
-    if base is not None and hasattr(base, "layers"):
-        base(input_ids=input_ids)
-    else:
+    try:
+        _, trunk = resolve_text_trunk(model)
+    except RuntimeError:
         model(input_ids)
+        return
+    trunk(input_ids=input_ids)
+
+
+def causal_logits(model, input_ids):
+    """Causal-LM logits for sanity checks, on architectures where ``model(ids)``
+    isn't a causal forward (DiffusionGemma's top-level forward wants canvas
+    inputs; its *encoder* + tied ``lm_head`` is the causal path)."""
+    trunk_path, trunk = resolve_text_trunk(model)
+    if trunk_path == "model":
+        return model(input_ids).logits
+    hidden = trunk(input_ids=input_ids).last_hidden_state
+    head = getattr(model, "lm_head", None)
+    if head is None:
+        raise RuntimeError(f"no lm_head on model with trunk {trunk_path!r}")
+    return head(hidden)

src/quant_tuner/calibrate/_iq2_grids.py

@@ -1,6 +1,6 @@
 """E8-lattice codebooks for the IQ2_* quants — GENERATED, do not edit.
 
-Extracted from llama.cpp ggml-common.h at commit 9e58d4d692ed3d350591cc86d06c73c61c122509
+Extracted from llama.cpp ggml-common.h at commit c84e85af61011f9fbfcf41479381d5ed1661a564
 by scripts/gen_iq2_grids.py. Each grid entry is 8 magnitude bytes
 (component order); signs live outside the codebook. For IQ2_XXS and
 IQ2_XS only sign patterns with an EVEN number of negatives per group