feat(engine): MTP self-speculative decoding — CPU implementation (#253) - #285
Merged
Conversation
…nk NumLayers (#253) Multi-Token Prediction (MTP) self-speculative decoding research (llama.cpp PR ggml-org/llama.cpp#22673): a checkpoint's MTP head is appended as extra trailing GGUF block(s) beyond the trunk (block_count = num_hidden_layers + mtp_num_hidden_layers). Add ModelConfig.NextnPredictLayers and extract it for Qwen3HybridDense/Qwen3MoeHybrid, subtracting it back out of NumLayers so every existing NumLayers consumer (KV-cache sizing, hybrid layout, per-layer arrays) keeps seeing exactly the trunk stack it already expects. Defaults to 0 for every checkpoint without the key — zero behavior change.
Design decision (see docs/SPECULATIVE.md): MTP self-speculative decoding needs a parallel interface, not an ISpeculativeDecoder overload, because there is no second IModel — the "draft" is the target model's own extra head sharing its weights file, seeded from the target's own hidden state rather than a second model's independent forward pass. IMtpState carries the MTP head's own tiny KV-cache (sized for just the trailing MTP block, not the trunk) plus the pending-hidden-state handoff row and the captured-rows buffer a verify-phase Forward call populates — mirrors llama.cpp's common_speculative_state_draft_mtp translated into dotLLM's per-sequence-state idiom (see IRecurrentSequenceState). IModel gains SupportsMtp (default false), CreateMtpState() (default null), a Forward(..., IMtpState?) overload that captures pre-final-norm hidden state as a pure side effect (default ignores it), and ForwardMtp (default throws). All zero-cost/zero-behavior-change for every model that doesn't override them.
CPU-first implementation of the MTP head confirmed against llama.cpp PR ggml-org/llama.cpp#22673's src/models/qwen35.cpp (load_block_mtp / graph_mtp): - MtpHeadWeights: the MTP block's own decoder-layer weights (structurally a full-attention Qwen3HybridDense layer) plus the four nextn.* tensors (eh_proj/enorm/hnorm/shared_head_*), with documented fallback to the trunk's own token_embd/output/output_norm when the optional head-local tensors are absent. - CpuMtpState: concrete IMtpState — unmanaged, 64-byte-aligned KV-cache sized for just the MTP block, pending-hidden handoff, captured-rows buffer. - Qwen3HybridDenseTransformerModel: loads the MTP head when config.NextnPredictLayers > 0 and the nextn.* tensors are present (null otherwise — zero behavior change); implements SupportsMtp/CreateMtpState/ ForwardMtp; the trunk Forward gains a hidden-state capture point (a pure side effect, proven byte-identical by test) right before the final norm overwrites the buffer in place. - SyntheticQwen35HybridDenseMtpGguf: tiny deterministic qwen35 GGUF fixture (with or without an MTP head) built from the confirmed tensor naming/layout — no real Qwen3.6-MTP-GGUF fixture is cached locally (see the issue's fixture-availability note), so this is what the tests below exercise the loader and forward math against.
IMtpSpeculativeDecoder.DraftAndVerify(targetModel, kvCacheTarget, mtpState, ...) — the self-speculative analogue of ISpeculativeDecoder.DraftAndVerify, sharing the same SpeculativeResult shape and greedy-only correctness gate (matches llama.cpp's own merged MTP draft sampler, which is top_k=1 today too — see the source citation in docs/SPECULATIVE.md). MtpSpeculativeDecoder's draft-verify-accept loop starts every round with a single-token "catchup" forward of lastToken (capturing its hidden state into mtpState) before drafting: neither a corrected token nor a bonus token from the previous round has ever been forwarded through the trunk as an input, so its hidden state does not exist yet and must be (re)computed before the MTP head's first draft step can be seeded correctly. This is a real correctness subtlety found and worked through this session — see the type's remarks and docs/SPECULATIVE.md for the full derivation. The MTP head's own KV-cache is reset every round (a documented simplification vs. llama.cpp's cross-round-persistent ctx_dft) since each round reseeds entirely from the target's own just-verified hidden state anyway.
…de (#253) Qwen3HybridDenseMtpTests (against SyntheticQwen35HybridDenseMtpGguf): MTP hparam detection + trunk-layer subtraction, zero-behavior-change for non-MTP checkpoints, hidden-state capture is a byte-identical-logits pure side effect, MTP head forward produces finite logits and advances its own KV-cache, determinism across fresh runs, and the trunk-fallback path when optional head-local nextn.* tensors are absent. MtpSpeculativeDecoderTests (synthetic MTP-capable mock model, isolating decoder mechanics from real forward math): constructor greedy-only gate, SupportsMtp guard, K=0 edge case, and the two tests demonstrating the issue's correctness bar — MTP self-speculative decoding produces the exact same output token sequence as plain greedy decode of the target alone, both when the MTP head always agrees with the target and when it is deliberately wrong at half the tokens (forcing rejections every other round). Full existing Models+Engine suite (1145 tests) still green.
Real citations (llama.cpp PR ggml-org/llama.cpp#22673 source, not a write-up), the parallel-interface design decision and rationale, the catchup-forward and per-round-KV-reset simplifications with their correctness arguments, what's tested and how, and what's explicitly left (CUDA, TextGenerator/CLI/server wiring, real-fixture end-to-end validation).
There was a problem hiding this comment.
Pull request overview
Adds CPU-side support for Multi-Token Prediction (MTP) self-speculative decoding (issue #253) by extending GGUF config extraction/loading to recognize the nextn_* head, adding model/state APIs (IMtpState, IModel.SupportsMtp/CreateMtpState/ForwardMtp), and introducing an engine-level MtpSpeculativeDecoder plus synthetic fixtures and unit tests.
Changes:
- Extend GGUF config extraction to detect
nextn_predict_layersand treatblock_countas trunk+MTP (trunkNumLayersis adjusted accordingly). - Implement CPU MTP head loading + forward path for
Qwen3HybridDenseTransformerModel, including hidden-state capture as a side effect of trunkForward. - Add
IMtpSpeculativeDecoder/MtpSpeculativeDecoderand synthetic GGUF fixtures + unit tests for loader/forward and decoder mechanics.
Reviewed changes
Copilot reviewed 13 out of 13 changed files in this pull request and generated 2 comments.
Show a summary per file
| File | Description |
|---|---|
| tests/DotLLM.Tests.Unit/Models/Architectures/Qwen3HybridDenseMtpTests.cs | Exercises MTP detection, trunk-layer subtraction, hidden-state capture side effect, and MTP head forward on a synthetic GGUF. |
| tests/DotLLM.Tests.Unit/Engine/MtpSpeculativeDecoderTests.cs | Validates draft/verify/accept semantics for the self-speculative decoder with a deterministic mock model. |
| src/DotLLM.Models/Gguf/SyntheticQwen35HybridDenseMtpGguf.cs | Builds a tiny deterministic Qwen3HybridDense GGUF fixture with optional trailing MTP (“NextN”) tensors. |
| src/DotLLM.Models/Gguf/GgufModelConfigExtractor.cs | Parses nextn_predict_layers and uses trunk-layer count for hybrid layout derivation. |
| src/DotLLM.Models/Architectures/Qwen3HybridDenseTransformerModel.cs | Loads optional MTP head weights, captures pre-final-norm hidden state, exposes SupportsMtp/CreateMtpState/ForwardMtp, and implements MTP head forward. |
| src/DotLLM.Models/Architectures/MtpHeadWeights.cs | Defines the loaded tensor bundle for the trailing MTP head. |
| src/DotLLM.Models/Architectures/CpuMtpState.cs | Implements the per-sequence CPU MTP KV-cache + pending hidden + captured hidden-row buffer. |
| src/DotLLM.Engine/MtpSpeculativeDecoder.cs | Implements MTP self-speculative draft/verify/accept loop (greedy-only) using ForwardMtp + one verify forward. |
| src/DotLLM.Engine/IMtpSpeculativeDecoder.cs | Introduces a parallel speculative-decoder interface tailored to MTP’s single-model state shape. |
| src/DotLLM.Core/Models/ModelConfig.cs | Adds NextnPredictLayers to represent trailing MTP head layer count. |
| src/DotLLM.Core/Models/IMtpState.cs | Defines the per-sequence state contract for MTP decoding. |
| src/DotLLM.Core/Models/IModel.cs | Adds MTP capability members with safe defaults (off/no-op/throw). |
| docs/SPECULATIVE.md | Documents MTP research, design choices, and current limitations/scope. |
Suppressed comments (1)
src/DotLLM.Engine/MtpSpeculativeDecoder.cs:237
- When all K draft tokens are accepted, the bonus token is written to
outputBufferbut the decoding constraint is not advanced for it. Since the caller is explicitly told not to advance constraints outsideDraftAndVerify, this leavesconstraintout of sync with the emitted token stream.
TokenMaskApplier.Apply(bonusLogitSpan, constraint.GetAllowedTokens());
int bonusToken = TensorPrimitives.IndexOfMax((ReadOnlySpan<float>)bonusLogitSpan);
outputBuffer[acceptedCount++] = bonusToken;
}
Comment on lines
+105
to
+109
| // Clamp K to remaining target KV-cache capacity and the MTP head's own KV-cache depth. | ||
| int maxPos = kvCacheTarget.MaxLength; | ||
| int k = Math.Min(numCandidates, maxPos - position - 1); | ||
| if (k <= 0) | ||
| return default; |
Comment on lines
+42
to
+45
| int nextnPredictLayers = architecture is Architecture.Qwen3MoeHybrid or Architecture.Qwen3HybridDense | ||
| ? (int)metadata.GetUInt32OrDefault($"{arch}.nextn_predict_layers", 0) | ||
| : 0; | ||
| int numTrunkLayers = numLayers - nextnPredictLayers; |
This was referenced Aug 7, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Adds Multi-Token Prediction (MTP) self-speculative decoding: a lightweight extra prediction head shipped in the same GGUF as the main model, verified in one extra forward pass, as opposed to dotLLM's existing two-model
ISpeculativeDecoder.Research
Confirmed the real tensor layout and decode-loop design directly from llama.cpp's actual merge (ggml-org/llama.cpp#22673), not secondary write-ups: MTP reuses the pre-existing DeepSeek-V3 "NextN" tensor group (
blk.{n}.nextn.*), driven by hparamnextn_predict_layers. The MTP block is a full-attention decoder layer appended after the trunk. llama.cpp's own merged-draft sampler is greedy-only (top_k=1), matching this project's existing speculative-decoder gate.What's implemented
nextn_predict_layers, subtracts from trunkNumLayers. Zero behavior change when absent.IMtpState,IModel.SupportsMtp/CreateMtpState()/ForwardMtp.Qwen3HybridDenseTransformerModel(eh_proj → gated attention → SwiGLU FFN → shared head), reusing existing trunk kernels.IMtpSpeculativeDecoder/MtpSpeculativeDecoder— a parallel interface toISpeculativeDecoder, not an overload, since MTP has no second model/KV-cache to plug into that signature. Rationale documented indocs/SPECULATIVE.md.Tested
11 new tests (GGUF loading, forward math, self-speculative decode) — including
DraftAndVerify_WithDisagreements_StillMatchesPlainGreedyDecode, which demonstrates the correctness bar: byte-identical output to plain greedy decode even when the MTP head is deliberately wrong at half the tokens. Full existing Models+Engine suite (1145 tests) still green, zero regressions — independently re-verified on a clean rebuild before this PR.Explicitly out of scope / left for follow-up
--mtpflag etc.).nextn_predict_layers > 1).Qwen3.6-27B-MTP-GGUF— no such fixture exists locally (checked~/.dotllm/test-cache/and the HF cache). All tests use a synthetic GGUF built from the confirmed-real tensor layout. Flagging per the issue's own instruction rather than downloading a 27B model unprompted — happy to fetch it in a follow-up once confirmed.Refs #253 (not closing — CUDA + real-model validation remain).