main
md 76 lines 2.47 KB
Rendered Raw
1 # MTP (Multi-Token Prediction)
2
3 MTP is a speculative decoding method where the target model includes native
4 multi-token prediction capability. Unlike draft-model-based methods, you do not
5 need to provide a separate draft model.
6
7 MTP is useful when:
8
9 - Your model natively supports MTP.
10 - You want model-based speculative decoding with minimal extra configuration.
11
12 ## Gemma 4 Assistant Models
13
14 Gemma 4 assistant checkpoints use vLLM's Gemma 4 MTP path. They are not generic
15 draft models, even though they are passed through the `model` field in
16 `--speculative-config`.
17
18 Use `"method": "mtp"` when serving Gemma 4 with an assistant checkpoint:
19
20 ```bash
21 vllm serve google/gemma-4-E2B-it \
22 --tensor-parallel-size 1 \
23 --max-model-len 8192 \
24 --speculative-config '{"method":"mtp","model":"gg-hf-am/gemma-4-E2B-it-assistant","num_speculative_tokens":1}'
25 ```
26
27 The E2B, E4B, 12B, 26B-A4B, and 31B Gemma 4 IT assistant checkpoints are supported.
28 Tower-based variants use `model_type: gemma4_assistant` and the encoder-free
29 Gemma 4 Unified variant (12B) uses `model_type: gemma4_unified_assistant`.
30 vLLM maps both to `Gemma4MTPModel` internally and wires the assistant layers
31 to share KV cache with the target model.
32
33 If an older vLLM release logs `SpeculativeConfig(method='draft_model', ...)`
34 for a Gemma 4 assistant checkpoint, that release is treating the assistant as a
35 generic draft model and may fail during initialization for multimodal Gemma 4
36 targets. Upgrade to a version with Gemma 4 MTP support instead.
37
38 ## Offline Example
39
40 ```python
41 from vllm import LLM, SamplingParams
42
43 prompts = ["The future of AI is"]
44 sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
45
46 llm = LLM(
47 model="XiaomiMiMo/MiMo-7B-Base",
48 tensor_parallel_size=1,
49 speculative_config={
50 "method": "mtp",
51 "num_speculative_tokens": 1,
52 },
53 )
54 outputs = llm.generate(prompts, sampling_params)
55
56 for output in outputs:
57 prompt = output.prompt
58 generated_text = output.outputs[0].text
59 print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
60 ```
61
62 ## Online Example
63
64 ```bash
65 vllm serve XiaomiMiMo/MiMo-7B-Base \
66 --tensor-parallel-size 1 \
67 --speculative-config '{"method":"mtp","num_speculative_tokens":1}'
68 ```
69
70 ## Notes
71
72 - MTP only works for model families that support MTP in vLLM.
73 - `num_speculative_tokens` controls speculative depth. A small value like `1`
74 is a good default to start with.
75 - If your model does not support MTP, use another method such as EAGLE or draft
76 model speculation.