1. 项目概述为什么我们需要加速 Gemma 4 的推理最近在折腾大语言模型推理部署的朋友估计都绕不开一个词延迟。模型能力越来越强参数越来越多但用户对响应速度的期待却丝毫没有降低。尤其是在线对话、代码生成、实时翻译这些场景慢个一两秒体验感就直线下降。我最近就在深度优化 Gemma 4 的推理性能发现了一个非常有意思且效果显著的思路利用多令牌预测Multi-Token Prediction的草稿模型Drafter来加速推测解码Speculative Decoding。简单来说这就像给一个深思熟虑的“大师”主模型配了一个思维敏捷的“助手”草稿模型。助手负责快速“预判”大师接下来可能会说的好几个词多令牌然后大师一次性审阅这些预判快速“批准”正确的部分从而跳过大量重复的计算。这和我们之前常见的、只能预测下一个词的草稿模型相比效率提升是数量级的。如果你正在为 Gemma 4 的推理速度发愁或者对推测解码的潜力感兴趣这篇从原理到实操、再到避坑的完整经验分享应该能给你带来不少启发。2. 核心原理拆解多令牌预测草稿模型如何工作要理解这个加速方案我们得先拆解几个核心概念推测解码、草稿模型以及本次的重点——多令牌预测。2.1 推测解码让大模型“抄近道”推测解码的核心思想是“验证比生成便宜”。让一个速度快、成本低的小模型草稿模型先走一步生成一段连续的令牌序列草稿然后让大模型主模型并行地对这段草稿进行验证。验证时大模型不是逐个生成而是并行计算每个位置在给定历史上下文下的下一个令牌概率分布。如果草稿的某个令牌与大模型“心中所想”的概率分布峰值一致就接受它一旦出现分歧就拒绝该令牌及之后的所有草稿用大模型自己生成的令牌替换分歧点然后继续。这个过程的关键在于大模型并行验证 K 个令牌的计算开销远小于它自己串行生成 K 个令牌。只要草稿模型的准确率足够高我们就能用一次大模型前向传播“兑换”出多个被接受的输出令牌从而大幅降低整体延迟。2.2 从单步到多步多令牌预测草稿模型的革新传统的草稿模型通常被训练成只预测下一个令牌Next-Token Prediction这和大多数LLM的训练目标一致。但在推测解码的上下文中这种单步预测存在局限它每次只能提供一个候选容错率低且无法充分利用草稿模型自身的序列生成能力。多令牌预测草稿模型则不同。它在训练时目标就是同时预测未来多个时间步的令牌。例如一个具有n_predict4能力的草稿模型其训练目标是基于当前上下文x同时预测(y_t, y_{t1}, y_{t2}, y_{t3})。这通常通过修改模型输出头来实现最后的线性层不再输出一个|V|维词表大小的向量而是输出n_predict个|V|维的向量每个对应一个未来位置的概率分布。这样做带来的巨大优势是更高的吞吐潜力一次前向传播直接产出多个连续的草稿令牌而不是一个。这显著减少了草稿阶段本身的耗时。更好的序列一致性由于是在一次前向传播中共同预测多个令牌这些预测基于完全相同的隐藏状态它们之间的内在一致性理论上比自回归地一个个生成更好。更匹配验证阶段主模型的验证本身就是并行的。一个能并行产出多令牌草稿的模型与验证过程在形式上更加对齐。2.3 系统工作流程全景结合了多令牌预测草稿模型的推测解码其单轮工作流程如下草稿阶段给定当前上下文多令牌预测草稿模型运行一次前向传播直接生成一个长度为γ例如 4 或 5的候选令牌序列[d1, d2, ..., dγ]。验证阶段主模型Gemma 4以同样的上下文为输入并行运行一次前向传播得到对于未来γ个位置的概率分布[P1, P2, ..., Pγ]。接受/拒绝决策从第一个位置开始比较草稿令牌d1是否等于主模型分布P1中概率最高的令牌即贪心解码下的选择。如果相等则接受d1并继续比较d2与P2的峰值令牌依此类推。如果在第r个位置出现分歧dr不等于Pr的峰值则拒绝dr及之后所有草稿。接受前r-1个令牌。用主模型在位置r的分布Pr中采样或取贪心得到的新令牌替换被拒绝的dr。更新与循环将本轮接受的所有令牌长度为r添加到已生成序列中并以新的序列作为上下文开始下一轮循环。注意这里描述的是“贪心验证”是最常见的方式。也可以进行随机采样验证但复杂度更高。多令牌预测的优势在于它极大增加了单轮循环中可能被接受的令牌数量上限。3. 模型训练与适配打造高效的“助手”要让这个方案落地首要任务就是获得一个高质量的多令牌预测草稿模型。直接拿一个现成的单步预测模型来用是不行的必须进行针对性的训练或适配。3.1 训练数据与目标构造训练数据来自主模型Gemma 4的预训练数据或高质量通用语料。关键在于如何构造训练目标。对于一个训练样本文本序列我们将其分割成多个重叠的上下文-目标对。例如对于序列[x1, x2, x3, x4, x5, x6, x7]设定n_predict3输入上下文[x1, x2, x3, x4] 目标[x5, x6, x7]滑动窗口后下一个样本可能是输入[x2, x3, x4, x5] 目标[x6, x7, x8]模型的结构需要调整。一种典型做法是骨干网络使用与主模型架构相似但层数少得多的小模型例如Gemma 4 有 60层草稿模型可能只用 4-8 层。直接复用主模型的词表嵌入层通常是个好主意可以保证词表空间一致。输出头这是关键。不再使用单一的LM Head而是使用n_predict个独立的线性层或一个共享权重但输出维度为n_predict * |V|的大层分别对应预测第1个、第2个…第n_predict个未来令牌。损失函数计算每个预测位置的交叉熵损失然后求和或平均。L (CE_loss(pos1) CE_loss(pos2) ... CE_loss(pos_n)) / n。3.2 知识蒸馏向“大师”学习单纯用文本数据训练小模型预测多个令牌效果可能有限。更强大的方法是引入知识蒸馏Knowledge Distillation让草稿模型直接学习主模型的“思考方式”。具体来说在构造训练数据时我们不仅提供真实的后续令牌作为硬目标更重要的是利用主模型Gemma 4来提供“软目标”。对于每个训练上下文让主模型运行一次前向传播得到未来n_predict个位置完整的概率分布而不仅仅是峰值令牌。这些分布包含了丰富的知识比如同义词的概率、语法结构的偏好等。训练草稿模型时损失函数由两部分组成硬目标损失与真实令牌的交叉熵确保基础准确性。软目标损失草稿模型输出的概率分布与主模型概率分布之间的KL散度Kullback–Leibler divergence。这迫使草稿模型模仿主模型的输出分布形态。通常在训练初期更侧重软目标损失让草稿模型快速学会主模型的风格训练后期则更侧重硬目标损失收敛到更精确的预测。实操心得蒸馏的温度参数Temperature设置很重要。温度 1 可以平滑主模型的分布让概率较小的令牌信息更明显便于草稿模型学习更丰富的知识。我们通常在训练早期使用较高的温度如 2.0后期逐渐降低到 1.0。3.3 架构选择与尺寸权衡草稿模型不需要强大的推理和知识能力它只需要在“给定上下文后接下来最可能是什么”这个任务上与主模型保持高度一致。因此架构优先选择与主模型相同的架构家族如都是Transformer Decoder这能最大程度保证行为一致性。用一个小号的 Gemma如 Gemma 2B作为草稿模型的起点是常见选择。尺寸这是一个权衡。模型越大预测准确率通常越高能带来更高的接受率。但模型越大其单次前向传播的耗时也越长可能会抵消加速收益。经验上草稿模型的参数量应为主模型的 1/10 到 1/100并且其单次推理延迟必须远低于主模型单次生成一个令牌的延迟。对于 Gemma 4可能为百亿参数级别一个几亿参数的草稿模型往往是合适的起点。层数 vs. 宽度在总参数量一定的情况下相对更“宽”隐藏层维度大而不是更“深”层数多的模型有时在并行预测任务上表现更好因为前向传播更浅延迟更低。4. 推理引擎集成与优化实战有了训练好的草稿模型下一步就是将其集成到推理引擎中并与主模型Gemma 4协同工作。这里以目前生态支持较好的 vLLM 和 TensorRT-LLM 为例讲解集成思路和优化要点。4.1 推理流程的工程实现整个推测解码流程需要在推理引擎层面实现主要控制逻辑如下# 伪代码展示核心循环逻辑 def speculative_decoding_with_multi_token_drafter(main_model, drafter_model, prompt, max_tokens, n_predict): generated_tokens [] current_context prompt while len(generated_tokens) max_tokens: # 1. 草稿阶段多令牌预测 draft_tokens drafter_model.generate_multi_token(current_context, n_predict) # 一次前向得n_predict个令牌 # 2. 验证阶段并行前向 # 将当前上下文与草稿令牌拼接作为主模型的输入 verification_input concat(current_context, draft_tokens) # 主模型一次前向传播得到对草稿每个位置的概率分布 main_probs main_model.forward_parallel(verification_input, len(draft_tokens)) # 3. 接受/拒绝逻辑 accepted_tokens [] for i, (draft_token, token_probs) in enumerate(zip(draft_tokens, main_probs)): top_token argmax(token_probs) # 贪心解码 if draft_token top_token: accepted_tokens.append(draft_token) else: # 发生分歧 replacement_token sample_or_greedy(token_probs) # 从主模型分布中取新令牌 accepted_tokens.append(replacement_token) break # 只替换分歧点后续草稿丢弃 # 4. 更新 generated_tokens.extend(accepted_tokens) current_context concat(prompt, generated_tokens) # 更新上下文继续循环 return generated_tokens4.2 与 vLLM 集成vLLM 通过其SpeculativeProposer接口原生支持推测解码。我们需要实现一个支持多令牌预测的Proposer。创建自定义 Proposer继承vllm.spec_decode.proposer.SpeculativeProposer。实现get_spec_proposals方法在这个方法中调用我们的多令牌预测草稿模型。这里的关键是vLLM 期望的proposal_len就是我们想要的n_predict。我们需要确保草稿模型的一次前向能返回恰好这个长度的令牌列表。批次处理vLLM 是高度批处理优化的。我们的草稿模型前向传播也必须支持批量输入即一次性为多个请求生成草稿。这要求草稿模型也能高效地进行批量推理。KV Cache 管理为了极致性能草稿模型也应支持 KV Cache避免每次都为整个上下文重新计算。由于草稿模型很浅其 KV Cache 的体积通常很小管理起来相对容易。配置示例片段from vllm import LLM, SamplingParams from my_custom_drafter import MultiTokenDrafterProposer # 初始化主模型和采样参数 llm LLM(modelgoogle/gemma-4, ...) sampling_params SamplingParams(temperature0.0, max_tokens512) # 贪心解码便于验证 # 创建并附加我们的多令牌草稿提议器 drafter_proposer MultiTokenDrafterProposer(drafter_model_pathpath/to/my_drafter, n_predict4) llm.set_speculative_proposer(drafter_proposer) # 正常使用 generate outputs llm.generate(prompts, sampling_params)4.3 与 TensorRT-LLM 集成TensorRT-LLM 在构建引擎时即支持推测解码。我们需要在构建阶段build阶段就将草稿模型的信息包含进去。构建草稿模型 TRT 引擎使用 TensorRT-LLM 的 API将我们训练好的多令牌预测草稿模型例如 PyTorch 格式编译成一个独立的 TensorRT 引擎drafter.engine。在编译时需要明确指定输出是n_predict个令牌的 logits。构建主模型 TRT 引擎编译 Gemma 4 模型在配置中启用推测解码选项并关联草稿模型引擎的路径。运行时加载在 TensorRT-LLM 的运行时 API 中加载主模型引擎时会自动识别并加载关联的草稿模型引擎。执行调用生成接口时推理过程会自动执行多令牌推测解码流程对用户透明。优势TensorRT-LLM 的集成在底层计算图级别进行融合优化可能获得比 Python 层调度更高的性能尤其是减少了主机与设备之间的数据传输开销。4.4 性能调优关键参数集成后需要通过基准测试来调优几个关键参数n_predict或gamma草稿模型每次预测的令牌数。这不是越大越好。增加n_predict能提高单轮吞吐上限但也会降低草稿的准确率预测越远越难导致拒绝率上升增加无效计算。需要找到一个平衡点通常对于 2-4B 的草稿模型配合百亿级主模型3-5 是一个常见的有效范围。草稿模型大小如前所述需要在准确率和延迟之间权衡。通过 A/B 测试测量不同大小草稿模型下的“平均接受长度”和“端到端生成延迟”来确定。批次大小Batch Size推测解码在批处理场景下收益更明显因为主模型并行验证的成本被均摊。但大批次也会增加内存压力和草稿模型的负担。需要根据 GPU 内存调整。主模型与草稿模型的执行重叠高级的优化会尝试让草稿模型为下一轮生成草稿时与主模型本轮验证过程在时间上部分重叠如果硬件资源允许以进一步隐藏延迟。5. 效果评估与典型问题排查部署完成后如何科学地评估加速效果遇到效果不理想的情况又该如何排查5.1 核心评估指标不要只看“快了多少”这种模糊感觉必须量化分析吞吐量Tokens/s在固定批次大小和生成长度下单位时间生成的令牌数。这是最直接的收益指标。延迟Time to First Token / Per-token Latency对于交互式应用首次令牌延迟TTFT和每个新令牌的延迟至关重要。推测解码通常能显著改善 TTFT。接受率Acceptance Rate平均每轮推测解码主模型接受草稿令牌的数量。理想情况是接近n_predict。计算公式总接受令牌数 / 推测解码轮数。这是衡量草稿模型质量的核心内部指标。加速比Speedup Ratio标准自回归解码耗时 / 推测解码耗时。注意要在相同硬件、相同配置如批次大小下对比。输出质量使用困惑度PPL或在特定任务如代码生成、问答上的评估分数确保加速没有损害生成文本的质量。5.2 效果不佳的常见原因与排查如果加速效果不明显甚至变慢了可以按照以下清单排查问题现象可能原因排查方法与解决方案接受率极低1.51. 草稿模型预测能力太差。2. 草稿模型与主模型词表/分词器不匹配。3.n_predict设置过大超出草稿模型能力。1.检查草稿模型质量单独测试草稿模型在验证集上的多令牌预测准确率Top-1 Acc。如果第一个令牌准确率就很低需要重新训练或使用更强的草稿模型。2.确认词表一致性确保两个模型使用完全相同的分词器。嵌入层是否对齐3.降低n_predict从 2 或 3 开始尝试逐步增加观察接受率变化。草稿阶段耗时过长1. 草稿模型本身太大或太深。2. 草稿模型未启用 KV Cache或实现低效。3. 批次处理开销大。1.分析 Profiling使用nsys或py-spy等工具分析时间主要消耗在草稿模型的哪一部分。2.优化草稿模型考虑量化INT8/FP8、使用更高效的注意力实现如 FlashAttention。3.调整批次大小对于小批次草稿模型开销占比可能变高可以尝试动态批次。加速比随序列长度下降推测解码的优势在生成初期最明显随着上下文变长每次验证都需要处理更长的序列计算量增加。这是固有特性。可以考虑动态调整n_predict在生成初期使用较大的值中后期减小。或者设定一个最大上下文窗口超过后回退到标准解码。输出质量下降1. 草稿模型引入了系统性偏差。2. 接受/拒绝逻辑在非贪心采样如 temperature0, top-p下有缺陷。1.质量评估定量计算 PPL 差异。如果下降明显检查草稿模型的训练数据是否有偏或尝试更强的知识蒸馏。2.采样验证如果使用采样验证逻辑需要比较随机采样的结果实现更复杂。确保你的验证逻辑与采样设置正确匹配。可以先在贪心解码temperature0下验证流程正确性。内存溢出OOM同时加载主模型和草稿模型以及它们的 KV Cache内存占用翻倍。1.使用量化模型为主模型和草稿模型应用 GPTQ/AWQ 等量化技术。2.优化 KV Cache使用 PagedAttentionvLLM等内存管理技术。3.降低批次大小或最大序列长度。5.3 一个真实的调优案例在我们针对代码补全场景优化 Gemma 4 时最初使用一个未经蒸馏的 1B 参数草稿模型n_predict5。结果发现接受率只有 1.8远低于预期。分析发现草稿模型在预测函数名、括号等代码特定结构时准确率尚可但在预测变量名、复杂表达式时很差。解决方案领域适应训练使用代码数据集如 GitHub Python代码对草稿模型进行继续预训练使其更熟悉代码分布。针对性蒸馏从 Gemma 4 在代码补全任务上的输出中进行蒸馏而不仅用通用文本。调整n_predict代码的局部连续性较强我们将n_predict降至 3接受率提升至 2.6整体吞吐量反而比n_predict5时提高了 40%。这个案例说明草稿模型与任务领域的匹配度至关重要通用草稿模型在特定任务上可能表现不佳。6. 进阶策略与未来展望在基础方案跑通之后还有一些进阶策略可以进一步压榨性能并值得思考未来的方向。6.1 动态推测与早停机制固定的n_predict可能不是最优的。我们可以让草稿模型在生成时动态决定要预测多少令牌。例如草稿模型可以输出一个“置信度”分数当置信度低于某个阈值时提前停止生成本轮草稿即使还没达到n_predict。这可以避免生成低质量的、注定被拒绝的后缀草稿节省验证开销。6.2 多草稿模型集成“三个臭皮匠顶个诸葛亮”的思路也可以应用在这里。同时使用多个不同架构或不同训练数据的轻量级草稿模型每个模型独立生成一份草稿。主模型并行验证所有这些草稿选择接受长度最长的那一条路径。这能有效提高单轮接受长度的期望值但代价是草稿阶段的计算量成倍增加需要精细权衡。6.3 草稿模型的持续学习与在线适配对于一个部署在特定应用如客服机器人、代码助手的模型其对话模式和代码风格是相对固定的。我们可以收集实际服务中主模型生成的日志在线微调Online Fine-tuning草稿模型让它越来越擅长预测主模型在该场景下的后续输出。这种持续学习的闭环能让加速效果随着时间越来越好。6.4 硬件感知的协同优化在芯片层面未来可能会有针对推测解码的硬件优化。例如NPU/GPU 是否可以提供一种“验证模式”让大模型核对小模型输出时只执行计算量更小的部分前向传播或者内存系统能否优化主模型与草稿模型之间 KV Cache 的交换这些硬件与软件的协同设计将是突破推理效率瓶颈的下一个前沿。从我实际的部署经验来看为 Gemma 4 这类大模型配备一个量身定制的多令牌预测草稿模型是目前性价比极高的推理加速方案。它不需要改变主模型本身不损害输出质量却能带来 2-3 倍甚至更高的吞吐提升。整个过程中最耗时的部分是草稿模型的训练与调优但一旦完成其收益是长期且稳定的。如果你正在面临大模型推理的成本或延迟压力我强烈建议你深入尝试这个方向它很可能成为你技术栈中的一个关键性能利器。