llama + spec: MTP Support #22673
Draft
am17an wants to merge 9 commits intoggml-org:masterfrom
Draft
Conversation
github-actions
Bot
added
model
Contributor
|
Nice, I think this is a fresh start better than my WIP #18886 (that I still never find the time to continue) There were some other attempts to add MTP support but they all heavily rely on host <--> device data copy. I assume you tried addressed this, right? (Maybe there was a discussion somewhere but I wasn't aware of) |
ngxson
reviewed
May 4, 2026
Contributor
ngxson
left a comment
ngxson
left a comment
There was a problem hiding this comment.
(not a review, but opening some discussions)
Comment on lines
+73
to
+74
| // number of recurrent-state snapshots per seq for rollback; tensors are widened to (1 + n_rs_seq) groups | ||
| uint32_t n_rs_seq = 0; |
Contributor
There was a problem hiding this comment.
not 100% sure but maybe the naming with _seq is a bit confusing (or I'm misunderstanding this)
I imagine that we want to keep a buffer ring style of recurrent-state(s), similar to SWA in KV cache, right? if that's the case, probably better call it n_rs_window
Contributor
Author
There was a problem hiding this comment.
Yes this is partly the review comment from here #22400 (comment)
|
|
||
| for (int il = 0; il < n_layer; ++il) { | ||
| // MTP/NextN layers are loaded as extra decoder blocks but not executed in the main pass. | ||
| const int n_transformer_layers = n_layer - (int)hparams.nextn_predict_layers; |
Contributor
There was a problem hiding this comment.
nits, but maybe call it n_main_layers, as technically nextn layer is also a transformer layer
Comment on lines
+811
to
+823
| //TODO: generalize if this is ok, we should load <arch_name>_mtp arch? | ||
| if (params_base.speculative.type == COMMON_SPECULATIVE_TYPE_MTP) { | ||
| SRV_INF("loading MTP head from '%s' (override_arch=qwen35_mtp)\n", | ||
| params_base.model.path.c_str()); | ||
|
|
||
| auto mparams_mtp = common_model_params_to_llama(params_base); | ||
| mparams_mtp.override_arch = "qwen35_mtp"; | ||
|
|
||
| model_mtp.reset(llama_model_load_from_file(params_base.model.path.c_str(), mparams_mtp)); | ||
| if (model_mtp == nullptr) { | ||
| SRV_ERR("failed to load MTP head from '%s'\n", params_base.model.path.c_str()); | ||
| return false; | ||
| } |
Contributor
There was a problem hiding this comment.
if you look at #18886, the better way is to move llama_graph_type to the public API, then load the context with the appropriate graph type
Contributor
Author
There was a problem hiding this comment.
Yes that seems like the correct way to do this if we want to support MTP in a generic way
Contributor
Author
|
@ngxson yes the h2d was discussed with GG, he's working on a refactor which will allow us to share tensors between two llama context |
Member
|
Great work, this should massively bridge the TG gap with vLLM, or maybe even surpass it together with tensor-parallel. |
Currently speculative checkpoint needs to restart from a checkpoint after some draft tokens are not accepted, this leads to some wastage in running the target again. This PR adds the ability to rollback upto `draft_max` by storing the GDN intermediates.
Contributor
|
in my opinion Qwen 3.6 is the most important thing that happened in open source models in a long time, this is going to be so valuable. ngram could be set to match only very strong and long candidates - for large repetitive paraphrasing |
|
" idea is that MTP should automatically start and we shouldn't need to distribute the MTP gguf separately but also it has it's own context/kv-cache etc." -> Does this mean MTP needs additional resources (RAM/VRAM?) If so, there should always be an option to remain to disable it. Right now on my system (6 GB VRAM, 32 GB RAM), speculative decoding just makes things much slower even on very small draft models because of that exact reason, they need own context and kv-cache. Such low to midrange systems already operate on the edge in terms of memory. |
|
I'm getting garbage responses running this PR on the Vulkan backend with an R9700 using llama-server. I'm using the GGUF you linked above. Interestingly, draft acceptance is only 0.01282. Prompt: "Hello!" |
Contributor
Author
|
@cmp-nct I'm not sure, but could be possible @Dampfinchen as of right now it is opt-in via @mbednarek360 |
|
Might it be possible/useful to run the draft model on a second GPU? Given that MTP weights model are relatively small this might provide a useful speedup on systems with a dedicated high-VRAM "AI" GPU with a cheaper low-VRAM "normal" GPU used for display output, etc... possibly prevent some degree of resource contention. |
|
Thank you, we are eagerly awaiting this to become stable, here automated test results for my machine; __
Result:
|
Contributor
Author
|
@cturan Thanks for testing, I'm aware of the issue for the prefill and will work on a fix. |
|
Might be a long shot, but any chance of supporting MTP with a reduced vocabulary? MTP layers are rather chonky and reducing token embeddings might help users with less VRAM by filtering out certain languages. Obviously the full model will still be able to produce those tokens if need be so it won't be gimped. |
|
Working on taking this for a spin with the Q4_K_M quant of Qwen3.6-35BA3B. I was gonna try to start from unsloth's quant since they already perform really well, but of course they don't have any mtp layers. @am17an Think it would work if I just "steal" the layers from your q8 quant and merge them into the unsloth quant? (add blk.40 and bump some top-level config like block_count and kv_count) |
|
only a quick test run, 1x 5090 qwen3.6-27b mtp 3, q4_0 quantized, kv also q4_0 same model, same config (except mtp) prompt „create a flappy bird clone“ (I‘m not creative, sorry) Great Speedup! |
|
Observing very nice speed boost on 5090 with RDson/Qwen3.6-27B-MTP-Q4_K_M-GGUF! I also tried loading this model https://huggingface.co/lyf/Qwen3.6-27B-uncensored-heretic-v2-NVFP4-MTP but I get |
M1 Ultra: Qwen3.6 27B Q8_0 went from 17 t/s to 25 t/s with |
Member
That message translates to ELI5 as "this GGUF has the MTP layers stripped". |
|
Google has just released the MTP layers for Gemma 4: https://huggingface.co/google/gemma-4-26B-A4B-it-assistant |
|
I just tested it with the iGPU of my R9 6900HX, using ROCm + force-host-allocation-APU + Qwen3.5-4B Q8: I manually converted Qwen3.5-4B to GGUF using convert_hf_to_gguf.py and these are the results: Prompt:Both gave the correct answer No MTP ~6.6 tok/sMTP `--spec-type mtp --spec-draft-n-max 3` ~11.6 tok/sEdit: Add GGUF link: https://huggingface.co/adriabama06/Qwen3.5-4B-q8_0-am17an-MTP |
|
Can confirm with latest changes in Vulkan I no longer have issues with token generation suddenly stopping. At 120k tokens, the t/s generation dropped to 45t/s so that's still a huge boost in performance! |
Merged
3 tasks
ggerganov
reviewed
May 5, 2026
Comment on lines
+243
to
+248
| void handle_mtp_for_ubatch( | ||
| int32_t n_tokens, | ||
| const llama_token * tokens, | ||
| const llama_pos * positions, | ||
| struct ggml_tensor * t_h_pre_norm); | ||
|
|
Member
There was a problem hiding this comment.
This logic has to be extracted into the server/speculative contexts. One of the main reasons is that just the token positions are not enough to be able to do this correctly. For example, for multi-modal use cases, this implementation does not take into account that multiple tokens have the same position and there is no way to resolve this inside llama_context.
Also, restoring prompts in a slot from the prompt cache does not work because we forget the computed embeddings for the prompt.
Hence, we have to actually extract the embeddings from the context (in a similar way as we extract the logits) and manage them along side the prompts (store them, prefix cache them, etc.). Quite a lot of work is needed, but I think that's the proper way to implement support for these methods.
For MTP, we can already use the existing embeddings API because MTP only needs the output embeddings and we already have the API for that llama_get_embeddings_* (edit: on second look, we actually need the embeddings before the output norm, so need API for that as well). But for Eagle3 this is not enough. For that I am preparing new embeddings API (#22728) that will be used to extract the internal layer embeddings. Feedback is welcome.
Contributor
Author
There was a problem hiding this comment.
@ggerganov IMO the only blocker to not run this inside the same llama graph was that kv-cache becomes weird to handle. Perhaps we can restructure the kv-cache such that there is an auxiliary cache for these "speculators", then everything becomes kind of already working from a sync perspective. Or I am maybe missing something
nybblr
pushed a commit
to nybblr/llama.cpp
that referenced
this pull request
May 5, 2026
Without this override, upstream PR ggml-org#22673 MTP context inherits n_ubatch=1024 from the main context, sizing its compute buffer at ~1.1 GiB. At 150k on 24GB, that triggers cudaMalloc OOM during MTP context init. Override drops the buffer to ~362 MiB at n_ubatch=256 and lets the MTP context fit. Decode bench post-patch reveals the real wall is MTP head weight format mismatch (3.4% accept on this lab's sidecar GGUF vs 72-82% in PR description), not VRAM. See RESULTS_UPSTREAM_MTP_22673.md in the lab repo.
|
Just as a note: Gemma 4 MTP support will likely require other changes. From the blog post:
In the transformers implementation, Google also uses a heuristic approach to dynamically adjust how many tokens are drafted at each step, which could be nice to have for all MTP models. |
|
Did a clean DFlash + DDTree reproduction on a single RTX 3090 (24GB) a few days ago, both Qwens, full Bench results
Both with their matching 3.5-1.5B / 3.6-1.5B drafts. Two things worth flagging on these numbers:
vs MTP on Qwen 3.6So on Qwen 3.6 specifically: ~2.56x mean for DFlash + DDTree vs the ~1.85x being reported here for MTP. Different mechanisms, different tradeoffs:
Whether they compose — MTP draft tokens fed into DDTree verification — is the interesting question, but I haven't tried it yet. Plan to reproduce MTP on the same machine when this lands. Full methodology + per-task tables at https://insiderllm.com/guides/dflash-rtx-3090-bench-both-qwens/ if useful. |
|
If the user invokes with a non-MTP model it dies with an assert: llama_model_create: overriding architecture qwen35 -> qwen35_mtp obviously I don't expect it to mtp in this case but it should probably politely reject the option and cleanly abort or ignore the spec option. Aside, Qwen3.6-27B-Q8 and a 256k context on RTX A6000-- ~20tok/s without MTP, 55tok/s with MTP. Seems to work. Would be nice if parallel was still possible, but one step at a time! :) |
ooooof that is bad, PP takes a absolute dump, hope you're not processing any long prompts or care about how responsive the model is. Need to be doubling the PP. |
|
After transplanting from your provided MTP file to unsloth/Qwen3.6-27B-Q4_K_M, and testing it on a 4090 on Windows 10, it seems to "silently" crash if GPU VRAM is full or will be full? I don't fully understand the mode of the crash but: Common to both: First scenario - "silent" crash-fit 163840 --fit-target 1536 --ctx-size 163840 Key item: Second scenario - successful start, choose small values to not overload GPU just to get it working-fit 81920 --fit-target 512 --ctx-size 81920 It shouldn't assign MTP to 0 VRAM GPU, instead it should ether fall back to CPU or emit a clear error like “cannot load second MTP model on GPU: no free memory; reduce offload or force MTP CPU.” |
Contributor
|
With all the current mainstream ggufs being without mtp would it maybe make sense to allow a secondary gguf with the mtp tensors to be loaded instead of just 1 file with both the main and mtp model? |
|
Decided to test this myself and report back about results - eh - I guess this is better off for more memory bound systems. So my setup:
Ran a baseline first then did MTP n_max 2 & 3 Aggregate
Per-prompt accept rate
Per-prompt tok/s
Couldn't get the OP's results myself. thought I'd share |
|
Apologies for the spam, but can anyone familiar explain how to create the secondary smaller model locally for qwen3.6 variants like this which don't yet have it? |
|
I just tested this on my Rog Flow z13 (Strix Halo), and this gave me a big boost. All models are Q6, except the 'normal' 27b, which was a Q5. Just used a local bench script that runs a simple test 3 times each. I believe the context size was at 128k with a q8 kV cache. The only issue I ran into was it was crashing for me when trying to run parallel slots. Model | Prefill (t/s) | Decode (t/s) ===================================== Qwen3.6-35B | 153.1 | 42.3 |
|
Results on a 7900XT w/ Vulkan, Qwen 3.5 27b IQ4_XS converted by me from hf (Q4_K_M doesn't fit w/ mtp):
Prompt being "explain this Rust code". I'm seeing ~60 with Python and ~40 in multi-turn multimodal/english workflows. Smashing work. With the pp regression fixed this will make these dense models usable on whole new classes of hardware. |
|
I ran an independent smoke benchmark of this PR using the prompt suite from the linked gist. This is not a full quality/eval run, just a quick check for whether MTP produces a measurable speedup on my local machine. Environment:
Aggregate result:
The result looks positive: MTP produced a clear speedup on all 9 prompts in this short completion suite. In this configuration, One resource note: observed RSS was about 37.2 GB for baseline and about 74.2 GB for the MTP profiles. That is fine on this machine, but probably worth calling out for users testing larger Qwen3.6 variants. Thanks for working on this. Happy to run a longer-output or longer-context variant if that would be more useful for the PR. |
You may use https://gist.github.com/buzz/1c439684d5e3f36492ae9f64ef7e3f67#file-convert-py with the 35B-A3B MTP file provided in this thread under the "How to use" subheading |
4 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
20 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Overview
This PR adds support for MTP (Multi Token Prediction) heads. I tested this on Qwen3.6 27B and Qwen3.6 35BA3B but in principle it should work for any MTP model. I've posted the detailed results below, but typically I see a steady-state acceptance of around 75% with 3 draft tokens, which is more than >2x speed-up over baseline. The design decisions I took to get to this stage are as follows:
ubatchPerformance
A simple bench for testing various prompts is here: https://gist.github.com/am17an/228edfb84ed082aa88e3865d6fa27090. Posting the results below:
Performance on DGX Spark 🧵
No MTP (baseline)
./llama-server -m ../qwen3.6-q8_0.gguf -np 1 --chat-template-kwargs "{\"preserve_thinking\": true}"MTP --spec-draft-max-n 3
./llama-server -m ../qwen3.6-q8_0-mtp.gguf -np 1 --chat-template-kwargs "{\"preserve_thinking\": true}" --spec-type mtp --spec-draft-n-max 3MTP --spec-draft-max-n 2
./llama-server -m ../qwen3.6-q8_0-mtp.gguf -np 1 --chat-template-kwargs "{\"preserve_thinking\": true}" --spec-type mtp --spec-draft-n-max 2Draft model (Qwen3.5 0.8B) with spec-draft-n-max 16 with partial rollback
llama-server -m ../qwen3.6/Qwen3.6-27B-Q8_0.gguf -hfd unsloth/Qwen3.5-0.8B-GGUF:Q8_0 --spec-draft-n-max 16 -np 1 --chat-template-kwargs "{\"preserve_thinking\": true}"Master with draft model with spec-draft-n-max 64 with no partial rollback
llama-server -m ../qwen3.6/Qwen3.6-27B-Q8_0.gguf -hfd unsloth/Qwen3.5-0.8B-GGUF:Q8_0 --spec-draft-n-max 64 -np 1 --chat-template-kwargs "{\"preserve_thinking\": true}"How to use
I've uploaded the GGUF which I made by using the
convert_hf_to_gguf.pychanges in this PR. Here is another GGUF for the MoE (35BA3B) modelRequirements