MTP: option to use re-quantized output tensor for better TG performance

[ikawrakow]

This PR improves TG performance for Qwen3.5/3.6 models by 5-10% when using MTP. The performance gain depends on the quantization type of the model output tenor (see below).

The output tensor of the Qwen3.5/3.6 models is giant due to the very large vocabulary (248320). For the 27B dense variant the output tensor has 1.27 billion elements! The MTP layer uses the output tensor of the main model. The other tensors in the MTP layer are small in comparison: about 425 million for the 27B dense variant. This means that when generating draft tokens the bulk of the time is spent in the matrix multiplication with the output tensor. As TG is memory bound, one can expect increased MTP performance by decreasing the size of the output tensor. But using fewer bits for the output tensor will also decrease the quality of the main model. Hence, it is better to create a new output tensor, to be used only by the MTP layer, that is quantized with fewer bits than the original output tensor. This is what this PR does.

Initially I was trying to add an additional output tensor for the MTP layer while quantizing (or re-quantizing) the model. But adding an extra tensor to an existing GGUF context proved more difficult than expected, so I (for now) abandoned this approach. Instead, the additional output tensor is created on-the-fly while loading the model when

--mtp-requantize-output-tensor new_type
or
-mtprot new_type

is added to the command line. new_type is the name of the quantization type as defined in ggml (the name must match exactly), e.g., iq4_ks, q4_0, etc.

Picking the quantization type is a matter of experimentation. The fewer bits are used, the faster the draft generation. But as the quality of the quantization decreases with decreasing bits, one might expect that the acceptance rate will also decrease, which will then reduce observed TG performance. In the testing I have done, using -mtprot iq4_ks provided in the range of 10% improvement for Qwen3.6-27B-Q8_0 (so, the output tensor is Q8_0) without actually decreasing acceptance rate within statistical uncertainty.

Beware

• As the output tensor is very large, re-quantizing it while loading the model can take a non-negligible amount of time. Hence, avoid using quantization types that are very slow (e.g., Trellis quants, IQ2_KS). I will try to add the ability to do the re-quantization offline with llama-quantize in a follow up PR.
• Adding a re-quantized version of the output tensor obviously requires additional RAM/VRAM. For instance, for dense Qwen3.6-27B with -mtprot iq4_ks about 645 MiB are required. Hence, don't use this option if short on RAM/VRAM (or, if you still decide to use it, don't complain if you get OOM).

[ahoplock]

I know it's already merged at this point but you have this as --mtp-requantize-output-tensor (missing 'd') in line 3036 and in the PR description.

[zeljkokalezic]

Pure magic, thanks for this, TG throughput jumped ~24% average — massive win 💪

Setup

• Model: Qwen3.6-27B-MTP-Q5_K_XL on dual GPU (5060Ti + 3090)

    | config                   | pp t/s | pp ttfr ms | tg t/s | peak tg |
    |--------------------------|--------|------------|--------|---------|
    | Q5_K_XL (baseline)       | 1009   | 9185       | 54.36  | 64.33   |
    | Q5_K_XL + -mtprot iq4_ks | 992    | 9354       | 64.09  | 74.67   |

Run command:

    ~/Desktop/ik_llama.cpp/build/bin/llama-server \
  -m /home/zeljko/.lmstudio/models/unsloth/Qwen3.6-27B-MTP-GGUF/Qwen3.6-27B-UD-Q5_K_XL.gguf \
  --host 0.0.0.0 --port 8080 \
  --split-mode graph --tensor-split 5,2 \
  --cache-type-k q8_0 --cache-type-v q8_0 \
  --ctx-size 200000 \
  --flash-attn on \
  --jinja --parallel-tool-calls \
  --temp 0.6 --top-p 0.95 --top-k 20 \
  -ngl 999 \
  -mtp --draft-max 15 -mtprot iq4_ks \
  --no-mmap \
  -smf16 \
  -grt f16 \
  -sas \
  -np 1