Sequential Causal MTP: DeepSeek's Implementation

Section 7.2 ended on a sharp note. Meta's naive parallel MTP saved compute by predicting all $D$ future tokens at once from a single shared trunk — and paid for it by silently breaking the causal chain that makes autoregressive language modeling work. Predictions for distant horizons drifted, the heads stopped agreeing with each other, and downstream metrics fell behind a single-token baseline at large model sizes. The lesson was unambiguous: if you want extra prediction heads, they must respect causality. DeepSeek's sequential MTP is the answer that fell out of taking that lesson seriously.

The promise of sequential causal MTP. Each depth is its own small transformer block. Depth $k$ at position $i$ reads the previous depth's hidden state at the same position, mixes it with the embedding of the token at position $i + k$, and predicts the token at position $i + k + 1$. The causal chain survives intact. The signal stays sharp. The trick is the architecture, not a loss-function patch.

Why DeepSeek Went Sequential

Recall the parallel design from section 7.2: $D$ independent output heads sit on top of the same final hidden state $h_i$, and head $k$ tries to predict the token at position $i + k$. The trunk gets a richer gradient signal — every position contributes $D$ cross-entropy terms instead of one — but the heads cannot talk to each other. Head 2 has no idea what head 1 just predicted. The two predictions become marginal distributions over the joint, not a chain rule factorization. Empirically, that gap kills the win.

DeepSeek reasoned backwards from what the chain rule actually wants. The true joint over future tokens is

$P(t_{i+1}, t_{i+2}, \dots, t_{i+D} \mid t_{\le i}) = \prod_{k=1}^{D} P(t_{i+k} \mid t_{\le i + k - 1})$

Notice what the conditioning set looks like. To predict $t_{i+k}$ properly, the model needs everything up to and including $t_{i+k-1}$ — that is, the previously predicted future tokens. Parallel MTP throws this away: it conditions every head on $t_{\le i}$ alone, which is the wrong marginal for $k > 1$. Sequential MTP gives every depth the conditioning set it needs.

Anatomy of a Single MTP Module

An MTP module is the smallest unit you can imagine that still does useful autoregressive work. It is, almost literally, "one transformer block plus a few wires." Let us list every part:

1. Two inputs per position. The previous depth's hidden state $h_i^{k-1}$ and the embedding of the token at position $i + k$, which we write $\mathrm{Embed}(t_{i+k})$. The hidden state is the past; the embedding is the next-known future token.
2. RMSNorm on each input. Two independent RMSNorms — one per stream — strip away differing magnitudes so the downstream projection sees inputs on a comparable scale.
3. Concatenation, then projection. The two normalized vectors are concatenated along the feature axis (giving a length-$2d$ vector) and then projected back down to $d$ by a learned matrix $M_k$. This is where past and future meet.
4. One transformer block. A single transformer block — attention + MLP + norms — is applied to the projected stream. It uses the same causal mask as the main model, so position $i$ still cannot see positions $j > i$ within the block.
5. Shared output head. The block's output is run through the main model's output head (the linear layer that maps $d$-dim hidden vectors to $V$-dim logits), producing a vocabulary distribution.

What is shared, what is per-depth

The sharing story is what makes the design cheap. Two pieces are shared with the main model across every depth:

• The embedding table $\mathrm{Embed} \in \mathbb{R}^{V \times d}$ — the same one the main model uses on its inputs.
• The output head $W_{\text{out}} \in \mathbb{R}^{V \times d}$ — tied to the embedding (weight tying) on most LLM setups.

Three pieces are owned per depth:

• The two RMSNorms ($\sim 2d$ parameters).
• The projection $M_k \in \mathbb{R}^{d \times 2d}$ ($2d^2$ parameters).
• One transformer block (~$12 d^2$ parameters at standard FFN expansion). This is the bulk of the per-depth cost.

The Math: A Causal Chain Through Depths

Fix a position $i$ in the sequence. Let $h_i^0$ denote the main model's final hidden state at position $i$ — this is the starting point for all MTP depths and is, by construction, a function only of $t_{\le i}$. For depths $k = 1, 2, \dots, D$, define the MTP module recursively:

$h_i^k = \mathrm{TRM}_k\Big(\, M_k \, \mathrm{concat}\big[\mathrm{RMSNorm}(h_i^{k-1}),\; \mathrm{RMSNorm}(\mathrm{Embed}(t_{i+k}))\big]\,\Big)$

Every symbol: $\mathrm{TRM}_k$ is a single per-depth transformer block; $M_k \in \mathbb{R}^{d \times 2d}$ is the per-depth projection; $\mathrm{RMSNorm}$ is the root-mean-square norm with two independent learnable scales (one per stream); $\mathrm{Embed}$ is the shared embedding table; and $t_{i+k}$ is the actual ground-truth token at position $i + k$ in the training sequence (known at training time because the whole sequence is given).

The depth-$k$ prediction at position $i$ is then

$p^k_{i+k+1} = \mathrm{softmax}\!\big(W_{\text{out}} \, h_i^k\big) \;\in\; \Delta^{V-1}$

— a distribution over the vocabulary, predicting the token at position $i + k + 1$. $W_{\text{out}}$ is the shared (tied) output head. The cross-entropy loss at this depth compares $p^k_{i+k+1}$ to the one-hot of $t_{i+k+1}$; we will write the full training objective in section 7.4.

Why the recursion is the whole story

The single equation above hides three independent design choices, each of which corrects one specific failure of parallel MTP:

1. The recursion through $h_i^{k-1}$. This is the chain rule made concrete. Depth $k$ conditions on everything depth $k - 1$ knew, plus the new token $t_{i+k}$. Parallel MTP had every depth read $h_i^0$ directly — no chain.
2. The embedding of the future token. At training time we know $t_{i+k}$, so we feed it in. Without it, depth $k$ would have no information beyond what depth $k - 1$ already used — the module would be a redundant copy of its predecessor.
3. The independent transformer block per depth. Each depth gets its own parameters because each depth solves a slightly different problem — depth 1 maps "past + next token" → "next-next token", depth 2 maps "past + next-next token" → "next-next-next token". The tasks are related but not identical.

Manual Numerical Walkthrough

Let us pin down a single depth-1 MTP step end-to-end with tiny numbers. We will use $d = 4$, $V = 6$, and look at the prediction made for position $i + 2 = 4$ from position $i = 2$.

M = [[ 0.10,  0.20, -0.10,  0.05,  0.15, -0.05,  0.10,  0.20],
     [-0.20,  0.10,  0.30, -0.10,  0.05,  0.20, -0.10,  0.10],
     [ 0.15, -0.10,  0.05,  0.20, -0.10,  0.10,  0.30, -0.05],
     [ 0.05,  0.20, -0.10,  0.15,  0.20, -0.10,  0.05,  0.10]]

W_out = [[ 0.5,  0.1, -0.2,  0.3],   # vocab id 0
         [-0.3,  0.4,  0.2,  0.1],   # vocab id 1
         [ 0.2, -0.1,  0.5, -0.2],   # vocab id 2
         [ 0.4,  0.3,  0.1,  0.6],   # vocab id 3 <- TARGET
         [-0.1,  0.2, -0.3,  0.4],   # vocab id 4
         [ 0.1, -0.4,  0.2, -0.1]]   # vocab id 5

Visualizing the Sequential Forward Pass

The diagram below walks through a six-token sequence with one main transformer and two MTP modules stacked on top. Use ▶ Play to watch the depths fill in left-to-right, then hover any cell to inspect what enters and what comes out. The key observation: cell $h_i^k$ always reads cell $h_i^{k-1}$ from the row above and the embedding of the token $k$ positions to its right — never any cell to its right at its own depth. Causality holds at every depth.

Plain Python: One MTP Module by Hand

Before reaching for PyTorch, let us implement a depth-1 MTP module in plain NumPy. The goal is to expose every shape, every matmul, every normalization — no autograd, no broadcast magic. If you understand these $\sim 30$ lines, you understand the architecture.

1import numpy as np
2
3# Tiny toy: V = 8 vocab tokens, d = 4 hidden dim, T = 5 sequence positions.
4np.random.seed(0)
5V, d, T = 8, 4, 5
6
7# (a) Shared embedding table — same one the main model uses.
8Embed = np.random.randn(V, d) * 0.1
9
10# (b) Shared output head — also identical to the main model's.
11W_out = np.random.randn(V, d) * 0.1
12
13# (c) The toy "main model": pretend its hidden states are already computed.
14#     h0[i] is what the main model produced at position i.
15h0 = np.random.randn(T, d) * 0.5
16
17# (d) Per-depth MTP parameters. We do D = 1 (one MTP module).
18M = np.random.randn(d, 2 * d) * 0.1        # projection: 2d -> d
19W_block = np.random.randn(d, d) * 0.1      # toy 1-layer "transformer block"
20
21# Token IDs in the sequence (which words appeared at each position).
22tokens = np.array([3, 1, 4, 0, 5])         # shape (T,)
23
24def rmsnorm(x, eps=1e-6):
25    rms = np.sqrt((x ** 2).mean(axis=-1, keepdims=True) + eps)
26    return x / rms
27
28# ---- Sequential MTP, depth k = 1 ----
29# For each position i, predict the token at position i + 2 (one beyond next).
30# Module 1 needs: (h0[i], Embed[token at position i + 1]).
31predictions = []
32for i in range(T - 2):                       # last 2 positions have no target
33    # (1) RMSNorm both inputs.
34    h_norm = rmsnorm(h0[i])                  # (d,)
35    e_norm = rmsnorm(Embed[tokens[i + 1]])   # (d,)
36
37    # (2) Concatenate along the feature axis and project back to d.
38    cat = np.concatenate([h_norm, e_norm])   # (2d,)
39    combined = M @ cat                        # (d,)
40
41    # (3) One "transformer block" — toy linear stand-in.
42    h1_i = np.tanh(W_block @ combined)        # (d,)
43
44    # (4) Shared output head → vocabulary logits → softmax.
45    logits = W_out @ h1_i                     # (V,)
46    z = logits - logits.max()
47    p = np.exp(z) / np.exp(z).sum()           # (V,)
48    predictions.append((i, p.argmax(), tokens[i + 2]))
49
50for i, pred, target in predictions:
51    print(f"pos {i}: predict t+2 -> {pred:<3d}  target -> {target}")

Three observations are worth pulling out. First: the loop body is position-independent — every position runs the exact same arithmetic with different inputs. Second: the only place the future-token information enters is line 33, the embedding lookup at $i + 1$. Third: the projection $M$ is the only weight that mixes the two streams; everything else either normalizes a single stream or operates on the already-mixed combined vector.

PyTorch: A Reusable MTP Module Class

In production code the MTP module is just an $\texttt{nn.Module}$. It is small, vectorized over $(B, T)$, and takes the shared embedding and output head as forward-time arguments so they cannot be silently duplicated.

1import torch
2import torch.nn as nn
3import torch.nn.functional as F
4
5
6class MTPModule(nn.Module):
7    """
8    One depth of DeepSeek-style sequential MTP.
9
10    Forward signature:
11        h_prev  : (B, T, d)   hidden states from the previous depth.
12        tok_fut : (B, T)      token IDs at position i + k for each i.
13
14    Outputs:
15        h_next  : (B, T, d)   hidden states at this depth.
16        logits  : (B, T, V)   distribution over vocabulary.
17
18    Sharing contract (set by the parent model, NOT by this module):
19      • embed        — shared with the main model (nn.Embedding(V, d))
20      • output_head  — shared with the main model (Linear(d, V) tied to embed)
21    Per-depth (owned by this module):
22      • proj         — Linear(2d, d)
23      • block        — one transformer block
24      • norm_h, norm_e — RMSNorm for each input stream
25    """
26
27    def __init__(self, d_model: int, block: nn.Module):
28        super().__init__()
29        self.norm_h = nn.RMSNorm(d_model)            # for the previous hidden state
30        self.norm_e = nn.RMSNorm(d_model)            # for the future-token embedding
31        self.proj = nn.Linear(2 * d_model, d_model, bias=False)
32        self.block = block                            # one transformer block
33
34    def forward(
35        self,
36        h_prev: torch.Tensor,                         # (B, T, d)
37        tok_fut: torch.Tensor,                        # (B, T)
38        embed: nn.Embedding,                          # shared with main model
39        output_head: nn.Linear,                       # shared with main model
40        attn_mask: torch.Tensor | None = None,        # causal mask for self-attn
41    ):
42        # (1) Normalize both streams BEFORE mixing them.
43        h_norm = self.norm_h(h_prev)                  # (B, T, d)
44        e_norm = self.norm_e(embed(tok_fut))          # (B, T, d)
45
46        # (2) Concatenate along the feature axis, then project back to d.
47        cat = torch.cat([h_norm, e_norm], dim=-1)     # (B, T, 2d)
48        x = self.proj(cat)                             # (B, T, d)
49
50        # (3) One transformer block. The causal mask MUST be active here —
51        #     the MTP module is still autoregressive within the sequence.
52        h_next = self.block(x, attn_mask=attn_mask)   # (B, T, d)
53
54        # (4) Logits via the shared (tied) output head.
55        logits = output_head(h_next)                  # (B, T, V)
56
57        return h_next, logits

How the parent model wires it together

The parent model owns the embedding, the main transformer stack, the output head, and a small list of $\texttt{MTPModule}$ instances. The forward pass looks roughly like:

# main model
h = embed(tokens)
for blk in main_blocks:
    h = blk(h, attn_mask)
h0 = h                                          # (B, T, d)
logits0 = output_head(h0)                       # main next-token logits

# MTP depths
h_prev = h0
mtp_logits = []
for k, mtp in enumerate(mtp_modules, start=1):
    tok_fut = shift_left(tokens, k)             # (B, T) tokens at i + k
    h_prev, logits_k = mtp(h_prev, tok_fut, embed, output_head, attn_mask)
    mtp_logits.append(logits_k)

Two PyTorch-specific subtleties show up here. First, $\texttt{shift\_left}$ is not just a Python slice — it must respect padding tokens so the loss is not computed on garbage positions; almost every real implementation passes a position mask alongside. Second, the optimizer sees each MTP module's parameters but NOT a duplicate copy of the embedding or output head, because those were passed positionally rather than registered as sub-modules. That is the whole point of the explicit-arg sharing convention.

What Changes at Massive Scale

At a scale where the main model has 61 transformer blocks, 7168 hidden dimensions, and 671B parameters in total (with about 37B active per token), the MTP module's budget changes character. Let us walk through the costs.

The memory story

The dominant memory cost during training is activations, not parameters. An MTP module carries one block's worth of activations per token — and those activations live until the backward pass for the MTP loss runs. With activation checkpointing the cost can be amortized, but the simpler picture is: budget for $D + 1$ blocks of activations instead of $L$ blocks. At $L = 61$ and $D = 1$, that is a 1.6% bump.

The communication story

In data-parallel training, gradient all-reduces dominate the network budget. MTP modules add their parameters to the bucket — about 1.6% more grad data per step. In tensor-parallel or pipeline-parallel layouts, the MTP module's transformer block can be placed on the same shard as the final main block; no extra cross-device traffic. In FSDP, the MTP module is one more wrap unit — trivial to integrate.

The throughput story

Sequential MTP is sequential, full stop. Depth $k$ waits for depth $k - 1$. There is no pipeline trick that breaks this dependency without breaking causality — the same dependency that made the parallel design tempting in the first place. At $D = 1$ the cost is one block of extra depth, which on H100s at 4k sequence length is in the low single-digit percent of step time. Past $D = 2$ the latency cost starts mattering and the per-depth quality return diminishes — which is exactly what DeepSeek's ablations show.

Engineering Reality and Gotchas

Sequence length and the boundary problem

Depth $k$ needs the token at position $i + k$. For the last $k$ positions of the sequence that token does not exist. Two correct ways to handle this:

• Mask the loss on positions $i \ge T - k - 1$ at depth $k$. Cleanest, costs a couple of token-positions per sequence at $D = 1$.
• Right-pad the input by $D$ dummy tokens so every position has a valid future. Simpler at the call site but wastes a hair of compute.

The mask-tying trap

The transformer block inside the MTP module needs its own causal mask — the same shape as the main model's. A common bug is reusing the main model's pre-computed mask after the embedding or positional bias has been folded in; the MTP block ends up attending differently than the main one, sometimes leaking future info, sometimes silently masking the present position. Make the MTP block compute its own mask, or pass an unbiased causal mask explicitly.

RMSNorm scale init

The two RMSNorm gain vectors inside an MTP module initialize to 1, same as everywhere else. But because RMSNorm sits on a stream whose magnitude depends on the main model's drift during training, empirical wisdom is to warm-start the gain on the hidden stream at a slightly lower value (e.g. $0.5$) for the first ~1000 steps. This prevents an early stage where the projection sees huge h-stream norms and tiny e-stream norms; without the warm-start the projection's gradient blows up.

Loss scaling and lambda

The full training objective will be detailed in section 7.4, but a preview: the MTP loss is added with a coefficient $\lambda$ typically in the range $[0.1, 0.3]$. Setting it too high crowds out the main next-token loss; setting it too low makes the MTP module learn slowly and underperform at inference-time speculation. DeepSeek-V3 uses $\lambda$ around 0.3 during early training and decays it linearly.

What "sharing" really means at training time

Weight sharing is a contract enforced by the parent module, not by the MTP class itself. The cleanest way to enforce it in PyTorch is to never store the embedding or the output head as a child module of the MTP instance — pass them in. Otherwise $\texttt{parameters()}$ double-counts them and the optimizer applies two updates per step. Most production bugs in MTP implementations come from accidentally registering the shared tensors twice.

The takeaway. Sequential causal MTP is what you get when you refuse to compromise on the chain rule and refuse to spend more than a few percent of the main model's budget. One small transformer block per depth, two RMSNorms, one projection, shared embedding and output head, run in a strict causal sequence. The rest of the chapter — the loss formulation in section 7.4 and the speculative-decoding payoff in section 7.5 — is just downstream of getting this one architectural choice right.