Skip to content
Open
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 20 additions & 0 deletions docs/source/en/assisted_decoding.md
Original file line number Diff line number Diff line change
Expand Up @@ -141,6 +141,26 @@ outputs = model.generate(**inputs, assistant_early_exit=4, do_sample=False, max_
tokenizer.batch_decode(outputs, skip_special_tokens=True)
```

## Multi-token prediction

Multi-token prediction (MTP) drafts candidate tokens with extra prediction layers trained into the model itself, so it does not need a separate assistant model. The MTP layers predict several future tokens from a single position and share the main model's embeddings and output head. The main model verifies the drafts in one forward pass, which speeds up generation without loading a second model.

MTP works only with checkpoints trained with MTP layers, such as [DeepSeek-V3](https://huggingface.co/deepseek-ai/DeepSeek-V3) and [GLM-4.5](https://huggingface.co/zai-org/GLM-4.5). These models have a `num_mtp_layers` config value and MTP weights. [`~GenerationMixin.generate`] raises an error if the checkpoint has no MTP layers.

Pass `use_mtp=True` to [`~GenerationMixin.generate`]. Like speculative decoding, MTP supports greedy search and sampling but not batched inputs.

```py
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("zai-org/GLM-4.5-Air")
model = AutoModelForCausalLM.from_pretrained("zai-org/GLM-4.5-Air", device_map="auto")
inputs = tokenizer("Hugging Face is an open-source company", return_tensors="pt").to(model.device)

outputs = model.generate(**inputs, use_mtp=True, do_sample=False, max_new_tokens=50)
tokenizer.batch_decode(outputs, skip_special_tokens=True)
```

## Universal assisted decoding

Universal assisted decoding (UAD) makes speculative decoding possible even when the main and assistant models have different tokenizers. It lets you pair any small assistant model with the main model. Candidate tokens are re-encoded and the algorithm computes the longest common subsequence so the continuation stays aligned.
Expand Down
Loading