推测解码的树形验证：剪枝做错，加速比全吐回去

投机的性能模型一直很简单：小模型草拟 K 个 token，大模型一次性验证，接受前缀越长时间越省。一旦把这条线从一串拉直改成树形，收益和代价就不再沿同一条链走。SpecInfer（Miao 等人，收录于 ASPLOS 2024）把多个小模型的候选组织成 Token Tree，一次性验证整棵树并保留最长有效路径，公开报告加速比 1.5–3.5 倍。EAGLE-2 引入动态草稿树和剪枝，报告约 4 倍。数字好看，但树形验证的工程账，往往被算成“多生成几个候选”那么轻。

树形验证真正改变的是形状，不是数量

线性投机里，第 3 个候选是在第 2 个候选成立的前提下采出来的，接受率推导能干净地乘起来。树的每个节点是另一层意思：同一层上的多个兄弟节点代表互斥的未来，验证时必须把它们放进一次前向里做树注意力。跟线性的损失度规（loss）相比，后一个位置能看到的层数仍然受限于 ancestor。这意味着三件事：

• 树变成 n 个分支的节点之间的关系要对 batch 做调整，一个原本干净的序列就变成稀疏的树注意力掩码，形状是稀疏的对角矩阵按任意祖先顺序排列。
• 所有路径共享根节点之后共同祖先，被频繁访问的内核会为每个 batch 单独做 vLLM 这种 paged attention block table 的映射；树上增加的节点只需调整 block 偏移量即可，每条路径彼此独立。
• 验证同时批出来“候选”，不只是接受最长路径上的一小个前缀；平白多出来的枝丫在大部分计算中只是观战——事后剪掉并不会多花钱。

所以关于“剪枝”的第一条判断是：剪枝节省的是后续轮的草拟和验证开销，省不了当前这轮的验证。如果你把候选树修剪得接近线性，但这轮验证仍按 n 个 node 的树注意力算，就只能贴着“经典投机”本可以做到的加速比，值不回来为树建立的调度开销。

剪枝要在哪里剪，决定它省不省钱

剪枝能省点啥，取决于你插在哪一层：

草拟阶段剪枝

EAGLE-2 走的是这条路：动态草稿树不必为每个位置扩展同样多的 token，而是按上一轮各位置的激活状态决定哪些枝继续扩、哪些提前退出；超出后续验证预算就直接不结。这种按预算剪枝收益比较实在，因为草拟本来就要付款，扩散少付。代价是树的结构搜索本身也要时间，不能超过省下的验证开销。

验证前剪枝

不少实现选择在将树送进目标模型之前，按置信分数截断：某一层 node 自己路径下前面的 token 都没有通过大模型的对数概率筛选，继续往上长没必要。这一步能在准备阶段丢掉一轮看起来挺大的无用枝——但不要把数字乐观地加总到端到端时间。因为在 GPU 上，批量处理一个有掩蔽结构的 attention 的耗时，几乎和掩蔽区域本身无关；剪枝只减少了你给注意力掩码和 indexed block 传递的元素数，并不改写 kernel 的时间公式。实测收益更多来自减少 kernel 之间的同步和切换开销，而不是减少 masked token 本身计算出来的 FLOPs。

真正的问题是这样：树缩得越小，越容易变成单链再加了一点杂枝。收益也就回归到线性经典的区间。树的取值为四舍五入的工程红利：每层冒出多个子节点之后被目标模型一次性确认成功交叉扩展，比某个兄弟路径多成长几到几十个 token。按需剪什么，怎样让更多层同时推进的收益落在最短的那条路上？常常是拿 beta 参数先打到叶节点下做剪枝，收益不显著而非推不下去。

kv cache 与回退的隐性成本

中奖的验收规则挂在真实 Token 的版本上面。批量验证后你保留“最长有效路径”，其他的从中重来。在推理引擎（如 vLLM）上线投机解码时，难点在于接纳后回退时的缓存一致性：接纳后的节点要在 KV Cache 上按树结构 clone 共享, 否则下一次前向重复 prefill 会直接吃掉上一轮加速。多路径共用祖先，page 结构被祖先路径拆分时要重新映射 block 的 ref count。接受子树很长时对应前缀可以原地保持不变；如果只接受短前缀，剩下的长序列马上扔掉。一旦接受树里最长路径分散在 N 个非连续父节点上, ref count 和 block map 的调整会乱。能到达“动态复制”是在推进这种内存管理的流程之前，就要在上线前设计好。常见的一种保守实现是不会每轮动态分配几十个新增的 KV block，所以干脆不构造太深的树——最后加速也只能停留在 1.x。

参数与步长的调法

性能加速条件仍然是 αK > 1 + ρK：K 是草稿步长，α 是接受率，ρ 是小模型与大模型单步耗时比。树形把 K 从线性位数变成树叶子所代表的候选数量，但上述条件相当于整棵树的预算被更分散地拆成许多兄弟子节点,如果根节点之后有 β 个低分兄弟只吃掩码计算而没有回报，折算就把有产出树干的有效宽度拉低。工程上我先算“树上最密路径期望长度”和“开启树的版图计数”, 在接受的子树上运行接受率 α_tree 与同一配置上线性的对比。如果一条路线上没有任何一条非根节点上的兄弟树能给出准确重测的收益——即最短路径所兑换 token 数 = 覆盖下游上线的数量，就别保留树。

步长同样别一步到位。引用的性能模型提醒：ρ=0.2、α=0.6，步长 K=4 时相对 Tokens 耗时是 0.89，即只有 1.1x 多点；K=5 时降到 0.67，但有峰值之前的样本和低温下的接受率更能带来加速。实际部署中先 K=2～4 单支跑到饱和，再逐步长分叉并落实“实际 token/s 增益”，每档与基线放在相同 batch 与延迟 SLA 内快速 A/B。把接受率仅当作干指标，当作多变量指标时才不容易被美化。

草稿越像目标模型，接受率越高；这在小模型上最容易用蒸馏换整体温度低运行获得。但同分之下，树提供的价值高在三件事——命中用户实际路径的响应与候选序列，并在一次大模型调用里确认多个 token。这三件事没有一条适合手写而不用树实现增长，问题在于你养不养得起那棵树在每一轮里完成分发式复用。

结论落在取舍上

树形验证不是线性推测解码的平替开关，它是一场对固定候选数量的结构性重新安排。剪枝只在“后续轮”增加价值——草拟阶段节省以及预算扩散；到验证阶段能剪的面板小。若缓存克隆和回退接受值的路径没有打通，先行约束会比线性更紧。推进树形时，我会在测量上看三件事：每次候选切分下的有效端到端 token/s、KV 版本恢复耗时，与不同低温设定下的接受率衰减。反正不是在任何业务上这些都能对齐。