SpecFold：折叠多分支冗余以实现扩散语言模型中更快的推测解码

摘要

扩散大语言模型 (DLLM) 通过迭代块去噪生成文本，多分支推测解码通过在单次前向传递中验证主分支和多个草稿分支来加速此过程。虽然之前的 DLLM加速方法主要利用去噪步骤中的时间冗余，但我们在每个推测的验证步骤中确定了一个互补的冗余轴：多分支计算冗余。在推测性验证期间，草案分支从其父级继承大部分token，同时揭露一小部分附加位置，导致大部分隐藏状态在分支之间保持高度相似。我们提出了 SpecFold，一种算法系统协同设计，利用这种多分支冗余来降低多分支推测验证的成本。在算法上，SpecFold 执行token级别的残差门控，并通过折叠注意力和 FFN 选择性地重用父计算，同时保留残差隐藏状态。系统地，Triton 内核实现通过高效的稀疏多分支执行将这种细粒度的重用转化为端到端吞吐量增益。 SpecFold 与时间缓存正交，并与现有的 DLLM推测策略兼容。在两个 DLLM系列、五个模型和五个标准基准测试中，SpecFold 的吞吐量比 Spiffy 最高达到1.64倍，比普通解码最高达到1.99倍，同时保持了相当的任务性能。

SpecFold：折叠多分支冗余以实现扩散语言模型中更快的推测解码的原论文方法或结果图 — 图 1：SpecFold 概述。 (a) 标准 DLLM解码通过重复去噪前向传播迭代地揭开token的屏蔽。 (b) 多分支推测 DLLM解码在一次批量前向传递中验证主分支以及多个草稿分支。如果草稿与后续的主分支状态匹配，则其预先计算的逻辑将被重用以跳过未来的去噪前向传递。 (c) SpecFold 在每个验证步骤中利用跨分支冗余。在每个 Transformer 层之前，父子残差门会识别折叠位置，从而重用父计算，而发散位置则通过折叠注意力和 FFN 执行重新计算，从而提高解码吞吐量。