投机解码实践:用小模型做大模型的“猜词助手“,速度翻倍的技术路径 投机解码实践用小模型做大模型的猜词助手速度翻倍的技术路径一、自回归的诅咒为什么 70B 模型生成一个 Token 和 100 个 Token 一样慢大模型的自回归解码有个根本矛盾每生成一个 token必须遍历全部 Transformer 层。70B 模型有 80 层、每层 8,192 维的矩阵运算。生成 100 个 token 100 次完整前向传播——耗时约 3~5 秒A100。核心瓶颈不在计算量在计算密度。单 token 生成时 GPU 的 Tensor Core 利用率通常只有 15%~25%。矩阵维度是[1, 8192] × [8192, 8192]——batch_size1 的矩阵乘法严重浪费 GPU 的并行能力。这被称为memory-bound状态计算单元在等显存带宽供数。投机解码Speculative Decoding的思路是绕开这个矛盾用一个 10 倍小、10 倍快的草稿模型一次生成 K 个候选 tokendraft然后大模型一次前向传播验证这些 token。验证通过的就直接接受不通过的保留到第一个不匹配的 token后续丢弃。这样大模型一次前向传播可以推进 K 个 token——吞吐提升接近 K 倍接受率相关。某 70B 模型的实测原生自回归吞吐 32 tokens/s。引入一个 1.5B 草稿模型draft_len5投机解码吞吐 87 tokens/s——提升 2.7 倍。延时也从 3.1s 降至 1.2s100 token 输出。代价是多加载一个 1.5B 模型到 GPU 显存额外占用约 3GB。二、草稿生成与批量验证的流水线架构sequenceDiagram participant Draft as 草稿模型 (1.5B) participant Target as 目标模型 (70B) participant KV as KV Cache Note over Draft,KV: 阶段一草稿生成并行 K 步 Draft-Draft: 生成 token_1 (~1ms) Draft-Draft: 生成 token_2 (~1ms) Draft-Draft: ... Draft-Draft: 生成 token_K (~1ms) Draft-Target: 提交 K 个候选 token Note over Draft,KV: 阶段二批量验证一次前向传播 Target-Target: 并行验证 K 个 token (~50ms) Target-KV: 更新 KV Cache Note over Draft,KV: 阶段三接受/拒绝判决 Target-Target: 对比概率分布 alt 全部接受 Target-Draft: 继续下一轮K 个 token 全部生效 else 第 i 个被拒绝 Target-Target: 接受前 i-1 个 token Target-Draft: 从第 i 个位置重新生成草稿 end投机解码有三个关键组件草稿策略。最简单的方案是独立训练一个小模型作为草稿器。但更高效的方案是Medusa heads——在大模型最后一层附加多个预测头直接预测未来 K 个 token。Medusa 不需要独立模型仅增加约 0.1% 的参数却能达到独立草稿模型的 85%~95% 接受率。验证方式。不是简单的 token 匹配而是基于概率分布的接受/拒绝采样。大模型前向传播后对每个位置计算 token 概率分布。如果草稿模型对该位置的预测与大模型的非归一化概率之比 ≥ 某个随机采样阈值通常来自均匀分布 U(0,1)则接受否则拒绝并从该位置重新采样。KV Cache 管理。投机解码的 KV Cache 管理更复杂。草稿模型的 K 个 token 都有对应的 KV 状态被拒绝的 token 需要回滚对应的 KV Cache。vLLM 通过fork机制实现——在草稿开始时分叉 KV Cache 树验证后裁剪被拒绝的分支。三、投机解码的 Go 调度器实现package specdec import ( context math/rand sync ) // SpecDecScheduler 投机解码调度器 // 协调草稿模型和目标模型的批次调度 type SpecDecScheduler struct { draftModel DraftInferrer // 草稿模型接口 targetModel TargetVerifier // 目标模型接口 draftLen int // 草稿长度 K每次生成的候选 token 数 // 统计指标 mu sync.Mutex totalTokens int64 // 累计生成 token acceptedRatio float64 // 滑动平均接受率 totalRounds int64 // 投机解码轮次 } // DraftInferrer 草稿模型推理接口 type DraftInferrer interface { // GenerateDraft 生成 K 个草稿 token // 返回 token IDs 和对应的概率分布用于后续验证 GenerateDraft(ctx context.Context, prefix []int) ([]DraftToken, error) } // TargetVerifier 目标模型验证接口 type TargetVerifier interface { // VerifyAndAccept 并行验证 K 个候选 token // 返回接受的 token IDs、接受位置、更新后的 KV Cache handle VerifyAndAccept(ctx context.Context, prefix []int, drafts []DraftToken) ([]int, int, error) } // DraftToken 草稿 token 含概率 type DraftToken struct { ID int Probs []float32 // 该位置的完整概率分布用于拒绝采样 } // Run 执行投机解码循环 func (s *SpecDecScheduler) Run(ctx context.Context, prompt []int, maxNewTokens int) ([]int, error) { output : make([]int, 0, maxNewTokens) prefix : prompt for len(output) maxNewTokens { select { case -ctx.Done(): return output, ctx.Err() default: } // Step 1: 草稿模型生成 K 个候选 token drafts, err : s.draftModel.GenerateDraft(ctx, prefix) if err ! nil { return output, fmt.Errorf(draft generation failed: %w, err) } if len(drafts) 0 { break } // Step 2: 目标模型并行验证 accepted, acceptPos, err : s.targetModel.VerifyAndAccept(ctx, prefix, drafts) if err ! nil { return output, fmt.Errorf(verification failed: %w, err) } // Step 3: 收集接受的 token output append(output, accepted...) prefix accepted[len(accepted)-s.draftLen:] // 保留最后 draftLen 个 token 作为新的 prefix // 实际实现中 prefix 应保留完整上下文这里简化处理 // Step 4: 更新统计 s.updateStats(len(accepted), len(drafts)) // 提前退出如果验证到 EOS token if containsEOS(accepted) { break } } return output, nil } // updateStats 更新投机解码统计 func (s *SpecDecScheduler) updateStats(accepted, total int) { s.mu.Lock() defer s.mu.Unlock() s.totalRounds s.totalTokens int64(accepted) // 指数加权移动平均EWMAα0.1 平滑波动 ratio : float64(accepted) / float64(total) s.acceptedRatio s.acceptedRatio*0.9 ratio*0.1 } // RejectionSample 拒绝采样辅助函数 // 判断草稿 token 是否应被接受 func RejectionSample(draftProb, targetProb float32) bool { if targetProb draftProb { return true // 目标模型比草稿模型更确信直接接受 } // 按概率比接受 threshold : targetProb / draftProb return rand.Float32() threshold } // containsEOS 检查是否包含结束符 func containsEOS(tokens []int) bool { for _, t : range tokens { if t eosTokenID { return true } } return false } const eosTokenID 2 // Llama 系列 EOS token ID代码的关键设计点接口抽象DraftInferrer和TargetVerifier解耦草稿与验证的具体实现。可以替换为独立小模型、Medusa heads 或 eagle 等不同策略。指数加权统计acceptedRatio通过 EWMA 平滑避免单次波动影响监控判断。接受率低于 0.5 时应调低draftLen。逐轮 prefix 管理每次迭代后prefix 更新为最后 K 个 token含新接受的 token。简化版本——生产实现需要完整维护 token 序列。四、接受率波动、额外显存与退化风险接受率的不确定性。投机解码的加速比 (1 K × acceptance_rate) / (draft_cost_ratio 1)。当草稿模型与目标模型的输出分布高度一致时接受率 0.8加速接近 K 倍。但当任务切换如从 coding 切换到 poetry时接受率可能骤降至 0.3——此时开销大于收益。需要监控接受率指标低于阈值 0.4 时自动降级为普通自回归。额外的显存占用。独立草稿模型占用额外显存。1.5B 模型约 3GB对于 A100-40G 问题不大。但对于 A100-80G 部署 70BGQA 模型的场景70B 模型本身 140GB KV Cache ~30GB 草稿模型 3GB——已接近 173GB两张 A100 的总显存 160GB 是不够的。需要使用 Medusa0 额外模型显存或 Eagle仅需 ~100MB 额外参数。首 Token 延迟增加。投机解码在 prefill 之后才启动草稿-验证循环。这在 prefill 很长的场景下如 RAG 检索的长文首 Token 延迟主要由 prefill 贡献——投机解码的加速只对后续 token 有效。不适用场景模型部署在边缘设备显存不足以加载双模型草稿与目标任务风格差异过大如草稿训练于英文、任务为中文接受率低于 0.3批量推理场景大 batch 已利用 GPU 并行度投机解码的增益有限。五、总结投机解码将大模型解码从串行N 次前向传播 N 个 Token转变为并行K 次草稿 1 次验证 K 个 Token。核心收益取决于草稿接受率和草稿模型推理成本。在典型的 70B 1.5B 草稿配置下吞吐提升 2~3 倍是可实现的。落地路径优先选择 Medusa最高接受率/额外成本比或 Eagle次优但更简单设置draftLen初始值 3~5根据接受率动态调整同时在网关层实现自动降级——通过滚动窗口统计接受率低于阈值时退出投机模式。