feat(cuda): implement MTP self-speculative decoding for Qwen3HybridDense (#253) - #286
Merged
Conversation
…nse (#253) CUDA follow-up to the CPU MTP implementation (PR #285): adds CudaMtpState (device-resident KV-cache + pending-hidden handoff), MTP head weight loading (nextn.* tensors, mirroring the CPU loader), and ForwardMtp/Forward-with-capture on CudaQwen3HybridDenseTransformerModel, composed from existing proven CUDA kernels (RmsNorm, RoPE, attention, SwiGLU, Gemm dispatch) rather than new kernels. IMtpSpeculativeDecoder/MtpSpeculativeDecoder were already backend-agnostic and required no changes to support the CUDA model. While building a real-model CUDA correctness test (the engine-layer decoder was previously only tested against a position-independent CPU mock), found and fixed two latent bugs in the shared, backend-agnostic MtpSpeculativeDecoder.cs that the mock could not have caught: - The "catchup" forward re-processed lastToken a second time as verify-batch row 0. Harmless for the position-indexed attention KV-cache, but corrupts GDN/recurrent trunk state (a pure sequential recurrence, not position-indexed per IGdnState's own doc) -- exactly the architecture MTP targets (Qwen3.6-27B/Bonsai-27B). Fixed by reusing the catchup call's own logits as the round's position-0 comparison basis, which also removes the redundant forward the original design explicitly traded away as a documented simplification. - Both MtpSpeculativeDecoderTests and this task's copied test harness initialized `position` to 1 instead of 0, off by one against DraftAndVerify's own documented invariant ("lastToken already occupies position"). Invisible on the CPU mock (its Forward ignores position for logit computation) but corrupts RoPE/attention on a real model. Separately identified and documented (not fixed -- pre-existing, out of scope): speculative decoding's batched verify forward has no way to roll back a rejected draft token's contribution to recurrent (GDN/Mamba) trunk state once the batch has run, since IGdnState has no position-based checkpoint/restore. Affects both SpeculativeDecoder and MtpSpeculativeDecoder identically on CPU and CUDA; needs its own design pass. The new CUDA integration test isolates itself from this via an all-full-attention synthetic fixture (fullAttnInterval: 1, a new optional parameter on SyntheticQwen35HybridDenseMtpGguf) to give a clean proof of the CUDA MTP work itself. 7 new CUDA tests added (GGUF detection, hidden-capture-is-pure-side-effect, ForwardMtp determinism/finiteness, trunk-fallback, and the real-model speculative-matches-plain-greedy integration proof) -- all pass, plus the existing CPU MTP tests continue to pass with no regressions (verified: full Engine suite 648/650, non-MTP CUDA Qwen3HybridDense suite unaffected). Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
…before CUDA MTP verification
There was a problem hiding this comment.
Pull request overview
Adds CUDA support for MTP (Multi-Token Prediction) self-speculative decoding for Qwen3HybridDense, mirroring the earlier CPU implementation and tightening engine correctness around catchup/verify behavior.
Changes:
- Implemented CUDA-side MTP state and Qwen3HybridDense MTP head loading +
ForwardMtpexecution path. - Fixed engine-layer MTP speculative verification to avoid redundant
lastTokenre-forwarding by reusing catchup logits for position-0 comparison. - Expanded/adjusted synthetic GGUF fixture + tests (including new CUDA integration coverage) and corrected an off-by-one position convention in the existing engine unit test.
Reviewed changes
Copilot reviewed 6 out of 6 changed files in this pull request and generated 3 comments.
Show a summary per file
| File | Description |
|---|---|
| tests/DotLLM.Tests.Unit/Engine/MtpSpeculativeDecoderTests.cs | Fixes test harness position initialization to align with the decoder/KV-cache contract. |
| tests/DotLLM.Tests.Unit/Cuda/CudaQwen3HybridDenseMtpTests.cs | Adds CUDA tests covering GGUF detection, capture purity, MTP forward correctness/determinism, and end-to-end decoder equivalence. |
| src/DotLLM.Models/Gguf/SyntheticQwen35HybridDenseMtpGguf.cs | Extends synthetic fixture generator to allow an all-full-attention trunk variant for isolation testing. |
| src/DotLLM.Engine/MtpSpeculativeDecoder.cs | Reuses catchup logits for position-0 verification and removes redundant lastToken verify row. |
| src/DotLLM.Cuda/Architectures/CudaQwen3HybridDenseTransformerModel.cs | Loads MTP head tensors, exposes SupportsMtp/CreateMtpState, captures pre-final-norm hidden state, and implements ForwardMtp. |
| src/DotLLM.Cuda/Architectures/CudaMtpState.cs | Introduces device-resident MTP KV-cache + pending-hidden handoff, with host-captured hidden-row storage. |
Suppressed comments (1)
src/DotLLM.Cuda/Architectures/CudaQwen3HybridDenseTransformerModel.cs:1172
- ForwardMtpCore allocates a new managed int[] (posHost) on every draft step just to copy a single position scalar to the device. This is avoidable per-step GC churn inside the K-step MTP loop.
// RoPE — partial-rotary NeoX, at this step's absolute round-relative position.
int[] posHost = [position];
fixed (int* pPos = posHost)
{
CudaDriverApi.cuMemcpyHtoDAsync_v2(s.PositionDevice, (nint)pPos, sizeof(int), streamH).ThrowOnError();
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Comment on lines
75
to
77
| { | ||
| var w = new GgufWriter(); | ||
| var rng = new SyntheticGemma4Gguf.Xorshift(seed); |
Comment on lines
+1104
to
+1116
| nint embedHostBase = mtpHead.EmbedTokensHostBase ?? _embedDataBase; | ||
| ulong embedDataOffset = mtpHead.EmbedTokensHostBase is not null ? mtpHead.EmbedTokensDataOffset : _embedDataOffset; | ||
| long embedRowBytes = mtpHead.EmbedTokensHostBase is not null ? mtpHead.EmbedTokensRowBytes : _embedRowBytes; | ||
| QuantizationType embedQt = mtpHead.EmbedTokensHostBase is not null ? mtpHead.EmbedTokensQt : _tokenEmbedQt; | ||
|
|
||
| float[] embedHost = new float[hiddenSize]; | ||
| nint rowSrc = embedHostBase + (nint)(embedDataOffset + (ulong)tokenId * (ulong)embedRowBytes); | ||
| Dequantize.ToFloat32(rowSrc, hiddenSize, embedQt, embedHost); | ||
| fixed (float* pEmbedHost = embedHost) | ||
| { | ||
| CudaDriverApi.cuMemcpyHtoDAsync_v2(s.Embed, (nint)pEmbedHost, | ||
| (nuint)((long)hiddenSize * sizeof(float)), streamH).ThrowOnError(); | ||
| } |
Comment on lines
+215
to
+225
| public IMtpState? CreateMtpState() | ||
| { | ||
| if (_mtpHead is null) | ||
| return null; | ||
|
|
||
| return new CudaMtpState( | ||
| hiddenSize: Config.HiddenSize, | ||
| numKvHeads: _mtpHead.Value.Layer.FullAttn!.Value.NumKvHeads, | ||
| headDim: Config.HeadDim, | ||
| maxSteps: MtpDefaultMaxDraftSteps); | ||
| } |
This was referenced Aug 8, 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.
Implements the CUDA side of MTP self-speculative decoding, following the CPU implementation merged earlier today (PR #285). Requested to unblock an honest CPU-vs-CUDA-vs-llama.cpp benchmark comparison.
Confirmed genuinely new work
No prior CUDA MTP work existed anywhere (branches, PRs, issue references) — #253 explicitly scoped CUDA as CPU-first follow-up.
What's implemented
CudaMtpState : IMtpState— device-resident GPU K/V cache for the MTP block's own attention, plus device-resident pending-hidden-state handoff (no host round-trip mid-round).CudaQwen3HybridDenseTransformerModel: MTP head weight loading (nextn.*tensors, mirroring the CPU loader),SupportsMtp/CreateMtpState/ForwardMtp, and aForward(...,IMtpState?)override capturing pre-final-norm hidden state as a pure side effect.ForwardMtpis composed entirely from existing proven CUDA kernels (RmsNorm, RoPE, attention, SwiGLU, theGemmdispatcher) — no new kernels needed.IMtpSpeculativeDecoder/MtpSpeculativeDecoderwere already backend-agnostic — zero required engine changes to make CUDA work.Two real bugs found and fixed in the shared engine layer
lastTokena second time as verify-batch row 0 — harmless for position-indexed attention KV-cache, but corrupts GDN recurrent state (IGdnStateis a pure sequential recurrence, no position addressing). Fixed by reusing the catchup call's own logits as the round's position-0 comparison basis, removing a redundant forward per round as a side benefit.positioninitialized to 1 instead of 0) — invisible to the CPU mock test (ignores position) but corrupts RoPE on a real model. Fixed in both the existing CPU test and the new CUDA test.Documented, not fixed (pre-existing, out of scope for this PR): speculative decoding's batched verify has no rollback for rejected tokens' effect on recurrent trunk state. Affects both
SpeculativeDecoderandMtpSpeculativeDecoder, CPU and CUDA identically — recommend filing separately.Tested
7 new CUDA tests (GGUF detection, hidden-capture purity,
ForwardMtpdeterminism/finiteness, trunk-fallback, full self-speculative-vs-plain-greedy integration proof against the real CUDA model) — all passing, independently re-verified on a clean rebuild. Full Engine suite (650 tests) independently re-run: 648 passed, 2 pre-existing unrelated skips, zero regressions.Refs #253 (not closing — real end-to-end 27B validation and CLI/server wiring are separate, in-progress work).