投机解码的树状注意力验证：一次前向如何并行校验多条候选路径

在线大模型服务里，一个请求从进入队列到吐出最后一个 token，时间被拆成许多段。其中最难压缩的一段是自回归解码：每生成一个 token，都要把整个模型的权重从显存搬运到计算单元一次。GPU 的算力往往没被喂饱，真正卡住的是内存带宽。投机解码试图用“先猜后验”绕开这个瓶颈，但猜出来的候选越多，验证环节的开销也越大。树状注意力验证要解决的正是这个问题：让多条候选路径共享同一次前向。

假设有一个面向代码补全的在线服务，用户敲下函数签名后，模型需要续写若干行。草稿模型一次提出多个候选 token，这些候选不只是排成一条直线，而是分叉成一棵树。目标模型要在一次前向里判断树中每一条从根到叶的路径是否成立，并选出被接受的最长前缀。直觉上，验证 K 个候选需要 K 次前向；树状注意力把这件事压到一次。理解它为什么可行，需要先看清线性链式验证在哪里碰壁。

线性链式验证的瓶颈

最朴素的投机解码流程是：草稿模型自回归生成 K 个候选 token，记作 t1、t2、…、tK，目标模型把这 K 个 token 拼在已有上下文后面，一次前向计算每个位置的目标分布，然后从前往后逐个比对。一旦某个位置被拒绝，它之后的所有候选都作废。

这条链的问题在于，K 个候选被绑成一根绳子。只要中间某个 token 猜错，后面的全部浪费。草稿模型在语法常规、模板代码、常见短语上命中率高，但在需要具体事实或罕见标识符的位置容易偏离。代码补全场景里，函数名之后往往跟着固定的参数列表，这一段草稿很容易猜对；可一旦进入函数体内部的具体逻辑，草稿模型开始不确定，链式验证的接受长度就急剧缩短。

链式验证还有一个隐性成本：它把候选排成序列，目标模型看到的是“草稿模型认为最可能的那一条路径”。如果草稿模型在第 3 个位置拿不准，它仍然只能输出一个 token，目标模型没有机会同时评估其他可能。验证的并行度受限于单条路径，草稿模型的不确定性直接转化为接受率的下降。

树状候选如何组织

树状验证改变了候选的组织方式。草稿模型在每个位置不再只给一个 token，而是给出 top-2 或 top-3 的候选，这些候选按前缀关系连成一棵树。根是当前已确认的上下文，每个节点是一个候选 token，从根到某个节点的路径就是一条候选续写。

以代码补全为例，当前上下文是 def merge(a, b):，草稿模型可能在第 1 个位置给出 return 和 if，在第 2 个位置分别续上 a 和 b，形成分叉。树中不同分支共享前缀，return a 和 return b 共用 return 这个节点。共享前缀意味着目标模型对 return 的计算只需要做一次，后续分支在此基础上展开。

树的结构决定了验证的并行度。节点总数是验证批大小的上限，树的深度是理论最大接受长度。分支越多，覆盖草稿模型不确定位置的可能性越大；但节点数增加也会推高单次前向的计算量和 KV Cache 占用。工程上通常需要在分支因子和深度之间做取舍，常见做法是限制每个节点的子节点数，同时保留若干层深度。

自定义注意力掩码让路径互不干扰

标准因果注意力掩码是一个下三角矩阵：位置 i 只能看到位置 j ≤ i 的 token。树状验证需要更精细的规则——每个节点只能看到它的祖先节点，不能看到兄弟分支或其他分支的 token。如果掩码写错，return a 分支会错误地注意到 return b 分支的 token，验证结果就不再对应任何真实路径。

掩码的构造分两步。先把树中所有节点按某种顺序展开成一个序列，通常是按层或按深度优先排列。然后为每个节点确定它的祖先集合，在注意力矩阵中只在这些位置置为可见。对于 return a 这个节点，它能看到根、return、a，但看不到 b；对于 return b，它能看到根、return、b，看不到 a。

这个掩码在实现上通常表示为一个二维布尔矩阵，形状是“树节点数 × 树节点数”。目标模型在一次前向中处理整个节点序列，注意力计算按掩码屏蔽非法连接。由于不同分支的节点在同一序列中，KV Cache 也按这个序列布局：每个节点在每一层都有一个 key 和 value，兄弟节点之间不共享 KV，只有祖先关系才允许读取。

掩码正确性的验证方法很直接：对树中任意两个节点，检查它们的可见关系是否与树的前缀关系一致。如果某个节点能看到非祖先节点，说明掩码过宽，验证结果会混入错误信息；如果某个节点看不到某个祖先，说明掩码过窄，该路径的上下文不完整。生产环境中，掩码通常在构建树时一并生成，并作为调试信息保留，便于在异常时回溯。

flowchart TD
    A[已确认上下文] --> B[候选 return]
    A --> C[候选 if]
    B --> D[候选 a]
    B --> E[候选 b]
    C --> F[候选 a]
    C --> G[候选 b]
    D --> H[目标模型一次前向]
    E --> H
    F --> H
    G --> H
    H --> I[按掩码逐路径校验]
    I --> J[选出最长接受前缀]
    J --> K[更新 KV Cache 并输出]

图中从根到叶的四条路径在一次前向中同时进入目标模型，掩码保证每条路径只看到自己的祖先。校验后选出最长接受前缀，其余分支的 KV 被丢弃。

接受长度、验证批大小与加速比

加速比并不等于树中节点数。它取决于每次前向平均能接受多少个 token，以及这次前向本身比单 token 解码贵多少。

设树有 N 个节点，单次前向的计算量大致随 N 增长，但增长不是线性的：注意力部分与 N 的平方相关，前馈部分与 N 线性相关。当 N 较小时，前向时间主要由权重搬运决定，N 的增加带来的额外计算可能被掩盖；当 N 增大到一定程度，计算开始成为瓶颈，单次前向的耗时明显上升。

接受长度是另一个变量。树的分支覆盖了草稿模型的不确定位置，但目标模型最终只接受一条路径。如果草稿模型在某个位置的分支中没有包含正确 token，该位置之后的所有路径都会被截断。接受长度短，意味着这次前向的大部分计算被浪费；接受长度长，摊到每个 token 上的前向成本就低。

工程上通常观察三个量：每次前向的平均接受 token 数、单次前向的耗时、以及单位时间输出的 token 数。前两个量的比值决定加速比，第三个量决定服务吞吐。如果接受长度下降而验证批大小不变，说明草稿模型在当前输入分布上偏离较大，继续增大树宽只会增加计算浪费。

维度	线性链式验证	树状注意力验证
候选组织	单条路径，K 个 token	多分支树，N 个节点
验证前向次数	一次前向验证一条路径	一次前向验证整棵树
中间拒绝的代价	后续候选全部作废	仅该分支被截断，其他分支继续
掩码复杂度	标准因果掩码	自定义祖先可见掩码
KV Cache 布局	连续序列	按树节点展开，分支不共享
接受长度上限	受单路径正确性限制	受最长正确路径限制
单次前向开销	随 K 线性增长	随 N 增长，注意力部分更高
适用场景	草稿模型高度确定	草稿模型存在局部不确定

表格中的对比是定性描述。实际收益取决于模型、任务和草稿质量，不能脱离具体配置直接套用。

实现中的 KV Cache 与批处理约束

树状验证对 KV Cache 的管理提出了额外要求。标准解码中，KV Cache 按 token 顺序追加，每个新 token 的 key 和 value 接在序列末尾。树状验证中，多个分支的节点同时写入，Cache 的布局需要与树的展开顺序一致。

一种常见做法是为树分配一块连续的 Cache 区域，按节点编号存储。每个节点记录自己的父节点编号，注意力计算时通过掩码而非物理位置来限制可见性。校验完成后，被接受的路径上的节点保留，其余节点的 Cache 被标记为可回收。如果服务框架使用分页式 KV Cache，回收表现为页的释放或复用。

批处理是另一个约束。在线服务通常同时处理多个请求，每个请求有自己的树。不同请求的树大小可能不同，合并成一个大 batch 时，需要把各棵树拼在一起，并为每个请求单独构造掩码。这增加了调度复杂度：树大的请求占用更多计算资源，可能拖慢同批次的其他请求。工程上常见的做法是限制单棵树的节点数上限，并在调度时按树大小做分组，避免大小差异过大的请求互相干扰。

显存方面，树状验证的 KV Cache 占用与节点数成正比。节点数越多，单请求的显存占用越高，能同时服务的请求数就越少。这与吞吐量目标存在直接冲突，需要在树宽和并发数之间找平衡。

低接受率下的失效边界

树状验证并非在所有场景下都有效。当草稿模型与目标模型的分布差异较大时，树中分支可能全部偏离正确路径，接受长度退化到接近零。此时这次前向只确认了一个 token，甚至可能一个都没接受，计算全部浪费。

高温度采样是典型场景。温度升高后目标分布变得平坦，草稿模型的 top-2 或 top-3 候选很难覆盖目标模型真正会采样的 token。树的分支虽然增加了覆盖概率，但覆盖的是草稿模型认为可能的位置，而不是目标模型实际偏好的位置。接受率下降，加速比可能低于 1，即投机解码比直接解码还慢。

跨家族模型搭配也会触发类似问题。草稿模型和目标模型如果分词器不同或训练数据分布差异大，草稿的候选空间与目标模型的偏好空间错位，树的分支再宽也难以命中。

短输出任务同样不适用。如果请求只需要生成几十个 token，树构建、掩码生成和 Cache 分配的固定开销来不及被加速收益摊薄。首 token 延迟可能因为树初始化而增加，对延迟敏感的服务需要单独评估。

显存紧张的环境下，树状验证的额外 Cache 占用可能直接导致无法启动或频繁换出。此时增大树宽反而降低可用并发数，整体吞吐不升反降。

可观测指标与部署判断

部署树状验证时，需要观察几类信号来判断它是否在正常工作。接受长度是最直接的指标：如果平均接受长度持续低于某个阈值，说明草稿模型在当前流量上不可靠，应考虑降低树宽或更换草稿策略。

单次前向耗时与节点数的关系也需要监控。如果节点数增加后前向耗时增长超出预期，说明计算已成为瓶颈，继续加宽树不再划算。此时可以尝试减小树宽、增加并发，或改用更轻量的验证结构。

KV Cache 的占用和回收情况同样重要。如果被拒绝分支的 Cache 不能及时回收，显存会逐渐被无效节点占满，表现为可用并发数下降或请求排队变长。分页式 Cache 下，需要确认释放逻辑覆盖了所有未接受节点。

掩码正确性可以通过单元测试和线上抽样验证。构造一棵已知结构的树，检查目标模型对每个节点的注意力分布是否符合祖先可见规则。线上可以定期抽样，比对树状验证的输出与逐 token 解码的输出，确认分布一致性。

部署边界上，树状验证适合草稿模型与目标模型同家族、采样温度较低、输出长度较长的场景。代码补全、结构化生成、模板化文本是典型受益场景。高温度创意生成、跨家族搭配、极短输出和显存受限环境需要谨慎评估，必要时回退到线性链式验证或直接解码。

树状注意力验证把投机解码的并行度从“一条路径”扩展到“一棵树”，代价是掩码构造、Cache 布局和调度复杂度的上升。它没有消除草稿模型质量这个根本变量，只是让不确定位置上的多个可能同时进入验证。当草稿模型可靠时，它把接受长度推向树深；当草稿模型偏离时，它把浪费控制在单次前向之内。判断是否采用，最终取决于流量分布、显存预算和延迟目标三者之间的实际约束。