Add MTP Support for Gemma 4

[SamuelOliveirads]

Google has released the MTP for its models as standalone files. The upside is that we don’t need to download the model again, the downside is all the work I had to do to get their architecture to work.

In any case, this PR enables the use of MTP for Gemma 4 models. As a prerequisite, download the MTP, you can download the ones I made for testing:

• gemma-4-31B-it-MTP
• gemma-4-E2B-it-MTP

Or you can make your own, I’ve updated the GGUF conversion files to make it easier.

To use MTP, use these arguments:

`--spec-type mtp`, `-md {mtp_model_file}`, `-ngld 99`, `--draft-max 3`, and `--draft-p-min 0.0`

If you want to test other draft-max and draft-p-min options, feel free to do so, or just use --spec-autotune, an algorithm that tests each generated token to find the best fit for your scenario.

Benchmark - 1 run with max-tokens 1000

Tested with the Bartowski model.

Run	Code	Extract	Story	Overall
Baseline no MTP	21.7 t/s	20.5 t/s	21.1 t/s	21.1 +/- 0.6 t/s
Q8 MTP --draft-max 1	37.5 t/s, 96.1% accept	35.4 t/s, 96.1% accept	32.4 t/s, 74.5% accept	35.1 +/- 2.6 t/s, 88.9% accept
Q8 MTP --draft-max 2	49.5 t/s, 94.7% accept	43.8 t/s, 89.1% accept	37.2 t/s, 63.7% accept	43.5 +/- 6.2 t/s, 82.5% accept
Q8 MTP --draft-max 3	56.1 t/s, 85.2% accept	51.9 t/s, 86.4% accept	37.8 t/s, 50.5% accept	48.6 +/- 9.6 t/s, 74.0% accept
Q8 MTP --draft-max 4	57.8 t/s, 76.0% accept	52.9 t/s, 77.3% accept	34.8 t/s, 38.4% accept	48.5 +/- 12.1 t/s, 63.9% accept

Arguments used:

GGML_CUDA_NO_PINNED=1 CUDA_VISIBLE_DEVICES=1,0 ./build_release/bin/llama-server \
    --model '/media/samuel/llm_models/gemma-4/google_gemma-4-31B-it-Q8_0.gguf' \
    --ctx-size 20480 --main-gpu 1 -ctk q8_0 -ctv q8_0 --n-gpu-layers 99 \
    -b 128 -ub 128 --no-mmap --jinja --seed 42 \
    --threads 24 --threads-batch 32 --host 127.0.0.1 --port 18080 \
    --no-warmup --temp 0 --top-k 1 --top-p 1 \
    --spec-type mtp \
    -md '/media/samuel/llm_models/gemma-4/gemma-4-31B-it-assistant-Q8_0.gguf' \
    -ngld 99 --draft-max N --draft-p-min 0.0

Next steps:

• Currently -sm graph is broken, since it’s getting late, I’ll leave that for tomorrow, but I want to open the PR now so others can test it.
• Models with embeddings are supported, but I still need to test them, so don’t set your expectations too high for now.

[pestopoppa]

Pure-CPU EPYC bench: 21.02 t/s vs 7.05 t/s baseline = 2.98× speedup, 84.3% per-token acceptance

Tested PR #1744 on AMD EPYC 9655 (96 cores, 1.13 TB DDR5-5600, no GPU) with Gemma 4 31B-it Q4_K_M target + your gemma-4-31B-it-assistant-Q8_0.gguf drafter, draft-max=3, draft-p-min=0.0.

First attempt segfaulted on first request because of a chicken-and-egg in params_use_gemma4_external_mtp:

 static bool params_use_gemma4_external_mtp(const gpt_params & params_base) {
     return params_base.has_mtp &&
-        params_base.speculative.type == COMMON_SPECULATIVE_TYPE_MTP &&
         model_has_arch(params_base.speculative.model_dft, "gemma4_mtp");
 }

The params_base.speculative.type == COMMON_SPECULATIVE_TYPE_MTP precondition is itself set as a consequence of this helper returning true (slot init at L357 sets it inside the if (has_external_mtp) branch). With the precondition still in place, the helper always returns false → gate at L356 falls through to the "Disabling speculative" warning → downstream NULL deref. Removing the type check (the has_mtp flag plus a gemma4_mtp drafter arch is sufficient evidence of intent) clears the segfault.

After the fix:

Run	Throughput	Acceptance	Speedup
Baseline (no MTP)	7.05 t/s	—	1.0×
MTP draft-max=3	21.02 t/s	84.3% per-token (91/108)	2.98×

Notably higher speedup than your mixed-CPU/GPU bench (2.3×) — the small 500 MB drafter amortizes well against the slow dense 31B target on CPU; MTP is relatively more impactful when the target is BW-bound.

Two minor nits remaining: the size-accounting iteration at src/llama.cpp:2509 emits Oops: tensor with strange name mtp_post_proj.weight / mtp_pre_proj.weight for the two top-level (non-blk.*) tensors. Cosmetic — create_gemma4_mtp_tensors loads them correctly, so generation is unaffected. Could be silenced by special-casing those two names alongside token_embd.weight / output_norm.weight in that loop.

Hardware: AMD EPYC 9655 (96C/192T Zen 5), 1.13 TB DDR5-5600 ECC, NPS4, ~460 GB/s aggregate BW. Settings: taskset -c 0-95, numactl --interleave=all, OMP_PROC_BIND=spread OMP_PLACES=cores OMP_WAIT_POLICY=active, -t 96 -fa 1 --no-mmap, deterministic prompt + --temp 0 --top-k 1 --top-p 1 --seed 42.

[pestopoppa]

Cosmetic follow-up: silencing the two MTP-related "Oops: tensor with strange name" warnings

After my pure-CPU bench comment above, one nit remained: the size-accounting iteration at src/llama.cpp:2509 doesn't recognize the two non-blk.* top-level tensors that PR #1744 introduces (mtp_pre_proj.weight, mtp_post_proj.weight) and emits a warning for each:

Oops: tensor with strange name mtp_post_proj.weight
Oops: tensor with strange name mtp_pre_proj.weight

The warnings are cosmetic — create_gemma4_mtp_tensors loads these correctly via the proper LLM_TENSOR_MTP_PRE/POST_PROJ wiring; this iteration is just a memory accounting accumulator that special-cases the existing top-level tensors (token_embd.weight, output.weight, output_norm.weight) and warns on anything else.

One-line patch to silence them, alongside mtp_centroids / mtp_token_ordering for defensive future-proofing (these are TENSOR_NOT_REQUIRED in create_gemma4_mtp_tensors so they're absent from the current Q8_0/F16 drafters but may appear in future variants):

diff --git a/src/llama.cpp b/src/llama.cpp
@@ -2504,6 +2504,12 @@
         if (name == "output_norm.weight") {
             continue;
         }
+        // Top-level (non-blk.*) tensors introduced by Gemma4Assistant (gemma4_mtp arch).
+        // Loaded by create_gemma4_mtp_tensors; not per-layer, so they don't participate in this accounting.
+        if (name == "mtp_pre_proj.weight" || name == "mtp_post_proj.weight" ||
+            name == "mtp_centroids.weight" || name == "mtp_token_ordering.weight") {
+            continue;
+        }
         auto pos = name.find("blk.");
         if (pos != 0) {
             LLAMA_LOG_WARN("Oops: tensor with strange name %s\n", name.c_str());

Verified locally on top of 0f728f380 + my earlier params_use_gemma4_external_mtp fix: rebuild clean, Oops: ... mtp_post_proj.weight and Oops: ... mtp_pre_proj.weight both gone, MTP runtime unaffected (still MTP needs embeddings on decode, enabling + using external MTP assistant context).

Note: there is a separate pre-existing Oops: tensor with strange name rope_freqs.weight warning that fires for the target-side Gemma 4 31B base GGUF on this same code path — that's unrelated to PR #1744 (it would fire on main too) and probably belongs in a separate cleanup. Out of scope here.

---

Bonus: in-house GGUF conversion for the 26B-A4B drafter + bench result

While testing, I also exercised PR #1744's convert_hf_to_gguf.py Gemma4AssistantModel class on google/gemma-4-26B-A4B-it-assistant (which doesn't have a community GGUF yet — only 31B does via Radamanthys11). Conversion is clean: 840 MB safetensors → 839 MB BF16 GGUF (48 tensors) → 441 MB Q8_0 via llama-quantize. No errors, all tensor mappings resolved.

26B-A4B + MTP bench on the same EPYC 9655 setup:

Run	Throughput	Acceptance	Speedup
26B-A4B baseline (no MTP)	41.49 t/s	—	1.0×
26B-A4B + MTP draft-max=3	44.12 t/s	58.7% per-token (81/138)	1.06×

Much weaker than 31B Dense (2.98×, 84% accept) — confirms the MoE-batch=1 expert-routing-cancellation behavior that's been reported elsewhere. The smaller 16/8-head drafter struggles to predict MoE expert routing, AND the verifier loads up to K×8 distinct experts per K accepted tokens, eroding the bandwidth savings that make MTP win on dense.

Empirical conclusion: on pure CPU, dense Gemma 4 + MTP is the universally favorable case at batch=1; MoE batch=1 is marginal regardless of MTP correctness. Worth flagging in any deployment guidance the PR ends up shipping with.

[SamuelOliveirads]

Gemma 4 MTP is functional, however it does not currently run on a split graph due to the complexities inherent in this MTP architecture.

Here are the tests conducted for E2B:

Run	Code	Extract	Story	Overall
Baseline no MTP	136.1 t/s	123.5 t/s	133.2 t/s	130.9 +/- 6.6 t/s
Q8 MTP --draft-max 1	181.4 t/s, 70.6% accept	177.1 t/s, 76.5% accept	161.6 t/s, 54.7% accept	173.4 +/- 10.4 t/s, 67.3% accept
Q8 MTP --draft-max 2	172.7 t/s, 47.2% accept	158.8 t/s, 48.0% accept	141.8 t/s, 33.3% accept	157.8 +/- 15.5 t/s, 42.8% accept
Q8 MTP --draft-max 3	162.3 t/s, 34.3% accept	154.4 t/s, 34.0% accept	129.8 t/s, 22.5% accept	148.8 +/- 17.0 t/s, 30.2% accept
Q8 MTP --draft-max 4	133.9 t/s, 22.2% accept	134.3 t/s, 24.9% accept	119.6 t/s, 17.4% accept	129.3 +/- 8.4 t/s, 21.5% accept

Currently, MTP does not support the embedding layer, which may be linked to the drop in the acceptance rate. Since these are small models, I don’t see why we should invest so much effort in them, given that many should be focusing on 31B and 26B.

---

@pestopoppa Thank you for the tests, I’ve fixed the issues you pointed out.

[ikawrakow]

From the testing I have done it seems to be working. Only tested with Gemma4-31B

I'm seeing nearly 3X TG performance improvement running CPU-only, ~2.4X running with full GPU offload (and having the MTP output tensor on the GPU as per one of the comments).

We again suffer a significant PP performance loss. 1.5X compared to split mode layer, 2X+ compared to split mode graph when running on the GPU. On the CPU it is less but still not negligible (10% for 4 draft tokens, but also seems dependent on the the number of drafted tokens).

But overall this is a game changer for using Gemma4. Hopefully when merged this will motivate me to add SWA KV cache compression.

[ikawrakow]

We cannot use split mode graph for the main model and split mode layer for the draft model? I guess I'll see why as I go on with the review.

[SamuelOliveirads]

I believe I've fixed all the issues you pointed out, most of them were indeed errors.

The graph split was a bigger headache, so I was forcing it in layer mode, but now I've been able to spend more time building a K/V that MTP can use safely without triggering an illegal memory error.

[ikawrakow]

We want to find the last layer that matches is_sliding, right?
Then why not

int target_il = target_n_kv_layer-1;
for (; target_il >= 0; --target_il) {
    if ((target_hparams.swa_layers[target_il] != 0) == is_sliding) {
        break;
    }
}

?

[ikawrakow]

Do we have use cases where we want to return an empty graph? Because if not, we might as well throw or abort here. If we abort, we immediately know that we should never be building Gemma4 MTP graphs without input tokens. The way it is right now someone reading the code will need to go and check what happens when we get back an empty graph.

[ikawrakow]

Are there use cases where we want to run with the token embeddings being a copy of the hidden state? Because if not, we might as well abort here to reduce the mental burden for the reviewer and people reading this code later.

[ikawrakow]

The fact that it runs without hitting an assert tells me that n_tokens is never greater than lctx.mtp_target_ctx->kv_self.n (or we never have a target ctx). Because if n_tokens ever became the value of target_n_kv, we will crash later as n_kv needs to be padded to a multiple of 256 (FA on) or 32 (FA off). I.e., also this is an unnecessary mental burden.

[ikawrakow]

Do we anticipate to ever have mtp_target_ctx being not NULL and the model it holds being different from the model pointed to by mtp_target_model. Because if we don't, mtp_target_model should not be there as we can get the target model from the target context ptr.

[ikawrakow]

This means the MTP model output tensor is always on the CPU. Isn't this costly when running on the GPU? Normally, for models that use the token embedding tensor as the output tensor we use

    model.output = create_tensor(out_ctx, tn(LLM_TENSOR_OUTPUT, "weight"), {n_embd, n_vocab},
                                        llama_model_loader::TENSOR_NOT_REQUIRED);

    // if output is NULL, init from the input tok embed
    if (model.output == NULL) {
        model.output = create_tensor(out_ctx, tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab},
                        llama_model_loader::TENSOR_DUPLICATED);
    }

This will duplicate the tensor on the GPU holding the output layer.

OK, made the change and ran a quick test. TG increased from 97.7 t/s to 103.5 t/s.

[ikawrakow]

Brackets around has_mtp && embd

[ikawrakow]

This change is not related to the Gemma4 MTP functionality. Why do we need it?

[ikawrakow]

OK, this proves that the mtp_target_model member is not required.

[ahoplock]

Not sure if this is obvious to others or not but using --k-cache-hadamard --v-cache-hadamard tanked draft acceptance rate to basically zero when using gemma-4-26B-A4B on CPU-only (along with tanking t/s).

After removing those arguments, I'm seeing ~25% tg improvement on simple prompts and ~10% improvement in wall time on mtp-bench.py using --draft-max 2 (3 and 1 were both a bit slower).

[ikawrakow]

Not sure if this is obvious to others or not but using --k-cache-hadamard --v-cache-hadamard tanked draft acceptance rate to basically zero when using gemma-4-26B-A4B on CPU-only (along with tanking t/s).

Thank you for pointing this out. No, it is not immediately obvious. Yes, Hadamard transforms do not work at this point. The draft model uses the KV cache of the main model, so the keys and values have a Hadamard transform applied to them when using -khad, -vhad, but the draft model does not take this into account, hence we get nonsense results from the draft model (and correspondingly MTP performance tanks). I noticed this while reviewing, but thought it is better to first get this PR merged and then worry about Hadamard transforms when using MTP.

[TomTheWise]

Is there any consents if gemma 4 KV Cache quantization is a good or bad idea anyway? Are there new benchmarks available, because AFAIK the perplexity benchmark is worthless with the IT models.

Ministral-3 for example did not care at all, others are super sensitive - and with gemma 4 im currently not using KV Cache quant purely because im uncertain how bad it truly is on long contexts.

[ikawrakow]

@TomTheWise As you point out, one cannot deduce from PPL if quantizing the KV cache has a major impact for Gemma4-it. But if the base Gemma4 models are any indication, it should be OK to quantize down to Q4_0 with Hadamard.

[SamuelOliveirads]

Graph split is supported, so a new benchmark for 31B was run using the following arguments -sm graph -smgs -sas -mea 256 --split-mode-f32:

Run	Code	Extract	Story	Overall
Baseline no MTP	33.2 t/s	31.9 t/s	32.2 t/s	32.4 +/- 0.7 t/s
Q8 MTP --draft-max 1	50.5 t/s, 96.1% accept	41.2 t/s, 95.7% accept	41.5 t/s, 77.3% accept	44.4 +/- 5.3 t/s, 89.7% accept
Q8 MTP --draft-max 2	57.8 t/s, 92.3% accept	42.7 t/s, 89.3% accept	40.2 t/s, 63.4% accept	46.9 +/- 9.5 t/s, 81.6% accept
Q8 MTP --draft-max 3	64.2 t/s, 86.2% accept	45.2 t/s, 86.1% accept	37.3 t/s, 49.4% accept	48.9 +/- 13.8 t/s, 73.9% accept
Q8 MTP --draft-max 4	63.8 t/s, 76.0% accept	44.0 t/s, 77.9% accept	32.9 t/s, 38.7% accept	46.9 +/- 15.6 t/s, 64.2% accept

In the current setup, the main model runs in graph-split mode while the MTP remains in layer mode.

[danganbenpa]

Tried this PR on a single-GPU setup (RTX 3090, Windows, CUDA, build 4472 / commit 57e949c). Crashes consistently at build_gemma4.cpp:395 with GGML_ASSERT(has_target_ctx && "Gemma4 MTP assistant requires a bound target context") failed, immediately after the warning ggml_backend_cuda_context: a context for device 0 already exists?. Tested with KV f16, KV q8_0, KV q4_0, with and without GGML_CUDA_NO_PINNED=1, with and without explicit -mtp flag. Target loads fine, draft loads fine, fails at MTP context binding step. Only difference vs your benchmark setup is single-GPU vs CUDA_VISIBLE_DEVICES=1,0 — looks like the context-binding logic might assume distinct contexts for target and draft.

[SamuelOliveirads]

@danganbenpa You shouldn't have any MTP context issues without enabling it with -mtp, but you could try it with the new commits anyway? I removed that assert.

[TomTheWise]

This is a geamechanger!
On my PC with the 26B-A4B and using --cpu-moe as my GPU has not enough RAM, I get roughly 40% increase in tg!!

So in certain szenarios even for the MoE it is definitely worth it!!