Weaver：因子化先验自回归草稿解码提速4.37倍

背景

自回归语言模型（如 GPT、Llama）在生成文本时，每一步只能输出一个 token，导致交互式应用（如聊天、代码补全）的延迟较高。推测解码（Speculative decoding）通过在一次前向传播中并行生成多个候选 token，以计算换时间，显著提升了交互性。其中，分解式草稿模型（factorized draft models）因其能并行预测未来 token 的边缘分布而格外高效，但其独立性假设导致随着推测预算（即一次生成的候选 token 数量）增大，接受率急剧下降。本文针对这一限制进行了深入分析，并提出了一种新的解决方案。

核心内容

论文首先指出，分解式草稿模型虽然并行效率高，但由于预测的每个未来 token 条件独立，当推测序列变长时，候选 token 之间的依赖缺失导致接受率（即被主模型接受的 token 比例）快速衰减。为了在保持并行性的同时恢复条件依赖，作者引入了 Weaver——一个轻量级的自回归适配器。Weaver 从分解式草案器的 top-K 边缘分布中构建提议树（proposal tree），而不是直接使用所有可能的 token。具体来说，Weaver 在每一步预测多个可能的 token 分支，形成一个树状结构，从而在避免全词汇投影（full-vocabulary projection）的前提下，恢复提议 token 之间的条件依赖关系。

此外，为了支持对具有 Gated Delta Net 层的模型进行快速验证，论文推导了一个无需回滚（rollback-free）的树验证算法，并在 SGLang 中实现了优化的 CUDA 内核。该算法允许在一次前向传播中同时验证整棵提议树，不需要像传统方法那样在验证失败时回退到某个位置重新开始。通过结合模型和系统层面的贡献，Weaver 在自回归解码的基础上实现了 4.37 倍的加速，并比高度优化的 DFlash 基线提升了 24.7%。

关键要点

• 分解式草稿模型因独立性假设，在推测预算增大时接受率严重下降，这是其核心瓶颈。
• Weaver 是一个轻量级自回归适配器，通过从分解式草案器的 top-K 边缘分布构建提议树，恢复 token 间的条件依赖。
• 提议树结构避免了全词汇投影，使 Weaver 保持高效，同时减少了推测长度增加带来的接受率损失。
• 针对 Gated Delta Net 层的模型，作者推导了无需回滚的树验证算法，并在 SGLang 中实现优化的 CUDA 内核，支持整棵树的并行验证。
• 整体方法在标准自回归解码上实现了 4.37 倍加速，相比 DFlash 基线有 24.7% 的额外提升。

意义与影响

本文的工作在推测解码领域具有重要意义。首先，它揭示了分解式草稿模型在长推测预算下的固有缺陷，并提出了一个实用且轻量的解决方案——通过构建提议树来恢复依赖关系，同时保持并行性。这为未来设计更高效的草案模型提供了新思路。其次，Weaver 的系统贡献（如树验证算法和 CUDA 内核优化）可以直接集成到现有的推理框架（如 SGLang）中，提升实际部署中自回归模型的响应速度。最后，4.37 倍的加速比表明，在保持生成质量的前提下，交互式语言模型可以显著降低延迟，这对于实时对话、代码生成等场景具有直接的应用价值。此外，该工作也可能启发在更广泛的序列生成任务中应用类似的自适应树结构草案方法。