llama: avoid copying logits during prompt decode in MTP

[am17an]

Overview

Avoid copying the logits for every token in the batch when doing prompt processing for MTP since it only requires the pre-norm. This reduces memory traffic quite a bit and in turn increases PP speed with MTP.

Additional information

Requirements

• I have read and agree with the contributing guidelines
• AI usage disclosure: YES, for debugging and reviewing

[ggerganov]

A quick bench on RTX 5090 with Qwen3.6 27B Q4_K

[d-r-e]

A quick bench on RTX 5090 with Qwen3.6 27B Q4_K

Are the legend colors swapped?

[pwilkin]

@d-r-e no, MTP does negatively impact prompt processing, but under this PR the negative impact is halved.

[cb88]

2xMI50 qwen 27b Q4_1 does see some improvement with this PR
MI50 without MTP = 500t/s
with MTP = 250t/s
with MTP this PR = 300t/s

[0cc4m]

@d-r-e no, MTP does negatively impact prompt processing, but under this PR the negative impact is halved.

Why does it affect prompt processing?

[Mithras]

Made sure this PR is included and re-tested:
unsloth/Qwen3.6-27B-MTP-GGUF/Qwen3.6-27B-UD-Q5_K_XL.gguf

| model     |             test |             t/s |     peak t/s |          ttfr (ms) |       est_ppt (ms) |      e2e_ttft (ms) |
|:----------|-----------------:|----------------:|-------------:|-------------------:|-------------------:|-------------------:|
| qwen36-27 |  pp2048 @ d16384 | 1843.93 ± 14.82 |              |   9098.06 ± 124.65 |   9097.48 ± 124.65 |   9098.06 ± 124.65 |
| qwen36-27 |   tg128 @ d16384 |    74.46 ± 3.24 | 84.00 ± 0.82 |                    |                    |                    |
| qwen36-27 |  pp2048 @ d65536 |  1449.72 ± 9.78 |              |  42344.98 ± 292.27 |  42344.40 ± 292.27 |  42344.98 ± 292.27 |
| qwen36-27 |   tg128 @ d65536 |    61.97 ± 2.35 | 68.33 ± 4.78 |                    |                    |                    |
| qwen36-27 | pp2048 @ d131072 |  1075.30 ± 2.36 |              | 112238.48 ± 281.78 | 112237.90 ± 281.78 | 112238.48 ± 281.78 |
| qwen36-27 |  tg128 @ d131072 |    48.40 ± 2.53 | 55.00 ± 0.00 |                    |                    |                    |

pretty much the same as #22673 (comment) which probably had the PR already. Still almost 50% pp hit

[pwilkin]

Why does it affect prompt processing?

Due to the embeddings copy, most likely.

[jtjstock]

This particular commit seems to regress the acceptance rate, I lose about 5% at n = 4 in coding tasks. Otherwise the prefill improvement is amazing

[am17an]

@jtjstock this commit just stops copying stuff we don't ever read, in a way it's a free optimization and it should not affect anything except the prefill

[jtjstock]

@am17an Just did a bunch of runs, one at latest(b9254): draft acceptance = 0.75134 ( 3913 accepted / 5208 generated)
and one with this and a dependent commit reverted: draft acceptance = 0.83896 ( 4376 accepted / 5216 generated)

Same prompt. Same Seed. Temp 0. Same hardware minutes apart. Same compiler. The only oddity seems to be the outputs from the latest(b9254), they are slightly different each time, but highly similar, where as with mainline commits (3e12fbd + dependent 49c21f9, included this last one for a clean revert and build) reverted build they are identical each time.

I'm running this on windows, the command used:
llama-server.exe ^
-m models/localweights/Qwen3.6-35B-A3B-MTP-IQ4_XS-Q8nextn-GGUF/Qwen3.6-35B-A3B-MTP-IQ4_XS-Q8nextn.gguf ^
--no-mmap --ctx-size 16384 --port 12345 ^
--flash-attn on --cache-type-k q8_0 --cache-type-v q8_0 ^
--spec-type draft-mtp --spec-draft-n-max 4 ^
-np 1 --temp 0 --seed 0 ^
-sm layer -ngl 999 --tensor-split 20,18 ^
-t 4 -tb 8 --no-warmup --metrics

The prompt: write fizzbuzz in 16 different programming languages

[jtjstock]

@am17an Also just tested the official build of b9254(llama-b9254-bin-win-cuda-13.1-x64), seeing the same variability in the output and similar degraded acceptance rate.

Maybe it's specific to my config? It's running on 2x 5060ti 16GB

Edit: did some more testing. Issue still present without quantized kv, but does disappear when I switch to -sm tensor, output from -sm tensor is internally consistent across runs.

[jtjstock]

@am17an I think the issue I see, which still presents at head, is the MTP heads are being fed the raw pre-norm residual, but Qwen seems to do better with the post-norm. Before this PR, Qwen was effectively reading a post-norm because t_h_pre_norm aliased the buffer that the final norm was writing to. So it was better before because of a side effect.

It is a very small change to fix. Move the final build_norm before res->t_h_pre_norm = cur in the main graph, and move res->t_h_pre_norm = cur after the build_norm in the mtp one, in qwen35moe.cpp and qwen35.cpp. Two lines moved in each file.

I could be wrong, but it does result in a noticeable improvement to acceptance with the corresponding bump in token generation. I see acceptance move from ~75% to ~80% at n-max 4, and generation from ~203 t/s to ~225 t/s on tensor split. Prefill was of course unaffected.

I'm not looking at the sm layer issue I noted before, as I'm not using layer split anymore, but it's still present when I tested.

[am17an]

@jtjstock the model's MTP head was trained on the pre-norm hidden state. I don't understand how the post-norm residual would make it better.

[jtjstock]

Try it?

…

On Tue, Jun 2, 2026 at 12:32 AM Aman Gupta ***@***.***> wrote: *am17an* left a comment (ggml-org/llama.cpp#23198) <#23198 (comment)> @jtjstock <https://github.com/jtjstock> the model's MTP head was trained on the pre-norm residual. I don't understand how the post-norm residual would make it better. — Reply to this email directly, view it on GitHub <#23198?email_source=notifications&email_token=AEZEUVEXU7OY322LAT3LAED45ZKFJA5CNFSNUABFM5UWIORPF5TWS5BNNB2WEL2JONZXKZKDN5WW2ZLOOQXTINJZHA3TGNBRGYY2M4TFMFZW63VHNVSW45DJN5XKKZLWMVXHJLDGN5XXIZLSL5RWY2LDNM#issuecomment-4598734161>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AEZEUVGP5WQDWNYH62W6W3T45ZKFJAVCNFSM6AAAAACZBIXW32VHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHM2DKOJYG4ZTIMJWGE> . Triage notifications, keep track of coding agent tasks and review pull requests on the go with GitHub Mobile for iOS <https://github.com/notifications/mobile/ios/AEZEUVGI5QZBVZFFJEDVV7D45ZKFJA5CNFSNUABFM5UWIORPF5TWS5BNNB2WEL2JONZXKZKDN5WW2ZLOOQXTINJZHA3TGNBRGYY2M4TFMFZW63VHNVSW45DJN5XKKZLWMVXHJKTGN5XXIZLSL5UW64Y> and Android <https://github.com/notifications/mobile/android/AEZEUVHEOD6EYGRQTSJJZ4T45ZKFJA5CNFSNUABFM5UWIORPF5TWS5BNNB2WEL2JONZXKZKDN5WW2ZLOOQXTINJZHA3TGNBRGYY2M4TFMFZW63VHNVSW45DJN5XKKZLWMVXHJLTGN5XXIZLSL5QW4ZDSN5UWI>. Download it today! You are receiving this because you were mentioned.Message ID: ***@***.***>

[am17an]

No thanks. Please stop posting irrelevant comments

[jtjstock]

Sorry I am bothering you, but I don't see any source that specifically says Qwen 3.5/3.6's MTP is trained on pre-norm hidden state, I do see a paper that says the training method is undisclosed(https://arxiv.org/html/2605.09992v1), and vLLM seems to use the post-norm hidden state in the branch qwen 3.5's mtp follows.

Again, I do apologise for bothering you, you've done a tremendous amount of work on this and it is not my intention to waste your time.

[am17an]

Okay I just checked vLLM and you're right, they do pass in the post-norm hidden state, my assumption was that is was a deepseek like MTP which passes the pre norm hidden state. Thanks for pointing this out