boxwrench/gemma-4-qat-mtp-assistant-heads

Gemma 4 QAT MTP Assistant Heads — Atomic GGUF (Q8_0)

Three draft heads for speculative decoding with the official Gemma 4 QAT Q4_0 models. Converted from Google's published unquantized QAT assistant checkpoints into the GGUF format used by the Atomic TurboQuant llama.cpp fork. Using QAT-matched heads instead of the generic non-QAT heads raised acceptance rates from as low as 56.9% up to 91.8% on the 26B-A4B model.

Compatibility: These files use the gemma4_assistant architecture shape. They load on: - The Atomic TurboQuant llama.cpp fork - Stock llama.cpp once PR #23398 ("llama: add Gemma4 MTP") merges — it uses the same architecture name and tensor layout They will not load on stock llama.cpp before that PR, and are not compatible with the ik_llama format used by ji-farthing/gemma-4-qat-q4_0-MTP-assistants-ik-llama-GGUF. PARALLEL=2 note: The same n_tokens reshape bug that caused 2-slot crashes on the Atomic fork is also present in PR #23398 at the time of writing. Until that is patched upstream, use --n-parallel 1 regardless of which build you're on.

Background — What Is a Draft Head and Why Does Matching Matter?

Speculative decoding (called MTP — Multi-Token Prediction — in llama.cpp) is a speed technique. A small draft head guesses one or more tokens ahead of the main model. The main model then verifies those guesses in a single forward pass. Any correctly-predicted tokens are accepted for free. If 80% of draft tokens are accepted, you are effectively getting 1.8× the output work per pass.

The catch: the draft head needs to predict like the main model. If the draft head was trained on full-precision weights but the main model was quantized using QAT (Quantization-Aware Training — where the model was explicitly trained to behave correctly at Q4_0 precision), their output distributions diverge. The draft head guesses what the full-precision model would have said. The QAT model disagrees more often and accepts fewer of those guesses.

Switching to QAT-matched heads — draft heads trained against the same QAT checkpoint — closes that gap substantially:

For the 26B-A4B model the gap was especially striking — nearly 35 percentage points of acceptance rate were being lost purely to the head mismatch, not to any fundamental limit of speculative decoding.

Files

These are the draft heads only. The main models (doing the actual generation) come from Google's official repos listed in the "Pairs with" column.

Measured Performance (AMD Strix Halo, Vulkan/RADV)

Hardware: AMD Ryzen AI Max+ 395 (Strix Halo APU), 128 GB LPDDR5X (~256 GB/s bandwidth), 96 GiB GTT pool. Backend: Atomic TurboQuant llama.cpp, b9360-era Vulkan/RADV.

Wall time is normalized to a standard 1,150-token input / 2,000-token output task — a useful cross-model comparability metric. Prefill measured at ~1,400 tokens using a cache-busted random prompt (unique salt per run, so no KV cache hits).

Why is 26B-A4B so much faster than 31B? The 26B-A4B is a Mixture-of-Experts (MoE) model — it only activates ~4B parameters per output token, so the GPU reads far less memory per step. The 31B is a dense model that reads all 31B parameters from RAM on every single token. On memory-bandwidth-limited hardware like this APU, that difference dominates.

The 26B-A4B plain row wins the 2-slot aggregate because MTP with PARALLEL=2 was crashing on the upstream build at the time of measurement (see note below). A fix has been submitted upstream in PR #26 — once merged, MTP should be competitive in the 2-slot column as well.

Usage

Build and install the Atomic TurboQuant fork, then pass the draft head via --mtp-model:

LLAMA_PIPELINE_DEPTH2=0 llama-server \
  --model /path/to/gemma-4-26B-A4B-it-qat-q4_0.gguf \
  --mtp-model /path/to/gemma-4-26B-A4B-it-qat-assistant-MTP-Q8_0.gguf \
  --mtp-draft-n 3 \
  --draft-p-min 0.75 \
  --ctx-size 12288 \
  --cache-type-k q8_0 \
  --cache-type-v q8_0 \
  --n-parallel 1 \
  --port 8080

What each flag does:

PARALLEL=2 Status

--n-parallel 2 crashed in llm_build_gemma4_mtp with an assertion failure in ggml_reshape_3d on the upstream build at the time of these benchmarks. The root cause was n_tokens being used as the third reshape dimension when the MTP draft step always processes exactly one token column — regardless of how many server slots are active.

A fix (3 files: gemma4-assistant.cpp, llama-graph.cpp/h, llama-context.cpp) has been tested and submitted upstream as PR #26.

Until PR #26 is merged: use --n-parallel 1. For concurrency, run two separate single-slot servers behind a load balancer. After PR #26 merges: --n-parallel 2 should work and the MTP 2-slot numbers will be worth re-running.

How These Were Made

Source repos — Google's official QAT assistant checkpoints, published as unquantized safetensors on Hugging Face:

• —google/gemma-4-12B-it-qat-q4_0-unquantized-assistant
• —google/gemma-4-26B-A4B-it-qat-q4_0-unquantized-assistant
• —google/gemma-4-31B-it-qat-q4_0-unquantized-assistant

Process:

1. 1.Download the unquantized assistant checkpoints from the Google repos above
2. 2.Convert to GGUF using the Atomic fork's convert_hf_to_gguf.py with the gemma4_assistant architecture path
3. 3.Quantize to Q8_0 with llama-quantize

12B note: Google's 12B source repo uses a newer config class name that the converter does not recognize. The 26B-A4B and 31B both have "model_type": "gemma4_assistant" and convert without modification. The 12B has:

"model_type": "gemma4_unified_assistant"
"architectures": ["Gemma4UnifiedAssistantForCausalLM"]

Fix: temporarily edit those two fields in config.json to match the older names before running convert_hf_to_gguf.py, then restore the file afterward.

# in config.json, change:
#   "model_type": "gemma4_unified_assistant"  →  "model_type": "gemma4_assistant"
#   "architectures": ["Gemma4UnifiedAssistantForCausalLM"]  →  ["Gemma4AssistantForCausalLM"]
# then convert, then restore config.json

No weight values are edited — only two string fields in the metadata file so the converter recognizes the architecture.

The underlying weights are Google's. These GGUFs are a format and quantization conversion, not original model training.

License

These files are subject to the Google Gemma license. See the Gemma Terms of Use. By downloading or using these files you agree to those terms.

Related

About

These benchmarks and conversions come out of work on running local LLM infrastructure for water-utility and R&D applications — where data stays on-premise and the model has to earn its place on real operational tasks, not just leaderboard numbers.

If that framing interests you:

• —[tesla_agent](https://github.com/boxwrench/tesla_agent) — the benchmark runbook, reproducibility matrix, and supervised water-agent starting point behind these numbers. Includes the full Strix Halo hardware guide, evaluation harnesses, and an interactive web dashboard for comparing models.
• —[of-agents-and-aquifers](https://github.com/boxwrench/of-agents-and-aquifers) — writing and thinking on local AI agents in the context of water systems and public infrastructure.

GitHub: github.com/boxwrench Blog: Title 22 — water, systems, strategy