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. Download the unquantized assistant checkpoints from the Google repos above
2. Convert to GGUF using the Atomic fork's convert_hf_to_gguf.py with the gemma4_assistant architecture path
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 — 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 — 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