MTP

MTP#

MTP (multi-token prediction) uses a full MoE decoder layer as the draft model, trained with test-time training (TTT). It is designed for models whose architecture includes a native MTP head — the draft is a complete decoder layer (grouped-query attention with QK-norm, a sigmoid-router mixture-of-experts, and a shared expert) plus the fusion modules that combine the token embedding with the target’s hidden state.

Unlike the DFlash family, MTP is not block-parallel; it predicts one token at a time and is independent of the Eagle3 code path.

How it works#

  • Input fusion. Each step normalizes the token embedding and the target’s last hidden state (enorm / hnorm), concatenates them, and projects back to the hidden size (eh_proj) before the decoder layer.

  • TTT loop with a KV cache. The output hidden of each step is fed back as the next step’s input, and the inputs are shifted so that step k predicts the token k positions ahead. A per-step key/value cache lets each step reuse the previous steps’ attention state.

  • Single last hidden state. MTP consumes only the target’s final hidden state — there is no multi-layer auxiliary fusion.

  • Tied lm_head. The draft shares the target model’s frozen lm_head for both the CE target and the KL teacher, so it adds no lm_head parameters.

  • Loss. Cross-entropy against the target’s argmax plus KL divergence against the target’s soft distribution, with per-step weights that decay over the prediction horizon.

Attention backends#

MTP’s cached attention has three interchangeable backends, selected by attention_backend:

Backend

Use case

sdpa

Default. Eager cached attention; simplest, but memory grows with sequence length.

fa

Long sequences. The first block runs through compiled FlexAttention (linear memory in sequence length); later cached steps merge in with log-sum-exp.

usp

Very long sequences. fa plus Ulysses sequence parallelism to shard a 128k+ sequence across GPUs.

Initialization#

MTP can be initialized two ways:

  • Warm-start from a target’s MTP layer. If the target checkpoint contains a trained MTP layer, its weights are loaded as the draft’s starting point (continual training).

  • From scratch. For targets without a native MTP layer, define a dimension-matched draft (same hidden size and expert count as the target’s decoder layers) and train from random initialization.

In both cases the draft must be dimension-matched to the target model that will serve the hidden states, so acceptance-rate metrics are meaningful end to end.

Serving#

Training checkpoints are converted to HuggingFace format for serving. The converter is the inverse of the initialization step: it unfuses the expert weights back to the per-expert layout and writes the keys where the serving engine expects the MTP layer. The embedding and lm_head are borrowed from the main model at load time, matching the tied/frozen design. See Checkpoint Conversion.