Research: FR-Spec-style draft-vocab trimming for native MTP speculative decoding

[avifenesh]

Research Stage

• Background Research
• Hypothesis Formed
• Strategy / Implementation Forming
• Analysis of results
• Debrief / Documentation

Previous existing literature and research

The bellow is a result of trying to reduce my decode bottleneck and profiling it, noticing that a large piece of it is coming from the draft itself.

For large-vocabulary models, the draft model's LM-head projection (hidden -> vocab_size) is a significant fraction of per-step draft cost in speculative decoding, independent of the draft transformer's own size. FR-Spec (Zhao, Pan, Han et al., ACL 2025, arXiv:2502.14856) addresses this directly: restrict the draft's candidate search to a frequency-ranked subset of the vocabulary (e.g. top 32,768 of a 100k+ vocab), cutting LM-head compute by ~75% while the target model still verifies over the full vocabulary, so the output distribution is provably unchanged (lossless, same guarantee as standard speculative sampling).

SGLang already ships this for EAGLE-style drafters via --speculative-token-map: a hot_token_id tensor gathers the draft's lm_head rows down to the hot set at load time, and after the draft's topk sampling, predicted indices are remapped back to real vocab IDs (topk_index = hot_token_id[topk_index]) before being used downstream (eagle_worker_v2.py, lines ~336-338 and ~594-595 as of this writing).

llama.cpp's own EAGLE-3 implementation (src/models/eagle3.cpp) already has the equivalent machinery for a different reason — handling a draft model whose native vocabulary is smaller than the target's. It carries a d2t ("draft-to-target") tensor, auto-sized from the GGUF (d2t_meta->ne[0]), and after the draft head's matmul it scatters the compressed logits into a full-vocab tensor (filled -inf elsewhere) via ggml_set_rows, so every downstream consumer (sampler, verify path) only ever sees full-vocab-shaped, correctly-indexed logits.

The native MTP path (--spec-type draft-mtp, landed in #22673) has no equivalent — every MTP-capable architecture I checked (qwen35.cpp, qwen35moe.cpp, step35.cpp, glm4-moe.cpp, bailingmoe2.cpp) hardcodes the MTP head's output tensor (nextn.shared_head_head, or model.output when that's absent) to the trunk's full n_vocab.

Hypothesis

d2t is architecture-agnostic infrastructure, not EAGLE-3-specific. If a standalone MTP-only draft GGUF ships a d2t tensor, the same trim-at-load + scatter-after-matmul pattern EAGLE-3 already uses should let native MTP drafters use a frequency-ranked vocab subset too — with the same lossless guarantee (target still verifies full vocab) and the same kind of LM-head-cost reduction FR-Spec reports.

Strategy / Implementation

Implemented and tested for qwen35.cpp (covers Qwen3.5/Qwen3.6 dense MTP). Branch: https://github.com/avifenesh/llama.cpp/tree/frspec-mtp-vocab-trim (commit 047bfa508).

The change is ~30 lines, gated so it's a no-op for every existing GGUF:

• In load_arch_tensors: if the draft is MTP-only (mtp_only) and a d2t tensor is present in the GGUF, size output.weight from d2t->ne[0] instead of n_vocab, and load d2t via the same LLM_TENSOR_D2T tensor name EAGLE-3 already uses.
• In the MTP head's graph build (after build_lora_mm(head_w, ...)): if model.d2t is set, scatter the compressed logits back to full-vocab shape with ggml_set_rows, identical to the existing code in eagle3.cpp.
• No changes needed in common/speculative.cpp — the scatter happens inside the compute graph, so the draft sampling loop never has to be aware the head was trimmed (verified empirically: the EAGLE-3 draft loop already has zero d2t-aware code, for the same reason).

A draft GGUF with a trimmed output.weight + a d2t tensor was produced offline with gguf-py (gather the output.weight rows by a frequency-ranked token-ID list, store that list as d2t). No convert_hf_to_gguf.py producer path exists yet — this issue is about the consumer/loader side and whether the approach is worth pursuing before building that out.

The frequency map itself was built from two sources for reproducibility:

1. A 1GB, 50/50 code-prose corpus from public HF datasets (codeparrot/codeparrot-clean-valid, codeparrot/github-jupyter-code-to-text, Salesforce/wikitext wikitext-103, allenai/c4 en, HuggingFaceFW/fineweb-edu), tokenized with the target model's own tokenizer, top-32768 by raw frequency.
2. A code-only corpus from my own agent/coding-session logs (not reproducible by others, included only as a private-workload data point, not part of the proposal).

Analysis

Hardware: Qwen3.6-27B-NVFP4 trunk + Q6_K MTP draft GGUF (--spec-type draft-mtp --spec-draft-n-max 3), single RTX 5090 Laptop, 128k ctx, K=q8_0/V=q5_1 cache.

Kernel-level (Nsight Systems, isolated decode window, same prompt before/after):

	baseline (full 248,320-vocab head)	trimmed (32,768 hot vocab)	delta
mul_mat_vec_q (Q6_K, draft LM-head, batch=1)	407.1 ms / 655 calls	61.6 ms / 675 calls	-84.9%
avg latency per call	621 µs	91 µs	-85.3%

That matches FR-Spec's reported ~75% LM-head-compute reduction (32768/248320 ≈ 7.6x fewer rows; measured ≈6.6x time reduction — close given other fixed per-call overhead).

End-to-end decode throughput, wide bench (6 code + 4 prose prompts, temp=0, n-max=3):

draft / vocab map	code mean tok/s	prose mean tok/s	overall mean
baseline (no trim)	88.4	77.2	83.9
trimmed, public 50/50 corpus	89.6	78.3	85.1
trimmed, private code-heavy corpus	93.4	76.1	86.5

The public, reproducible 50/50 map improves both code and prose over baseline. A code-weighted map trades some prose acceptance for a bigger code win — expected, since the trim only ever costs speed (a miss falls through to a lower-confidence/rejected draft token), never correctness.

Correctness: at temperature=0, generated text is byte-for-byte identical between the baseline and trimmed drafts across every prompt tested. This is the expected result, not a surprise — the target always verifies over the full vocabulary regardless of what the draft proposed, so the lossless guarantee holds independent of the draft's hit rate. I'm noting it because I checked it, not because it was in doubt.

Architectural generality (claimed, not yet run): step35.cpp, glm4-moe.cpp, and bailingmoe2.cpp declare the identical nextn.shared_head_head / {n_embd, n_vocab} tensor pattern this patch targets in qwen35.cpp, so the same ~30-line change should apply to those architectures unchanged. I have not verified this on real hardware: I could not find a current, non-Qwen model that is (a) small enough to run on a single 24GB GPU, (b) ships an MTP/nextn head in its public GGUF, and (c) is loadable by llama.cpp's existing MTP path — every candidate I checked failed at least one of those three (e.g. a popular GLM-4.7-Flash GGUF reports num_nextn_predict_layers: 1 in its source config.json, but converts to llama.cpp's deepseek2 architecture, whose loader has no nextn/MTP code at all, and the actual GGUF bytes confirm no MTP tensors were written). I'm flagging this gap rather than papering over it.

Before building a convert_hf_to_gguf.py producer path and turning this into a PR, I'd like to check this is a direction the project wants, and get input on a few open questions: reusing the existing d2t tensor/convention for this vs. an MTP-specific key (is there a reason EAGLE-3 and MTP shouldn't share it?); where the frequency map should come from in a real workflow (baked into the converter from a bundled/standard corpus vs. a separate user-supplied map file, closer to SGLang's --speculative-token-map); and whether the missing cross-architecture validation above is a blocker before this is worth a PR, or fine to defer to whoever has hardware for one of the matching architectures. Happy to do the convert_hf_to_gguf.py work and write it up properly if this is a direction worth pursuing.

Relevant log output

mul_mat_vec_q<(ggml_type)14, (int)1, (bool)0, (bool)0>  [baseline]
  Total Time: 407,053,180 ns   Instances: 655   Avg: 621,455.2 ns

mul_mat_vec_q<(ggml_type)12, (int)1, (bool)1, (bool)0>  [trimmed, same kernel slot]
  Total Time: 61,598,225 ns    Instances: 675   Avg: 91,256.6 ns

draft acceptance (temp=0, deterministic, baseline vs trimmed, identical prompt):
  baseline: draft 215/251 accepted, mean len 3.56
  trimmed:  draft 215/251 accepted, mean len 3.56   <- identical, confirms determinism preserved

Disclosure: The research and the writing were done with help from an LLM, I noticed the requirement not to use AI assistants heavily, and in case that's unacceptable, I will close the issue with understanding.
However, LLM usage done under provision was steered and direct, reviewed and approved by me and I'm taking ownership and responsability over it.

[PierpaoloPernici]

A small independent check of the FR-Spec MTP vocab trim (Qwen3.8-27B)

First of all, thanks for the detailed write-up and for the clean d2t implementation. The cherry-pick was painless to build and worked on the first try.

I was curious to see how the shipped FR-Spec draft would behave on a somewhat different workload, so we ran a few tests on our own hardware, mostly using generic prompts and technical/prose output rather than agentic sessions.

Setup

Everything was tested with the custom FR-Spec build (10752 / 044b56f77), loading the draft cross-file with -md + --spec-type draft-mtp. The draft came from a different quant of the same base model.

Parameter	Value
Target	esatapedico/Qwen3.8-27B-NVFP4-MTP-LOW (Q5_0 lm_head, IQ4_XS internal MTP head)
Bench	9 generic prompts (code/prose/QA) + a real 70k-token technical document
--spec-draft-n-max / ctx / KV	4 / 192K / q8_0
Sampling	temp 1.0, seed 42 (same seed everywhere)
Lossless check	outputs byte-identical to the target's internal MTP at the same seed

Results

Draft on LOW	Long-ctx acc	Long-ctx tok/s	Short wall (acc)
LOW internal MTP (IQ4_XS, full 248k head)	68.1%	109.7	11.42–11.52s (49.4%)
frspec sxc (agentic trim 32k, NVFP4)	60.8%	105.6	10.76–11.31s (46.1%)
frspec prose (prose trim 32k, NVFP4)	60.6%	104.5	11.74s (44.3%)
cross-file full head (Q5_K, 248,320 rows)	70.9%	112.6	11.39–11.81s (49.4%)
cross-file full head (NVFP4 requant)	70.9%	109.6	11.81s (49.4%)
cross-file full head (BF16, 5.66 GiB)	70.9%	110.3	12.17s (49.4%)

A few things I found interesting:

1. The d2t machinery works really well. Cross-file loading from a different quant of the same base model worked without any surprises, and the runs were lossless: same output tokens as the target's internal MTP with the same seed. For me, this alone makes the d2t path very useful.

2. The 32k trim didn't work particularly well for our workload. On the long technical test it lost roughly 10 percentage points of acceptance compared with the full head, and around 5pp on the shorter prompts. That's quite a bit more than the ~3pp reported at K=3 on the agentic workload. It did give us a small wall-time win on short interactive prompts (~5%), but that was basically the only case where it came out ahead.

3. Changing the ranking source didn't change much. I rebuilt the trim using the prose-ranked vocabulary just to see what would happen, but the result was almost the same: 60.6% vs 60.8%. So at least on these tests, it looks more like a consequence of reducing the vocabulary to ~13% than a problem with the particular ranking used for the original trim.

4. The full head gave another interesting result. Acceptance topped out at 70.9%: IQ4_XS was lower, while Q5_K, NVFP4 and BF16 all landed at basically the same point. BF16 was actually a little slower, which isn't too surprising given that it means dragging a 5.66 GiB head through every draft pass. Q5_K ended up being the nicest compromise for us at about 1.9 GiB.

One thing I wasn't expecting: the cross-file Q5_K full head actually did slightly better than the target's own internal IQ4_XS head on the long-context run: 70.9% / 112.6 tok/s vs 68.1% / 109.7 tok/s. On the short test they were basically tied. So just improving the head quantization seems to be worth roughly 3pp of acceptance here.

Full write-up, numbers and scripts: gist

And, in the spirit of the experiment: I used Qwen3.8-27B itself to help me put together some of the test scripts. :)

[avifenesh]

A small independent check of the FR-Spec MTP vocab trim (Qwen3.8-27B)

First of all, thanks for the detailed write-up and for the clean d2t implementation. The cherry-pick was painless to build and worked on the first try.

I was curious to see how the shipped FR-Spec draft would behave on a somewhat different workload, so we ran a few tests on our own hardware, mostly using generic prompts and technical/prose output rather than agentic sessions.

Setup

Everything was tested with the custom FR-Spec build (10752 / 044b56f77), loading the draft cross-file with -md + --spec-type draft-mtp. The draft came from a different quant of the same base model.

Parameter Value
Target esatapedico/Qwen3.8-27B-NVFP4-MTP-LOW (Q5_0 lm_head, IQ4_XS internal MTP head)
Bench 9 generic prompts (code/prose/QA) + a real 70k-token technical document
--spec-draft-n-max / ctx / KV 4 / 192K / q8_0
Sampling temp 1.0, seed 42 (same seed everywhere)
Lossless check outputs byte-identical to the target's internal MTP at the same seed

Results

Draft on LOW Long-ctx acc Long-ctx tok/s Short wall (acc)
LOW internal MTP (IQ4_XS, full 248k head) 68.1% 109.7 11.42–11.52s (49.4%)
frspec sxc (agentic trim 32k, NVFP4) 60.8% 105.6 10.76–11.31s (46.1%)
frspec prose (prose trim 32k, NVFP4) 60.6% 104.5 11.74s (44.3%)
cross-file full head (Q5_K, 248,320 rows) 70.9% 112.6 11.39–11.81s (49.4%)
cross-file full head (NVFP4 requant) 70.9% 109.6 11.81s (49.4%)
cross-file full head (BF16, 5.66 GiB) 70.9% 110.3 12.17s (49.4%)
A few things I found interesting:

1. The d2t machinery works really well. Cross-file loading from a different quant of the same base model worked without any surprises, and the runs were lossless: same output tokens as the target's internal MTP with the same seed. For me, this alone makes the d2t path very useful.
2. The 32k trim didn't work particularly well for our workload. On the long technical test it lost roughly 10 percentage points of acceptance compared with the full head, and around 5pp on the shorter prompts. That's quite a bit more than the ~3pp reported at K=3 on the agentic workload. It did give us a small wall-time win on short interactive prompts (~5%), but that was basically the only case where it came out ahead.
3. Changing the ranking source didn't change much. I rebuilt the trim using the prose-ranked vocabulary just to see what would happen, but the result was almost the same: 60.6% vs 60.8%. So at least on these tests, it looks more like a consequence of reducing the vocabulary to ~13% than a problem with the particular ranking used for the original trim.
4. The full head gave another interesting result. Acceptance topped out at 70.9%: IQ4_XS was lower, while Q5_K, NVFP4 and BF16 all landed at basically the same point. BF16 was actually a little slower, which isn't too surprising given that it means dragging a 5.66 GiB head through every draft pass. Q5_K ended up being the nicest compromise for us at about 1.9 GiB.

One thing I wasn't expecting: the cross-file Q5_K full head actually did slightly better than the target's own internal IQ4_XS head on the long-context run: 70.9% / 112.6 tok/s vs 68.1% / 109.7 tok/s. On the short test they were basically tied. So just improving the head quantization seems to be worth roughly 3pp of acceptance here.

Full write-up, numbers and scripts: gist

And, in the spirit of the experiment: I used Qwen3.8-27B itself to help me put together some of the test scripts. :)

@PierpaoloPernici
https://huggingface.co/Avifenesh/Qwen3.8-27B-NVFP4-MTP-GGUF
The gain is not acceptance, it is the overall tok/s
Qwen3.8-27B-NVFP4-Q5K-mtp.gguf is the one used
Your machine matters a lot, if you are decode bound, that gives you a strong benefit, but if memory, that's not going to help a lot.

[PierpaoloPernici]

Totally — I get your point now: the relevant win is overall tok/s, not acceptance rate. The decode-bound vs. memory-bound distinction also explains our results: the trim gave us a small win on the short, decode-heavy prompts, but it did not pay off on the 70K-token workload, where the larger context puts much more pressure on memory and prefill. Thanks again for clarifying this

[drluoto]

Independent check on a different architecture and much slower memory: Qwen3.8-Flash-Next (qwen4exp) on Strix Halo (Ryzen AI Max+ 395, 128 GB, Vulkan/RADV), where the draft LM head is a much bigger share of the step than on a 5090.

Per drafted token the Q8_0 MTP sidecar reads ~829 MB, of which output.weight (248,320 × 2560) is 675 MB, i.e. 81 %. At ~230 GB/s that is ~6 ms per draft token, so with n-max 3 the head alone costs more than half of the speculation overhead.

Setup: trunk requantized to Q5_K (+Q8_0 router), MTP sidecar Q8_0 with output.weight trimmed to 65,536 frequency-ranked rows (agent traces weighted ×5 plus a small code/prose corpus), d2t as I64. Consumer side ported from the qwen35 patch above. Two things worth knowing for qwen4exp:

• the scatter has to go in graph_mtp's head (the logits = build_lora_mm(model.output, final, …) around line 500), not the trunk graph's head further down; in the wrong place it silently leaves logits at compressed indices, which looks like "works only for identity-like maps".
• 32k rows was too small for this workload (Swedish + code): acceptance fell from 0.87 to 0.72 on short prompts. 65k kept it.

Results, production-like server (3 slots, 262k ctx, vision, MTP n-max 3, p-min 0.7), 20 real agent traces, same binary for both rows:

sidecar	warm decode	cold	acceptance
original 248,320 rows	27.4 tok/s	29.3	0.78
trimmed 65,536 rows	29.9 tok/s	28.9	0.81

+9 % warm, 8k-context single prompt 33.9 → 36.8 tok/s, output unchanged (target verifies the full vocab), 0 corrupt answers, our quality gate passed. Sidecar shrinks 4.14 → 3.64 GB.

Branch with the qwen4exp port, the trimming script and the frequency map: https://github.com/drluoto/llama.cpp/tree/frspec-qwen4exp-strix (scripts/frspec/). The ready-made 65k sidecar is a 16-second run of that script on the existing Q8_0 head from https://huggingface.co/drluoto/Qwen3.8-Flash-Next-MTP-GGUF; I will add the file there as well. Done together with Claude Fable 5.1 as part of a longer optimization pass on this box; happy to test a convert-side producer path if that is the direction.