跳到主要内容
返回时间线
arXiv来源发表:

WaveFront Decoding 把循环语言模型的草稿与验证合并进同一次循环调用,在 Ouro-2.6B 与 Huginn-3.5B 上分别取得 2.42 倍和 3.54 倍解码加速

核心概要

该工作提出免训练的 WaveFront Decoding(WFD)自推测解码框架,利用循环语言模型中间递归输出可作草稿、权重共享可让不同位置与不同递归深度的状态在同一次批处理递归块调用中处理这两个性质,把混合深度状态组织成对角波前,使草稿与验证在同一批递归调用中并行进行,并用全深度预测纠正被拒绝的草稿;在 Spec-Bench 六个任务类别上相对自回归解码取得 Ouro-2.6B 2.42 倍、Huginn-3.5B 3.54 倍加速,结合跨递归 KV 共享后 Huginn-3.5B 加速最高达 4.81 倍。

AI-generated editorial illustration: WaveFront Decoding: Parallelized Self-Speculative Decoding for Looped Language Models

深度剖析

WFD 把循环语言模型的草稿与验证从两个分离阶段融合为一条连续流水线:新位置在浅深度起草,较早位置同时被推进到全深度验证,二者共享同一次批处理递归块调用。 此前工作已指出循环模型天然存在草稿—验证分解,并据此可实现两阶段的 draft-then-verify(DtV)基线;WFD 的增量在于消除阶段边界,让起草与验证在同一批调用中并发进行。 论文给出调度算法与波前状态示意,并在 Ouro-2.6B(full-stack 型)与 Huginn-3.5B(P/R/C 型)两类架构上验证,说明该调度不依赖特定架构族。

在 Spec-Bench 六个任务类别上,WFD 相对自回归解码取得 Ouro-2.6B 2.42 倍、Huginn-3.5B 3.54 倍的整体加速,并在每个任务上优于 DtV。 相对 DtV 整体提升约 27%,且在接受率较低时优势更明显,例如 Ouro 翻译任务上 DtV 在接受率 0.78 时仅 1.06 倍,而 WFD 仍保持 1.92 倍。 端到端吞吐在单张 NVIDIA RTX A6000、bf16、贪心解码、用户批大小为 1、生成上限 512 token、提示上限 1024 token 的设定下测得;任务级接受率介于 0.78 至 0.97,整体接受率 Ouro 为 0.92、Huginn 为 0.94。

论文给出解码成本模型,说明 WFD 的加速来自把串行递归块调用转为批处理调用,而非减少每个有效位置所需的递归计算,并指出 WFD 用有限波前宽度即可达到无限长 DtV 草稿的渐近每 token 延迟。 该分析把加速归因于内存带宽受限解码下权重读取的摊销,并解释 DtV 需要在验证摊销与拒绝代价之间权衡草稿块长度,而 WFD 的拒绝最多冲刷波前宽度内的在途位置。 分析基于 roofline 式算术强度估计与实测延迟分解;论文报告 Ouro-2.6B 仅达到 roofline 吞吐上限的 18%,其中 GPU 空闲占每 token 延迟的 65%,主要来自 kernel 启动开销,而 Huginn-3.5B 达到 74%。

跨递归 KV 共享可缓解 WFD 在长上下文下因混合深度位置访问不同递归 KV 槽而增加的 KV 流量,把 Huginn-3.5B 上的加速从 2.54 倍提升到 4.81 倍(GSM8K),MATH-500 上从 2.73 倍提升到 4.46 倍。 KV 共享并非 WFD 必需,但它改变了不同递归深度位置批处理的成本;论文将其作为互补手段单独评估,并指出该组合是近似扩展,因为混合深度位置可能观察到与顺序自回归不同的共享 KV 更新状态。 Huginn 被证明可容忍推理时 KV 共享;把 32 个递归 KV 槽减到 4 个或 1 个后,自回归准确率基本保持,WFD 接受率在 GSM8K 上从 0.88 升至最高 0.98、MATH-500 上从 0.90 升至 0.96;与相同 KV 共享配置下的自回归相比,WFD 准确率差距最多几个百分点。Ouro 未经跨递归 KV 共享训练,朴素事后共享无法保持其基线准确率,因此保留原有递归 KV 缓存组织。

启示与展望

该结果面向使用循环语言模型(full-stack 型如 Ouro、P/R/C 型如 Huginn)进行推理的服务场景,尤其是小批量、内存带宽受限的解码阶段;方法免训练、无需辅助草稿模型或权重修改,因此可直接叠加在已有检查点上。跨递归 KV 共享进一步把适用场景扩展到长上下文,论文报告在共享 KV 配置下可支持最长 64k token 上下文,而原始递归 KV 缓存在 Ouro 上 32k、Huginn 上 16k 即耗尽显存。论文把自适应草稿深度、树状草稿以及 P/R/C 三部分在专用硬件上的解耦列为后续方向,说明该方法被设计为可与这些调度策略组合。

评估限于单张 NVIDIA RTX A6000 与两个公开检查点,其他循环模型与其他硬件上的表现仍是开放问题。跨递归 KV 共享下的 WFD 被论文自身描述为近似扩展,因为混合深度位置可能在不同更新阶段观察到共享 KV 状态;fp32 下无 KV 共享时 WFD 与自回归产生完全相同的 token 序列,而共享 KV 时两者不再完全一致,尽管准确率差距最多几个百分点。bf16 下无 KV 共享时观察到的准确率差异被归因于批形状与归约顺序带来的有限精度效应,而非算法层面的输出质量变化。此外,Ouro 上自回归仅达到 roofline 吞吐上限的 18%,GPU 空闲占每 token 延迟的 65%,说明该模型的实测加速受 kernel 启动开销影响,CUDA graph 捕获后自回归吞吐提升 148%,此时 WFD 相对优化后自回归的加速为 2.30 倍,与 eager 模式下的 2.78 倍不同。自适应递归深度与 WFD 的交互也尚未实证:降低平均递归深度会收窄活跃波前,可能削弱 WFD 的相对加速,尽管绝对延迟可能改善。

来源