Misc. bug: current draft MTP implementation very slow input tokens digestion

[claude-eric-steiner]

Name and Version

I just ran some speed tests with MTP.
branch feature/turboquant-kv-cache, 2b61ea2

Operating systems

No response

Which llama.cpp modules do you know to be affected?

No response

Command line

Problem description & steps to reproduce

The output speed gain is real (and most excitedly far bigger in percentage gain for big context sizes), model quality tests I've also ran point to no quality degradation.
But for some reasons the input token digestion is severely decreased, which bloats up the wall time.

Just as a reference: I'm on a RTX 3090

Input tokens	qwen36_35b_a3b_iq4nl_no_mtp	qwen36_35b_a3b_iq4nl_draftmtp2	qwen36_35b_a3b_iq4nl_draftmtp3	qwen36_35b_a3b_iq4nl_draftmtp4
1,024	input 1461.19 tok/s, output 98.36 tok/s, elapsed 2.0s	input 480.90 tok/s, output 115.37 tok/s, elapsed 3.3s	input 471.46 tok/s, output 126.36 tok/s, elapsed 3.2s	input 455.82 tok/s, output 143.21 tok/s, elapsed 3.2s
8,192	input 3277.07 tok/s, output 87.38 tok/s, elapsed 4.0s	input 1568.59 tok/s, output 110.49 tok/s, elapsed 6.4s	input 1626.07 tok/s, output 109.00 tok/s, elapsed 6.2s	input 1671.36 tok/s, output 142.32 tok/s, elapsed 5.8s
32,768	input 2867.22 tok/s, output 61.33 tok/s, elapsed 13.6s	input 1704.29 tok/s, output 100.82 tok/s, elapsed 20.5s	input 1805.61 tok/s, output 119.21 tok/s, elapsed 19.3s	input 1805.14 tok/s, output 128.24 tok/s, elapsed 19.2s
65,536	input 2398.14 tok/s, output 44.47 tok/s, elapsed 30.3s	input 1463.43 tok/s, output 85.90 tok/s, elapsed 46.4s	input 1602.67 tok/s, output 101.13 tok/s, elapsed 42.2s	input 1600.79 tok/s, output 107.77 tok/s, elapsed 42.2s
131,072	input 1813.04 tok/s, output 28.48 tok/s, elapsed 77.0s	input 1194.38 tok/s, output 70.47 tok/s, elapsed 111.7s	input 1294.03 tok/s, output 77.20 tok/s, elapsed 103.1s	input 1294.96 tok/s, output 86.81 tok/s, elapsed 102.9s
196,608	input 1451.00 tok/s, output 20.60 tok/s, elapsed 142.0s	input 1079.42 tok/s, output 57.42 tok/s, elapsed 184.6s	input 1077.91 tok/s, output 64.31 tok/s, elapsed 184.6s	input 1087.32 tok/s, output 69.02 tok/s, elapsed 182.9s

The server currently treats MTP's hidden-state requirement as full embedding/output mode. That causes prompt batches to mark every token as an output row and enables normal embeddings on the target context. As a result, prefill pays for extra output buffer work and unnecessary final-head work that MTP does not consume.

There is also inherent MTP work during prefill: the prompt must be mirrored through the MTP draft context so that draft-time attention/cache state is available. That part cannot be removed. The problem is that the current implementation adds avoidable all-token output/logit/head overhead on top of that inherent sync.

The fix is to split hidden-state extraction from normal output extraction:

1. Keep llama_set_embeddings_pre_norm() as the MTP hidden-state tap.
2. Do not make MTP set normal llama_set_embeddings() on the target context.
3. During prompt batching, mark only actual output/logit rows in batch.logits.
4. Allocate and copy pre-norm hidden states by token row, independent of n_outputs.
5. In Qwen35/Qwen35MoE graphs, when pre-norm extraction is enabled, compute the full final hidden state for the prompt, then slice to output rows before final output norm and LM head.
6. In MTP prompt-sync batches with no requested logits, skip the final MTP head and expand only the hidden-state graph needed to update draft-context state.
7. Avoid copying draft-context pre-norm hidden states during prompt sync; the draft context only needs host-visible pre-norm rows during draft generation.

Expected result: MTP prefill should remain slower than no-MTP because it maintains the draft MTP context, but the severe all-token output/head overhead should be removed. Decode speed should remain improved.

This analysis was assisted by GPT 5.5

First Bad Commit

2b61ea2

Relevant log output

Logs

[atul7017128880]

please give me your cmd how you used it im new to this

[claude-eric-steiner]

So, with the fix in place it looks much better 👍

Input tokens	qwen36_35b_a3b_iq4nl_no_mtp	qwen36_35b_a3b_iq4nl_draftmtp2	qwen36_35b_a3b_iq4nl_draftmtp3	qwen36_35b_a3b_iq4nl_draftmtp4
1,024	input 1288.02 tok/s, output 118.10 tok/s, elapsed 1.9s	input 1190.01 tok/s, output 147.63 tok/s, elapsed 1.7s	input 1145.07 tok/s, output 152.45 tok/s, elapsed 1.7s	input 1277.92 tok/s, output 149.66 tok/s, elapsed 1.7s
8,192	input 3945.67 tok/s, output 104.42 tok/s, elapsed 3.3s	input 3528.47 tok/s, output 148.37 tok/s, elapsed 3.2s	input 3558.33 tok/s, output 157.01 tok/s, elapsed 3.1s	input 3367.78 tok/s, output 156.47 tok/s, elapsed 3.3s
32,768	input 3452.69 tok/s, output 76.57 tok/s, elapsed 11.2s	input 3065.30 tok/s, output 121.60 tok/s, elapsed 11.8s	input 2942.63 tok/s, output 128.69 tok/s, elapsed 12.2s	input 2922.91 tok/s, output 129.17 tok/s, elapsed 12.2s
65,536	input 2877.87 tok/s, output 54.10 tok/s, elapsed 25.2s	input 2550.80 tok/s, output 107.55 tok/s, elapsed 27.0s	input 2446.31 tok/s, output 100.64 tok/s, elapsed 28.2s	input 2413.37 tok/s, output 117.04 tok/s, elapsed 28.3s
131,072	input 2168.86 tok/s, output 34.91 tok/s, elapsed 64.3s	input 1889.45 tok/s, output 80.39 tok/s, elapsed 71.1s	input 1858.88 tok/s, output 87.21 tok/s, elapsed 72.1s	input 1808.79 tok/s, output 88.51 tok/s, elapsed 74.1s
196,608	input 1716.97 tok/s, output 25.70 tok/s, elapsed 119.8s	input 1478.25 tok/s, output 64.89 tok/s, elapsed 135.2s	input 1479.37 tok/s, output 69.51 tok/s, elapsed 135.0s	input 1410.82 tok/s, output 72.61 tok/s, elapsed 141.4s

It even seems the fix has improved the no_mtp results too 👍

Because the prompt must be mirrored through the MTP draft context so that draft-time attention/cache state is available,
there is still some overhead on the input side that we can't get rid off. The wall times for this speed tests still favor no_mtp.

But my test focused on the input size side, the output token side was aimed to be small.
The tides can turn in favor of using MTP when a small prompt needs to produce a big output.

I will conduct some multi dimensional forced input and output size tests and post them here...

[claude-eric-steiner]

I've ran some more tests with the fix. The task for the LLM was to code.
It got a specially crafted prompt (of the size target input tokens) containing the request to write very long code.
The response was then restricted to max_tokens (output tokens) and the wall time (and other stats) have been measured.

The tables below show the time savings in percent over the non MTP run both for MTP 2 and MTP 3.

qwen36_35b_a3b_iq4nl_draftmtp2

Target input tokens	max_tokens=1024	max_tokens=2048	max_tokens=4096	max_tokens=8192	max_tokens=16384
	Time saved (%)	Time saved (%)	Time saved (%)	Time saved (%)	Time saved (%)
1’024	-12%	-12%	-18%	-19%	-25%
8’192	-15%	-18%	-28%	-26%	-32%
32’768	-18%	-22%	-33%	-36%	-35%
65’536	-15%	-24%	-37%	-40%	-41%
131’072	-12%	-21%	-29%	-40%	-45%
196’608	-8%	-19%	-28%	-36%	-45%

qwen36_35b_a3b_iq4nl_draftmtp3

Target input tokens	max_tokens=1024	max_tokens=2048	max_tokens=4096	max_tokens=8192	max_tokens=16384
	Time saved (%)	Time saved (%)	Time saved (%)	Time saved (%)	Time saved (%)
1’024	-4%	-13%	-16%	-15%	-31%
8’192	-8%	-14%	-27%	-19%	-39%
32’768	-16%	-24%	-32%	-31%	-43%
65’536	-19%	-21%	-35%	-35%	-49%
131’072	-12%	-20%	-25%	-38%	-51%
196’608	-9%	-20%	-23%	-42%	-52%

[claude-eric-steiner]

PR: #149

But it will probably not be accepted because of this:

This project does not accept pull requests that are fully or predominantly AI-generated.

You can find a fork containing the fix in a separate branch here: https://github.com/claude-eric-steiner/llama-cpp-turboquant/tree/codex/mtp-prefill-hidden-state-fix

[utajum]

with the latest MTP code changes models run 50% slower

non-MTP Qwen 3.6 specifically, after I pulled the latest changes my existing setup started running 50% slower, I can revert to an older commit, but just wanted to report.