[Piggidragon]

Problem

An MTP draft step projects onto the whole output head and samples over the whole vocabulary. For Qwen3.8 27B (248k tokens, Q6_K head, 4070 + 3060 -sm tensor) that is 1.37 ms of GPU time per drafted token, at the memory bandwidth limit of both cards, plus ~0.7 ms of logits copy and CPU sampling on the critical path of the draft chain (nsys: the GPU idles 1.07 ms after every draft head). Two drafts per step make that about 10% of a decode step. It also means a third draft token does not pay: depth 3 measured no faster than depth 2.

Change

• llama_model_params::draft_vocab: a list of token ids. After load, the model gets output_draft.weight, a copy of those rows of the output head. The rows are read from the model file, keep the head's type and buffer type, and are split by rows like the head under split mode tensor.
• llama_context_params::draft_vocab: only for an MTP context.
  • Its draft head projects onto the copy: build_lora_mm substitutes the output head through the graph params, in that context's graphs only.
  • Its logits come out through the sampled-candidates API (llama_get_sampled_logits_ith / _candidates_ith), so common_sampler samples over the subset. Backend sampling is not used then.
  • The target still verifies every draft against the whole vocabulary, so the output does not change; a token outside the list is never drafted.
• --spec-draft-vocab FILE: token ids, whitespace separated, # comments.
• llama-draft-vocab (tools/draft-vocab): builds the list from a text corpus. It takes the most frequent tokens covering --coverage of the corpus, plus the control and end-of-generation tokens.
• common_sampler_clone/copy copy only the live candidates instead of the n_vocab scratch buffer (the server clones before every verify).

The idea follows the draft vocabulary of Strata by @Niko1221 (MIT). It is credited in the README features table.

Results

Depth 2 with an English + German list

The production choice is depth 2, because German drafts worse than code and depth 3 made German slower. The list is 66k tokens from English, German (Wikipedia), code and man pages.

Qwen3.8 27B UD-Q4_K_S, production config (2 slots, q5_1 KV, MTP depth 2), 384 tokens, median of 2 runs, after the second review:

	no list, temp 1.0	list, temp 1.0	no list, greedy	list, greedy
4 English/code tasks (mean)	59.1 t/s	63.6	62.1	66.4
edit quoted code	67.8	73.0	71.5	75.9
German letter	46.9	50.0	51.5	55.2

Acceptance is unchanged (German 0.390 vs 0.386 at temp 1.0). VRAM: the 66k list takes ~266 MiB over both GPUs.

Depth sweep (first version, list from code/English only)

Mean of 4 synthetic tasks (write code, edit quoted code, explain, list), 384 tokens, median of 3 runs. The lists come from public text: the llama.cpp sources, the Python stdlib, wikitext test, Strata's docs.

	temperature 1.0	greedy
full head, depth 2 (main)	57.8 t/s	62.0 t/s
full head, depth 3	58.9	60.2
34k list (99.5%), depth 2	61.1	66.6
43k list (99.8%), depth 3	64.5	69.9
50k list (99.9%), depth 3	64.7	70.9
55k list (99.95%), depth 3	64.7	70.9
43k list, depth 4	64.3	72.1

• Code editing: 65.2 -> 80.3 t/s at temperature 1.0, 71.1 -> 87.2 t/s greedy.
• English explanation: 52.9 -> 55.8 t/s (the corpus is mostly code).
• VRAM: the 43k list takes 173 MiB over both GPUs, the 50k list ~200 MiB.

Validation

• The first 30 draft candidates of a greedy request match with and without the list: ids and probabilities (top 3 per step, debug log). So the subset head gives the full head's logits for its tokens.
• Acceptance stays near the full head's (temperature 1.0, 43k list): 0.599 -> 0.584 write, 0.804 -> 0.792 edit, 0.530 -> 0.498 explain at depth 2. The drop is the tokens outside the list.
• All targets build.
• No unit test: there is no small MTP model here, and the random test models have a 128-token vocabulary.

Notes

• Only MTP contexts change, and only the heads that are the model's output head. A model whose MTP layers have their own head (nextn.shared_head_head) does not use the list (warning at load).
• The list is as good as its corpus: a language it lacks gets few drafts. While more than --spec-draft-vocab-max-miss (0.10) of the last 64 tokens of most drafting sequences are outside the list, the context drafts over the whole vocabulary; German with a code/English list then runs as without a list. The tool takes any text; the ids depend only on the tokenizer.
• With a list, draft sampling runs on the CPU, also during that fallback. With --spec-draft-backend-sampling (not available with -sm tensor) the fallback is about 2% slower than no list.
• The other Strata drafting idea I tried, coupled draft sampling (the draft samples with the target's random numbers), gave nothing on this model. The MTP head is much sharper than the target after top-p, so it is not part of this PR.

🤖 Generated with Claude Code

[Piggidragon]

Review fixes (5 findings) in bc24d81:

• Prism-Hadamard / precision lookups: the draft head now uses the output head's input transform and activation precision. Both are looked up with the output head. LoRA adapters of the output head are not applied to the draft head; they cover all of its rows.

• Language outside the list: a sequence whose recent text is often outside the list (a language the corpus lacks) now drafts over the whole vocabulary.

  • The drafter checks the last 64 tokens of the drafting sequences. It switches with llama_set_draft_vocab() while more than --spec-draft-vocab-max-miss (default 0.10) of them are outside the list, and switches back below half of that.

  • Measured with a code/English list at depth 3:

    	German, temp 1.0	German, greedy	code edit, temp 1.0	code edit, greedy
    no list	40.6	48.0	72.9	78.5
    list, no fallback	38.5	43.0	80.5	86.6
    list + fallback	40.6	48.0	79.9	86.4

• Separate MTP draft model: only that model gets the draft head, not the target.

• Fit (no_alloc): the dry run now counts the draft head.

• llama-draft-vocab: also adds user-defined tokens (<think>, <tool_call>, ...).

Docs: the server README rows and the tool README describe what happens to a language outside the list.

🤖 Generated with Claude Code

[Piggidragon]

Second review (5 findings, none changed the output) fixed in cda23ab, main merged in 7e6cf41:

1. Medium: the MTP context reserved its compute buffers for the draft head only. A switch to the whole vocabulary reallocated them, and GGML_SCHED_DEBUG_REALLOC=1 aborted. The reserve now uses the whole head (~1.6 MB more at --parallel 2). Check: German with a code/English list (the fallback switches) under GGML_SCHED_DEBUG_REALLOC=1 gives 0 reallocations, no abort.
2. Backend draft sampling was off for every --spec-draft-vocab. That also covered a list the model skipped. It is now off only when the model has a draft head. CPU draft sampling with a list (also during the fallback, ~2% slower than no list where backend sampling exists) is documented.
3. The switch was an OR over all drafting sequences. It now needs most of them, with per-sequence hysteresis. One conversation in a language the list lacks no longer takes the gain of the others.
4. Models whose MTP layers have their own head (nextn.shared_head_head) skip the list with a warning instead of allocating an unused copy.
5. Fit counts the draft head of a separate MTP draft model.

The README row and the PR description now show depth 2 (the production choice) with a 66k English + German + code + man-page list: +7% at temperature 1.0 and greedy, German included (46.9 -> 50.0 t/s).

The review also found a crash in main with -sm tensor --concurrent-prefill (not caused by this PR), fixed in #43.

🤖 Generated with Claude Code