Speculative decoding (MTP)

Speculative decoding lets a smaller assistant propose future tokens while the target model verifies them in parallel. mistral.rs exposes this through the MTP (Multi-Token Prediction) API: attach an assistant checkpoint and the engine drafts several tokens per target step.

Output stays exact: every accepted token is verified by the target model before it is emitted.

Caution

The target must run with paged attention (the default on CUDA). Loading fails with an error if paged attention is off.

CLI

mistralrs run -m google/gemma-4-E4B-it --quant 8 \
  --mtp-model google/gemma-4-E4B-it-assistant \
  --mtp-n-predict 6

--mtp-model accepts a Hugging Face id or a local path. The same flags work with mistralrs serve. See run and serve flag references.

Python

from mistralrs import Runner, Which

runner = Runner(
    which=Which.MultimodalPlain(model_id="google/gemma-4-E4B-it"),
    in_situ_quant="8",
    mtp_model="google/gemma-4-E4B-it-assistant",
    mtp_n_predict=6,
)

Gemma 4 loads via Which.MultimodalPlain; it is currently the only model family with MTP assistant checkpoints.

Rust

let model = mistralrs::ModelBuilder::new("google/gemma-4-E4B-it")
    .with_mtp_model("google/gemma-4-E4B-it-assistant", Some(6))
    .build()
    .await?;

For full control, with_mtp_config(MtpConfig { model, n_predict }) is equivalent. The MTP builder methods exist on the text, multimodal, and auto-detecting model builders.

--mtp-n-predict controls how many assistant tokens are proposed per step. If it is omitted, mistral.rs reads num_assistant_tokens from the assistant’s generation_config.json and falls back to 6.

Supported models

Mode	Target models	Assistant model	Status
MTP	Gemma 4	Gemma 4 assistant checkpoints (--mtp-model)	Supported with paged attention
MTP	Qwen3.5, Qwen3.8	Built into the checkpoint (--mtp)	Supported with paged attention
DFlash / DFlash 2	Qwen3.5, Qwen3.8	DFlash draft checkpoints (--mtp-model)	Supported with paged attention

Legacy target/draft speculative decoding has been removed. New speculative decoding features use the MTP proposer/target path.

Qwen3.5 / Qwen3.8

Qwen3.5 and Qwen3.8 checkpoints ship a multi-token prediction head (mtp.* weights) sharing the target’s embeddings and lm_head. Pass --mtp to load it; no separate assistant model is needed:

mistralrs serve -m Qwen/Qwen3.8-27B --quant 4 --mtp

The head keeps its own paged KV cache alongside the target’s, so the KV budget grows by one attention layer. --mtp-n-predict defaults to 2 for these heads. Greedy output is equivalent to non-MTP decoding up to floating-point tie-breaks in the hybrid recurrent layers. UQFF artifacts must include the mtp.* tensors to use --mtp.

DFlash and DFlash 2 (Qwen3.5 / Qwen3.8)

DFlash drafters are small block-diffusion models that propose a whole block of tokens in one pass, conditioned on hidden states tapped from the target’s intermediate layers. They reach much higher acceptance than the built-in single-token MTP head. Pass a DFlash checkpoint (a HF id or local path) as --mtp-model; mistral.rs detects DFlash vs DFlash 2 from the checkpoint’s architectures and loads the path selector and dynamic convolutions when present:

mistralrs serve -m Qwen/Qwen3.8-27B --quant 4 \
  --mtp-model incoai/Qwen3.8-27B-DFlash2

By default the draft depth is adaptive: it starts at min(block_size - 1, 7) and settles per workload on a tier from {3, 5, max} driven by an EMA of the measured acceptance, so easy content (math, code) keeps deep drafts while hard content (thinking-mode reasoning) drops to shallow ones instead of paying for rejected verify rows. Adaptation applies at low concurrency (up to 2 concurrent decodes); larger batches always draft at full depth, where the batched drafter makes deep drafts nearly free. Passing --mtp-n-predict pins a fixed depth; MISTRALRS_DFLASH_ADAPTIVE=0 disables adaptation, =1 forces it even with an explicit depth (the explicit value becomes the maximum). The drafter’s projection weights are requantized to the target’s in-situ quantization type at load (override with the MISTRALRS_DFLASH_ISQ environment variable, e.g. q6k or none for bf16).

Acceptance is strongly content-dependent: incoai/Qwen3.8-27B-DFlash2 accepts ~7 of 7 drafts on math/code and much less on thinking-mode reasoning traces, which are out of its training distribution. As everywhere in this engine, drafts are verified by the target: greedy output is byte-identical to non-speculative decoding.

Gemma 4

Gemma 4 assistant checkpoints are MTP drafters for Gemma 4 target models. See the google/gemma-4-E4B-it-assistant model card for the upstream checkpoint. A downloaded checkout works too:

mistralrs run -m google/gemma-4-E4B-it --quant 8 \
  --mtp-model ./gemma-4-E4B-it-assistant \
  --mtp-n-predict 6

Non-paged KV-cache MTP is intentionally disabled for now, which is why paged attention is required (see the note at the top).

The target and assistant configs must match where the implementation requires it, including vocabulary size and target hidden size. Mismatches fail at load, before generation starts.

Notes

MTP remains exact: accepted output is verified by the target model before it is emitted. Throughput gain depends on how many proposed tokens the target accepts and on the cost of the target verification pass.

MTP supports batched generation and constrained decoding.

MTP is configured at launch time only (CLI flags, Runner(...), or the model builder). There is no per-request HTTP field to toggle it; load the server with the assistant attached and every request uses it.