Gemma 4 MTP 推理加速实战：开源模型推理速度提升 3 倍的秘密

推理速度，是开源大模型从”能跑”到”能用”之间最关键的一道坎。

一个 31B 参数的模型，在单张 GPU 上标准自回归解码只能做到 8-10 tokens/sec——这意味着生成一段 500 字的回答需要等待将近一分钟。对于实时对话、代码补全、在线翻译这类对延迟敏感的场景，这个速度远远不够。

2026 年 5 月 5 日，Google 为 Gemma 4 全系列模型发布了专属的 MTP（Multi-Token Prediction）drafter 模型。通过 speculative decoding 机制，推理速度可以提升至原来的 3 倍——31B 模型在消费级 GPU 上达到约 24 tokens/sec。

本文将从 MTP 的工作原理讲起，手把手演示如何在 vLLM 和 SGLang 中配置和使用 Gemma 4 MTP，并给出性能调优的实战建议。

MTP 原理解析：一次预测多个 token

传统自回归解码的瓶颈

标准的大语言模型采用自回归（autoregressive）方式生成文本：每一步只预测下一个 token，拿到结果后再预测下下个 token。这意味着生成 N 个 token 需要执行 N 次完整的前向传播。

问题在于，每次前向传播都需要加载数十亿参数的权重矩阵，但实际的计算负载（单个 token）远远没有把 GPU 的并行算力用满。这种访存瓶颈（memory-bound）导致 GPU 利用率极低，大量计算单元在空转。

Multi-Token Prediction 的核心思路

MTP 的解决方案非常直观——既然 GPU 有富余的算力，为什么不同时预测多个 token？

具体流程分为两个阶段：

阶段一：草拟（Draft）

一个轻量级的 drafter 模型快速生成 K 个候选 token。由于 drafter 模型参数量远小于主模型（例如 78M vs 31B），单次前向传播速度极快。

阶段二：验证（Verify）

主模型（target model）对这 K 个候选 token 进行一次性并行验证。验证的计算量约等于一次标准前向传播，但可以同时处理多个位置。如果候选 token 被接受，就直接输出；如果被拒绝，则从主模型的分布中重新采样。

传统解码：  [t1] → [t2] → [t3] → [t4] → [t5]    需要 5 次前向传播
MTP 解码：  [t1,t2,t3,t4,t5] → 验证 → 接受 4 个   只需 1+1 次前向传播

💡 关键性质：这种”草拟-验证”机制在数学上保证了输出分布与直接使用主模型完全一致，不会降低生成质量。

Gemma 4 MTP 与传统 Speculative Decoding 的区别

传统的 speculative decoding 通常用一个独立的小模型（比如用 Llama-8B 给 Llama-70B 做草拟），但存在几个问题：

1. 词表不匹配：两个独立模型可能使用不同的 tokenizer
2. 分布差异大：独立训练的小模型与大模型的概率分布差距较大，接受率低
3. 需要额外加载完整模型：小模型虽然参数少，但仍是一个完整的语言模型

Gemma 4 的 MTP drafter 不同——它是 Google 专门为 Gemma 4 训练的配套组件：

• 共享主模型的 tokenizer 和词表
• 使用 MTP 头（Multi-Token Prediction Head）而非完整的语言模型
• E2B/E4B 的 drafter 还使用了 Centroids Masking 技术，将词表空间从约 262K 压缩到约 4K 个候选 token，极大加速了 logit 计算
• 与主模型共享 KV Cache，进一步减少显存开销

Gemma 4 MTP Drafter 模型一览

Google 为 Gemma 4 全部四个主模型都发布了对应的 drafter 模型：

命名规则很简单：在主模型 ID 后面加上 -assistant 后缀就是对应的 drafter 模型。

⚠️ 注意：drafter 模型不能独立使用，它必须与对应的主模型配对才能工作。不要尝试单独加载 drafter 模型进行推理。

实战：Hugging Face Transformers 中使用 MTP

最直接的方式是通过 Hugging Face Transformers 库使用 MTP。以 Gemma 4 E4B 为例：

环境准备

pip install -U transformers torch accelerate

基础用法

from transformers import AutoProcessor, AutoModelForCausalLM

TARGET_MODEL_ID = "google/gemma-4-E4B-it"
ASSISTANT_MODEL_ID = "google/gemma-4-E4B-it-assistant"

# 加载主模型
processor = AutoProcessor.from_pretrained(TARGET_MODEL_ID)
target_model = AutoModelForCausalLM.from_pretrained(
    TARGET_MODEL_ID,
    torch_dtype="auto",
    device_map="auto",
)

# 加载 MTP drafter 模型
assistant_model = AutoModelForCausalLM.from_pretrained(
    ASSISTANT_MODEL_ID,
    torch_dtype="auto",
    device_map="auto",
)

# 构造对话
messages = [
    {"role": "system", "content": "你是一个有帮助的助手。"},
    {"role": "user", "content": "用 Python 写一个快速排序算法。"},
]

text = processor.apply_chat_template(
    messages,
    tokenize=False,
    add_generation_prompt=True,
)
inputs = processor(text=text, return_tensors="pt").to(target_model.device)
input_len = inputs["input_ids"].shape[-1]

# 关键：通过 assistant_model 参数启用 MTP
outputs = target_model.generate(
    **inputs,
    assistant_model=assistant_model,
    max_new_tokens=512,
)

response = processor.decode(outputs[0][input_len:], skip_special_tokens=True)
print(response)

核心就是 generate() 方法中的 assistant_model 参数——传入 drafter 模型即可自动启用 speculative decoding。

调节草拟 token 数量

Transformers 提供了精细控制草拟 token 数量的参数：

outputs = target_model.generate(
    **inputs,
    assistant_model=assistant_model,
    max_new_tokens=512,
    # 固定每次草拟 4 个 token
    num_assistant_tokens=4,
    num_assistant_tokens_schedule="constant",
)

num_assistant_tokens_schedule 支持两种模式：

• "heuristic"（默认）：根据历史接受率动态调整草拟 token 数量
• "constant"：固定数量，适合已知任务特征的场景

实战：vLLM 中部署 Gemma 4 MTP 服务

vLLM 是目前最流行的 LLM 推理框架之一，已原生支持 Gemma 4 MTP。

启动推理服务

vllm serve google/gemma-4-31B-it \
  --tensor-parallel-size 2 \
  --max-model-len 8192 \
  --speculative-config '{
    "model": "google/gemma-4-31B-it-assistant",
    "num_speculative_tokens": 4
  }'

核心参数说明：

• --tensor-parallel-size 2：使用 2 张 GPU 进行张量并行
• --max-model-len 8192：最大上下文长度
• --speculative-config：指定 drafter 模型和草拟 token 数量

Python 离线推理

from vllm import LLM, SamplingParams

llm = LLM(
    model="google/gemma-4-E4B-it",
    speculative_config={
        "model": "google/gemma-4-E4B-it-assistant",
        "num_speculative_tokens": 4,
    },
    max_model_len=8192,
    trust_remote_code=True,
)

sampling_params = SamplingParams(
    temperature=0.7,
    top_p=0.95,
    top_k=64,
    max_tokens=1024,
)

prompts = [
    "解释一下 Transformer 架构中的自注意力机制。",
    "用 Rust 实现一个简单的 HTTP 服务器。",
]

outputs = llm.generate(prompts, sampling_params)
for output in outputs:
    print(f"Prompt: {output.prompt[:50]}...")
    print(f"Generated: {output.outputs[0].text[:200]}...")
    print("---")

调用 OpenAI 兼容 API

vLLM 服务启动后，可以直接使用 OpenAI SDK 调用：

from openai import OpenAI

client = OpenAI(
    base_url="http://localhost:8000/v1",
    api_key="EMPTY",
)

response = client.chat.completions.create(
    model="google/gemma-4-31B-it",
    messages=[
        {"role": "user", "content": "什么是 speculative decoding？"},
    ],
    max_tokens=512,
    temperature=0.7,
)

print(response.choices[0].message.content)

实战：SGLang 中部署 Gemma 4 MTP 服务

SGLang 同样在第一时间支持了 Gemma 4 MTP drafter。

启动推理服务

SGLang 使用 --speculative-algorithm 和 --speculative-draft-model-path 参数启用 MTP：

python -m sglang.launch_server \
  --model-path google/gemma-4-31B-it \
  --speculative-algorithm NEXTN \
  --speculative-draft-model-path google/gemma-4-31B-it-assistant \
  --speculative-num-draft-tokens 4 \
  --tp-size 2 \
  --port 30000

参数说明：

• --speculative-algorithm NEXTN：使用 MTP（Next-N Token Prediction）算法
• --speculative-draft-model-path：drafter 模型路径
• --speculative-num-draft-tokens 4：每次草拟的 token 数量

调用 API

SGLang 同样提供 OpenAI 兼容接口：

from openai import OpenAI

client = OpenAI(
    base_url="http://localhost:30000/v1",
    api_key="EMPTY",
)

response = client.chat.completions.create(
    model="google/gemma-4-31B-it",
    messages=[
        {"role": "system", "content": "你是一位资深的 Python 开发者。"},
        {"role": "user", "content": "如何优化 pandas DataFrame 的内存使用？"},
    ],
    max_tokens=1024,
    temperature=0.7,
    stream=True,
)

for chunk in response:
    if chunk.choices[0].delta.content:
        print(chunk.choices[0].delta.content, end="", flush=True)

性能基准测试

不同模型的加速效果

以下是在 2 张 NVIDIA A100 80GB GPU 上的典型性能数据（输入 1000 tokens，输出 1000 tokens）：

💡 趋势：模型越大，MTP 带来的加速比越明显。这是因为大模型的前向传播开销更大，“空闲算力”更多，drafter 草拟的 token 能更有效地利用这些算力。

不同任务类型的效果差异

MTP 的加速效果与任务的”可预测性”密切相关：

代码生成和翻译场景效果最好，因为这些任务中 token 序列的确定性较高（比如代码中的关键字、语法结构），drafter 模型能准确预测大部分 token。

vLLM 基准测试命令

使用 vLLM 内置的基准测试工具进行性能对比：

# 不启用 MTP 的基准
vllm serve google/gemma-4-31B-it \
  --tensor-parallel-size 2 \
  --max-model-len 32768

vllm bench serve \
  --model google/gemma-4-31B-it \
  --dataset-name random \
  --random-input-len 1000 \
  --random-output-len 1000 \
  --num-prompts 16

# 启用 MTP 的基准
vllm serve google/gemma-4-31B-it \
  --tensor-parallel-size 2 \
  --max-model-len 32768 \
  --speculative-config '{
    "model": "google/gemma-4-31B-it-assistant",
    "num_speculative_tokens": 4
  }'

vllm bench serve \
  --model google/gemma-4-31B-it \
  --dataset-name random \
  --random-input-len 1000 \
  --random-output-len 1000 \
  --num-prompts 16

最佳实践与调优建议

1. 选择合适的 num_speculative_tokens

草拟 token 数量不是越大越好。设得太大，drafter 模型的开销增加，且后面的 token 被拒绝的概率更高，反而会拖慢速度。

• E2B/E4B 模型：推荐设为 2-4
• 26B-A4B / 31B 模型：推荐设为 4-8
• 代码生成场景：可以适当增大到 6-8
• 创意写作场景：建议设为 2-3

2. 批处理场景的权衡

在高并发批处理场景下，MTP 带来的加速比可能会下降。因为 drafter 模型的计算和主模型的验证都会竞争 GPU 资源。建议：

• 低并发、低延迟优先：开启 MTP，加速效果显著
• 高并发、高吞吐优先：测试对比有无 MTP 的吞吐量，按实际结果选择

3. 显存规划

以 31B 模型为例，加载主模型和 drafter 后的显存占用参考：

主模型 (BF16): ~62 GB
MTP Drafter:   ~1 GB
KV Cache:      视 max_model_len 而定
---
总计:          ~63 GB + KV Cache

drafter 模型仅增加约 1-2% 的显存开销。在使用 --tensor-parallel-size 2 时，每张 GPU 分摊约 32 GB。

4. 量化与 MTP 组合

MTP 可以与量化技术叠加使用。例如使用 FP8 量化后：

vllm serve google/gemma-4-31B-it \
  --tensor-parallel-size 1 \
  --quantization fp8 \
  --max-model-len 8192 \
  --speculative-config '{
    "model": "google/gemma-4-31B-it-assistant",
    "num_speculative_tokens": 4
  }'

FP8 量化 + MTP 的组合可以让 31B 模型在单张 24GB GPU 上运行，同时享受加速效果。

5. 生产环境推荐配置

# 生产环境推荐的 vLLM 启动命令
vllm serve google/gemma-4-31B-it \
  --tensor-parallel-size 2 \
  --max-model-len 32768 \
  --speculative-config '{
    "model": "google/gemma-4-31B-it-assistant",
    "num_speculative_tokens": 4
  }' \
  --enable-prefix-caching \
  --max-num-seqs 64 \
  --port 8000

--enable-prefix-caching 会缓存共享的前缀 KV Cache，对多轮对话场景有额外加速效果。

Gemma 4 的更多亮点

除了 MTP 加速，Gemma 4 本身也值得关注：

• 256K 上下文窗口：31B 模型支持最长 256K tokens 的上下文
• 多模态能力：原生支持文本和图像输入
• 140+ 语言支持：涵盖主流和小语种
• MoE 架构：26B-A4B 采用混合专家架构，激活参数仅 4B，推理更高效
• 推理模式：内置 step-by-step 思维链推理
• 函数调用：原生支持结构化工具调用

总结

Gemma 4 MTP 代表了开源模型推理加速的一个重要方向：不修改模型架构、不牺牲生成质量，通过专属 drafter 模型实现显著的速度提升。

核心要点回顾：

1. MTP drafter 是配套组件，不是独立模型。它与主模型共享 tokenizer、KV Cache，接受率远高于传统的独立小模型方案
2. 加速效果取决于任务。代码生成、翻译等高确定性任务可达 2.5-3x；创意写作等低确定性任务约 1.5-2x
3. 部署成本极低。drafter 模型仅增加 1-2% 的显存开销，vLLM 和 SGLang 都只需要一行配置即可启用
4. 可与量化技术叠加。FP8 + MTP 让 31B 模型在单张消费级 GPU 上也能流畅运行

对于正在部署 Gemma 4 的团队，MTP 是一个几乎”无脑启用”的优化选项——配置简单、零质量损失、显著提速。赶紧试试吧。