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:

Model	Non-QAT head acceptance	QAT-matched head acceptance	Change
Gemma 4 12B QAT Q4_0	71.3%	78.4%	+7.1 pp
Gemma 4 26B-A4B QAT Q4_0	56.9%	91.8%	+34.9 pp
Gemma 4 31B QAT Q4_0	42.5%	60.4%	+17.9 pp

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.

File	Pairs with	Size
gemma-4-12B-it-qat-assistant-MTP-Q8_0.gguf	google/gemma-4-12B-it-qat-q4_0	444 MiB
gemma-4-26B-A4B-it-qat-assistant-MTP-Q8_0.gguf	google/gemma-4-26B-A4B-it-qat-q4_0	441 MiB
gemma-4-31B-it-qat-assistant-MTP-Q8_0.gguf	google/gemma-4-31B-it-qat-q4_0	491 MiB

---

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.

Lane	Decode	Prefill	Wall time (std)	2-slot aggregate	Notes
26B-A4B plain, F16 KV	59.4 t/s	1,194.4 t/s	34.6 s	90.9 t/s	fastest 2-slot option
26B-A4B QAT MTP + Q8 KV	71.4 t/s	729.3 t/s	29.6 s	62.5 t/s	best single-stream lane
12B plain, F16 KV	25.7 t/s	666.5 t/s	79.5 s	47.6 t/s	plain baseline
12B QAT MTP + Q8 KV	45.6 t/s	539.9 t/s	46.0 s	43.5 t/s	single-slot; +77% vs plain
12B QAT MTP + Q8 KV, 2-slot	38.6 t/s (48.6 eff.)	550.6 t/s	53.9 s	62.5 t/s	post-fix; 88.6% acceptance; +31% vs plain 2-slot
31B plain, Q8 KV	11.0 t/s	204.2 t/s	187.4 s	20.0 t/s	dense, bandwidth-bound
31B QAT MTP + F16 KV	19.1 t/s	203.6 t/s	110.4 s	18.9 t/s	+74% vs plain

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:

Flag / env var	Value	Reason
LLAMA_PIPELINE_DEPTH2=0	env var	Required for Vulkan multi-slot stability. Prevents thread queue deadlocks when two slots are active simultaneously. Set before the binary.
--mtp-draft-n	3	Draft 3 tokens per step. Block-size 3 gave best net throughput in sweep; block-size 2 has higher acceptance but lower overall speed gain.
--draft-p-min	0.75	Only speculate when the draft head is ≥75% confident. Prunes wasted work on uncertain tokens.
--ctx-size	12288	KV cache allocation. Tune up or down based on your available GTT memory.
--cache-type-k/v	q8_0	Q8 KV cache reduces memory use with minimal quality impact, enabling longer context at a given memory limit.
--n-parallel	1 or 2	1 is the safe default. 2 works after the fix in PR #26 — see note below.

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

Resource	Link
Main Gemma 4 QAT Q4_0 models	google/gemma-4-12B-it-qat-q4_0 · 26B-A4B · 31B
ik_llama variant heads (different arch, not compatible here)	ji-farthing/gemma-4-qat-q4_0-MTP-assistants-ik-llama-GGUF
Atomic TurboQuant llama.cpp fork	AtomicBot-ai/atomic-llama-cpp-turboquant
Full benchmark data and methodology	boxwrench/tesla_agent

---

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