MTP as Speculative Decoding at Inference

Sections 7.1 through 7.4 made the case for Multi-Token Prediction as a training objective: it pulls more supervision out of every forward pass, sharpens the model's long-range planning, and costs almost nothing in extra parameters. But there is a second life for the same machinery that DeepSeek-V3 quietly cashes in at deployment time. The K small MTP heads that were trained to predict $2, 3, \dots, K$ tokens ahead already know how to draft K speculative tokens for any prefix. Hand those drafts to the main model for one verification forward pass and you get a wall-clock speedup that costs nothing to deploy. That is MTP-driven speculative decoding — and it is why a 671B-parameter mixture-of-experts model can serve tokens at a price that makes commercial sense.

The promise of MTP at inference. Reuse the training-time draft heads as a free, in-distribution draft model. One main-model forward pass per round, K draft tokens proposed in parallel, between 1 and $K+1$ tokens committed. The arithmetic is striking: with K=3 and a 75% acceptance rate, the same hardware emits roughly $3\times$ more tokens per second than the standard autoregressive loop — at identical sampling distribution.

The Inference Bottleneck

Standard autoregressive decoding does one forward pass per generated token. For a 67B-parameter dense model that is hundreds of gigabytes of parameter weights and KV-cache state pulled through GPU memory per token. The bottleneck on a modern H100 is not compute — the tensor cores are idling. The bottleneck is the memory bandwidth needed to stream the model weights and the cache into the SMs for each forward.

Concretely, a single H100 has roughly 3 TB/s of HBM bandwidth, and a 67B FP16 model is about 134 GB. If every token required a full model read, the ceiling is around 22 tokens per second per GPU — and that ceiling does not budge no matter how fast your tensor cores are. Compute utilisation at inference often sits below 5%. The expensive thing on the chip is doing nothing for 95% of every second.

Intuition: Two Models, One Forward Pass

The mental image is a typist and a proof-reader. The typist (the MTP heads) is quick and not always right — they bash out the next few words at speed. The proof-reader (the main model) is slow and authoritative — they read the whole proposal and stop at the first word that they themselves would not have written. Everything before that line is kept. The proof-reader then writes the corrected word and hands the draft back.

Crucially, the proof-reader does not re-read the page once for each word. They read the entire proposed continuation in a single sweep — and because the proof-reader is a causal transformer, that sweep already gives them their own prediction at every position simultaneously. One forward pass delivers K+1 verdicts. If the typist guessed even half-right, the throughput of the system as a whole rises by exactly that factor.

Two properties make this work without compromising output quality. First, the accept-or-reject rule is greedy and exact: a draft token is accepted if and only if it equals what the main model would have produced at that position. Second, the rejected position is repaired using the main model's own argmax at that position — which is sitting right there in the verification forward pass. So the committed sequence is byte-for-byte identical to what pure autoregressive decoding would have produced. Speculative decoding is a speedup, not an approximation.

The Math of Acceptance and Speedup

Let $p$ be the per-token probability that a draft from the MTP heads agrees with what the main model would have emitted at the same position, and let $K$ be the number of draft tokens proposed each round. Acceptance is sequential — the moment one draft is rejected, all subsequent drafts are discarded because they were conditioned on a prefix the main model would not have written. So the number of accepted drafts $J$ in a round follows a truncated geometric distribution:

$\Pr[J = j] = \begin{cases} p^{j}(1 - p) & 0 \le j < K \\ p^{K} & j = K \end{cases}$

Every round commits $J + 1$ tokens (the $+1$ being the main model's free prediction at the rejection point — even when $J = K$ we get one extra position out of the forward pass). The expected number of tokens committed per round is therefore:

$\mathbb{E}[J + 1] = 1 + \sum_{j=1}^{K} p^{j} = \frac{1 - p^{K+1}}{1 - p}$

That ratio is, to a first approximation, the wall-clock speedup over autoregressive decoding — because every round costs essentially one main forward pass and the standard loop would need $\mathbb{E}[J+1]$ such passes to produce the same number of tokens. Three regimes are worth burning into intuition:

Why the speedup saturates with K

The sum $\sum_{j=0}^{K} p^j$ is a geometric series: its limit as $K \to \infty$ is $1 / (1 - p)$. With $p = 0.8$ that ceiling is exactly $5\times$. Doubling K past 3 or 4 barely moves the needle because the probability that draft 5 is right given draft 4 was already right is multiplied by yet another factor of $p$. This is why DeepSeek-V3 trains only one MTP head and still gets most of the available speedup at inference — the marginal benefit of head 4 is much smaller than the marginal training cost of training it well.

Why the cost of a draft is not zero — but close

The above formula treated each round as one main-model forward pass. The MTP heads themselves do cost something — but each head is a few linear layers riding on the same trunk output, so the K-head draft costs roughly $K \cdot c_{\text{head}}$ where $c_{\text{head}} \ll c_{\text{trunk}}$. Empirically the overhead is on the order of 1–3% of one main-model forward. So a speedup of $E[J+1] = 2.95$ in the table above lands as about $2.85\times$ after honest accounting — still decisive.

Manual Numerical Walkthrough

Let us run three concrete rounds with $K = 3$, $p = 0.8$, writing the toy sentence The small boat drifted into the harbour. Every number is worked out by hand so the mechanism is fully exposed.

Visualizing One Verification Round

The visualizer below builds the sentence one round at a time. Each round shows the K draft tokens (green = main model would have written this, red = rejected) and the corrected main-model token tacked on at the rejection point. Toggle between Autoregressive (forced K=1) and Speculative to feel the round-by-round difference. Slide K between 1 and 4, and the acceptance probability between 0.30 and 0.95, to explore the regime around $p \approx 0.8$ where DeepSeek-V3 actually operates.

Three things to lock in. First, the "tokens / round (observed)" stat is the empirical version of $\mathbb{E}[J + 1]$. With p=0.8 and K=3 it hovers around 3, exactly as the geometric series predicts. With K=1 it can never exceed 2 — that is the speedup ceiling for any 1-token draft model. Second, every round commits at least one token, even when the very first draft was wrong: the main model's "corrected" token always lands. There is no scenario where speculative decoding produces fewer tokens than autoregressive would have. Third, the speedup is bounded above by $1 / (1 - p)$: push K beyond 4 with p=0.8 and the gain past 5× becomes vanishingly thin.

Plain Python: A Speculative Decoding Loop

Below is the entire mechanism in plain Python — no PyTorch, no model. We replace the main model with a deterministic lookup table (it "wants to write" a fixed sentence) and the MTP heads with a coin flip per token. The loop counts rounds versus tokens so the speedup is visible.

1import random
2random.seed(0)
3
4# A toy "vocabulary" — what the main model wants to say.
5TRUE = ["The", "small", "boat", "drifted", "into", "the", "harbour"]
6WRONG = ["a", "ship", "of", "with", "in"]
7
8def main_model_next(prefix):
9    """The expensive main model: returns the next token given a prefix.
10       Cost = 1 forward pass per token in the standard loop."""
11    return TRUE[len(prefix)]
12
13def draft_model_chunk(prefix, K, p):
14    """The cheap MTP heads: return K speculative tokens.
15       Each one is 'right' with probability p (it agrees with the main model)."""
16    out = []
17    for i in range(K):
18        right = TRUE[len(prefix) + i]
19        out.append(right if random.random() < p else random.choice(WRONG))
20    return out
21
22def main_model_verify(prefix, drafts):
23    """One main-model forward pass over prefix+drafts. Returns the main
24       model's next-token at every prefix length — exactly what attention
25       gives us for free in a single forward."""
26    return [TRUE[len(prefix) + i] for i in range(len(drafts) + 1)]
27
28def speculative_step(prefix, K=3, p=0.8):
29    # 1. Draft K tokens with the cheap MTP heads.
30    drafts = draft_model_chunk(prefix, K, p)
31    # 2. ONE forward pass through the main model verifies them all.
32    targets = main_model_verify(prefix, drafts)
33    # 3. Accept the longest matching prefix; commit the corrected token.
34    j = 0
35    while j < K and drafts[j] == targets[j]:
36        j += 1
37    committed = drafts[:j] + [targets[j]]
38    return committed, j
39
40prefix = []
41rounds = 0
42while len(prefix) < len(TRUE):
43    committed, j = speculative_step(prefix, K=3, p=0.8)
44    prefix += committed
45    rounds += 1
46    print(f"round {rounds:>2}: accepted {j}/3, committed {len(committed)} → {prefix}")
47print(f"done in {rounds} rounds vs {len(TRUE)} for autoregressive")

Two structural details deserve a second pass. First, the verification call in main_model_verify intentionally returns K+1 targets, not K. That extra slot — the main model's prediction at position L+K, conditioned on all the drafts — is the "free" token that gets committed even when every single draft was correct. Forget the +1 and your best-case round only commits K tokens instead of K+1.

Second, the accept rule is greedy-equality: drafts are accepted iff they exactly match the main model's argmax. This works for greedy decoding (temperature 0). For temperature sampling or top-p sampling, the rule is replaced by the speculative-sampling probability-ratio test from Chen et al., 2023: accept draft token $x$ with probability $\min(1, q_{\text{main}}(x) / q_{\text{draft}}(x))$ where $q$ are the per-position output distributions. That rule preserves the main model's sampling distribution exactly — sampled outputs are still drawn from the main model, just much faster.

Sanity check the formula. Run the script ten times with random.seed(0..9) and compute the empirical average tokens per round. With $K=3, p=0.8$ you should see $\approx 2.95$ tokens per round — the geometric series talking back to you in code.

PyTorch: Verifying K Drafts in One Forward Pass

The production version replaces the deterministic table with real tensors: argmax over the vocabulary at K+1 positions of a single forward pass. The accept-or-reject decision lives entirely on the GPU, with no Python loop inside the hot path.

1import torch
2import torch.nn.functional as F
3
4@torch.no_grad()
5def speculative_decode_step(
6    main_model,        # huge transformer; one call = one big forward pass
7    mtp_heads,         # K small heads, share trunk with main_model
8    input_ids,         # (1, L)  — current prefix
9    K: int = 4,
10    temperature: float = 0.0,
11):
12    """One round of MTP-driven speculative decoding."""
13
14    # 1) Draft K tokens with the MTP heads.
15    #    The heads ride on the trunk's last hidden state, so we run the
16    #    trunk ONCE to fill the KV cache up to position L, then read
17    #    K parallel logits from the K heads.
18    trunk_out = main_model.trunk(input_ids)                  # (1, L, d)
19    last_h    = trunk_out.last_hidden_state[:, -1, :]        # (1, d)
20    draft_logits = torch.stack(
21        [h(last_h) for h in mtp_heads], dim=1
22    )                                                         # (1, K, V)
23    draft_ids = draft_logits.argmax(dim=-1)                  # (1, K)
24
25    # 2) Verify: feed [input_ids, draft_ids] into the main LM head in ONE
26    #    forward pass and read K+1 next-token logits.
27    verify_in  = torch.cat([input_ids, draft_ids], dim=1)    # (1, L+K)
28    verify_out = main_model(verify_in)                       # (1, L+K, V)
29    main_logits = verify_out.logits[:, -(K + 1):, :]         # (1, K+1, V)
30    main_ids    = main_logits.argmax(dim=-1)                 # (1, K+1)
31
32    # 3) Accept the longest matching prefix; commit the corrected token.
33    agree = (draft_ids[0] == main_ids[0, :K])                # (K,) bool
34    j = int(agree.cumprod(dim=0).sum().item())               # first-False trick
35    committed = torch.cat(
36        [draft_ids[:, :j], main_ids[:, j:j + 1]], dim=1
37    )                                                         # (1, j+1)
38    return torch.cat([input_ids, committed], dim=1), j

Three subtleties worth marking, all about how this loop interacts with the rest of an LLM serving stack:

1. The KV cache is reused across rounds. The verification forward pass populates the cache for positions $L \dots L + K - 1$. Only the cache entries past the rejection point need to be invalidated — the entries for the $j$ accepted positions are still valid because they were computed conditioned on the actual committed prefix. A well-implemented serving stack saves the entire trunk recomputation cost for those $j$ tokens.
2. Sampling preserves the main-model distribution. Speculative decoding is not an approximation. With greedy decoding the committed sequence is byte-for-byte identical to standard decoding; with temperature sampling the accept rule is the probability-ratio test and the marginal distribution of every committed token is still exactly the main model's. You cannot tell from the output whether a server is using speculative decoding under the hood.
3. Batching changes the arithmetic. If the server batches N requests, the main-model forward pass already amortises bandwidth across N sequences — the memory-bandwidth bottleneck loosens and the tensor cores warm up. Speculative decoding still helps, but the speedup shrinks because the autoregressive baseline was already faster per token. The win is largest in single-stream / low-batch interactive serving, which is also where latency matters most.

At Massive Scale: Memory-Bandwidth Wins

Speculative decoding looks like a clever programming trick until you do the arithmetic for a real 671B-parameter MoE model. The DeepSeek-V3 numbers turn it into a deployment necessity.

Two patterns to read. First, every serving stack listed here uses speculative decoding in production — the question is where the drafts come from. Llama-3 70B uses a separately trained small Llama-3 8B draft model. DeepSeek-V3 uses the MTP heads that already exist for free because they were part of the training objective. The MTP route saves gigabytes of draft-model memory and avoids the distribution-shift problem that hurts external draft models (a separately trained 8B does not always agree with what the 70B would have written).

Second, the speedup is meaningfully larger for MoE models than the bare formula predicts. The reason is that an MoE forward pass only activates a sliver of the parameters (37B out of 671B for DeepSeek-V3) but still has to load all 671B of expert weights into HBM cache lines in the worst case. Speculative decoding amortises that loading cost over K+1 tokens instead of one. The HBM-bandwidth ceiling moves up by the same factor as your accept rate.

The interaction with expert parallelism

When experts are spread across GPUs (Chapter 5), each MoE forward pass triggers an all-to-all shuffle of tokens to the right experts and back. Speculative decoding multiplies the "tokens" in that shuffle by K+1 — which sounds bad until you remember that the all-to-all cost is almost entirely fixed-latency setup, not per-token data. Sending 1 token and sending 4 tokens through the same all-to-all costs almost the same wall-clock time. So speculative decoding turns the all-to-all latency into a smaller fraction of the per-token cost. The serving cluster gets cheaper, not more expensive.

The interaction with FP8 inference

DeepSeek-V3's FP8 training (Chapter 10) carries over to inference, cutting the model weight size roughly in half versus FP16. That shifts the memory-bandwidth ceiling upward by about $2\times$ on its own. Speculative decoding then stacks on top, so a serving node that produced 18 tokens/s/H100 in FP16 autoregressive mode can produce roughly $18 \cdot 2 \cdot 2 \approx 72$ tokens/s/H100 in FP8 + spec-decode mode. That 4× compound speedup is what actually makes a 671B-parameter MoE model affordable to serve.

Engineering Reality and Gotchas

MTP-driven speculative decoding looks like a clean win on paper. Three production failure modes earn their flags:

1. Draft / main distribution drift. The MTP heads were trained alongside the main model on the same data — but the main model continued to evolve through any fine-tuning, RLHF, or DPO stage that followed pre-training. If the alignment phase did not include the MTP heads in its optimisation, the heads' acceptance rate can degrade sharply on the new distribution (chat, instruction-following, refusal behaviour). The fix is to either re-tune the heads after alignment or back-propagate the alignment loss through them. DeepSeek-V3 reports a ≈4-point acceptance drop on chat traffic that they recover via head re-tuning.
2. K too large hurts more than it helps. Pushing K past 4 rarely improves throughput on real workloads, and it can hurt: every extra draft adds head-compute, and the geometric saturation means the last drafts are mostly wasted. Production schedulers monitor the running acceptance rate per request and cap K dynamically — typically K=2 for chat, K=3 for code, K=4 for highly templated outputs.
3. The verification forward pass needs careful masking. The K draft tokens added to the input must still be attended to causally — draft i sees drafts 0..i-1 but not i+1..K-1. If the attention mask accidentally lets later drafts leak into earlier positions, the verification logits become contaminated and the accept rate craters silently. This is the single most common bug in custom speculative-decoding kernels. Always unit-test by comparing a spec-decode forward pass against the autoregressive baseline at greedy temperature: they must produce byte-identical outputs.

The one sentence to carry forward: the MTP heads paid for themselves twice — once during training by improving sample efficiency, once at inference by turning every otherwise-idle tensor core into committed tokens — and that double-dip is the engineering reason a 671B-parameter MoE model can match dense-70B serving cost.