SpecFold:为扩散语言模型更快投机解码折叠多分支冗余

TL;DR SpecFold 利用多分支推测解码中的分支间冗余，通过折叠注意力/FFN 和 token 级残差门控，将扩散语言模型生成吞吐提升至 Spiffy 的 1.64 倍、vanilla decoding 的 1.99 倍。

问题

问题背景

扩散语言模型（DLLMs）通过迭代块去噪生成文本，推理步骤多、成本高。多分支推测解码（multi-branch speculative decoding）是当前主流加速手段：在单次前向中同时验证主分支与多个草稿分支，提升并行利用率。

现有方法局限

已有 DLLM 加速工作（如 Spiffy）主要利用跨去噪步骤的时间冗余（temporal caching），缓存逐步更新的 hidden states。但在每次推测验证步骤内部，草稿分支从父分支继承大部分 token，仅额外 unmask 少量位置，导致各分支的 hidden states 高度相似，却仍被 Transformer 各层完整重算。这种 多分支计算冗余 未被现有方法触及，造成验证成本高于理论下限，吞吐收益受限。

为什么这个问题难/重要

利用该冗余需要算法与系统协同：算法上需精确判断哪些 token 的 hidden states 可安全复用，并保持数值精度；系统上需用高效稀疏 kernel（如 Triton）实现细粒度门控与折叠注意力/FFN，避免 kernel 启动和内存搬运开销抵消收益。DLLM 推理效率是业界关注焦点，推测解码的验证阶段往往成为新瓶颈，任何降低验证成本的方法都能直接转化为端到端吞吐提升。

行业类比

类似于在图像扩散模型的逆向去噪过程中复用相邻步骤的特征图，SpecFold 在单步内折叠多分支的冗余计算，为基于 Transformer 的扩散模型推理提供了新的优化维度。

核心洞察

• **SpecFold** 识别出每步多分支投机验证内部的计算冗余：草稿分支继承父分支大部分 token，仅 unmask 少量位置，导致各分支隐藏状态高度相似。与以往仅利用跨去噪步时序冗余的方法不同，**SpecFold** 首次系统利用分支间冗余，通过 token 级残差门控与折叠 attention/FFN，跳过无需重算的位置。这一角度与时间缓存正交，可叠加现有 **DLLM** 投机策略。
• 细粒度重用必须配合高效稀疏执行才能转化为端到端收益。**SpecFold** 通过 Triton kernel 实现门控、压缩、源解析、折叠投影与 attention/FFN，避免朴素条件计算带来的 kernel launch 与内存开销。算法-系统协同设计证明：仅发现冗余不够，还需在 GPU 上显式管理稀疏多分支执行，才能将理论 FLOPs 节省变为实际吞吐提升（相比 **Spiffy** 最高 1.64x，相比 vanilla decoding 最高 1.99x）。

方法

输入与冗余来源

SpecFold 面向 多分支投机验证 (multi-branch speculative verification) 的单次前向步骤。输入包括一个主分支和多个草稿分支，每个草稿分支继承父分支的大部分 token，仅额外 unmask 少量位置。因此各分支的隐藏状态存在显著相似性，形成 multi-branch computational redundancy。

关键模块

1. token-level residual gating
   对每个 token 位置判断是否可以从父分支折叠计算：若该 token 的隐藏状态与父分支接近，则直接复用父分支的残差，避免在 attention 和 FFN 中重复计算。

2. folded attention
   拆分 query 为 recomputed query positions 与 folded query positions。折叠位置的 query 直接沿用父分支对应隐藏状态，仅对新增 unmask 的 token 重新计算注意力，从而降低 attention 的稀疏计算量。

3. folded FFN
   与折叠 attention 类似，对无需更新的 token 直接保留父分支的 FFN 结果，只对差异位置执行 FFN，进一步减少逐 token 的前馈开销。

4. Triton kernel 实现
   通过 gate and compaction、source resolution、projections、attention with state reuse、FFN and residual 等内核，在 GPU 上实现细粒度的稀疏多分支执行，将算法级折叠转化为端到端吞吐收益。

输出与效果

输出为每个分支的最终隐藏状态和预测 logits，供后续投机验证使用。SpecFold 保持与 vanilla decoding 相当的任务性能，同时与 Spiffy 等 temporal caching 方法正交，可叠加部署。

跟同类方法的差异点：SpecFold 聚焦单次验证步骤内的多分支计算冗余，而既有 DLLM 加速方法主要利用跨去噪步骤的时间冗余。

实验

实验设计

SpecFold 在两种扩散语言模型（DLLM）家族上验证，覆盖 Fast-dLLM-v2 与 Nemotron-Labs-Diffusion，共五个模型规模，在五个标准基准上评估，对比基线包括 Spiffy 与 vanilla decoding。评估指标为推理吞吐（tokens/s）与任务性能保持。

关键发现

SpecFold 通过折叠多分支冗余，在保持可比较任务性能的前提下，实现最高 1.64x 的 Spiffy 吞吐提升，以及 1.99x 的 vanilla decoding 加速。该方法与现有时间缓存正交，可叠加使用。

与基线对比解读

Spiffy 利用多分支推测解码，但未消除分支间的冗余计算；SpecFold 通过 token 级残差门控与折叠注意力 / FFN，精准复用父分支隐藏状态，将细粒度复用转化为端到端吞吐收益。工程上，Triton 内核的稀疏多分支执行保证了实际加速，避免了仅算法层面的理论收益。

行业影响

落地场景

SpecFold 针对 diffusion language models (DLLMs) 的多分支 speculative decoding 验证阶段,可显著降低单步计算冗余。在电商平台的智能客服与商品描述生成中,多分支草稿并行验证能更快产出候选标题、卖点与个性化回复;在代码助手或企业知识库生成场景,DLLM 的迭代去噪推理延迟被压缩,支持更长文本的实时补全。

商业价值

核心收益在降本:通过 token 级 residual gating 与折叠 attention / FFN,减少每次验证的 FLOPs 与显存带宽占用,同等 GPU 下吞吐最高提升 1.64×(对比 Spiffy)与 1.99×(对比 vanilla decoding)。按 token 计价的 API 业务可降低单请求成本,实时交互类产品则因 tail latency 下降带来更好的用户体验与留存。由于任务性能保持 comparable,加速不牺牲输出质量。

与现有栈的接口

SpecFold 以 Triton kernel 实现,可集成到 vLLM / TensorRT-LLM 等推理框架的 DLLM 分支中。它与 temporal caching 正交,兼容 Spiffy 等现有 speculation 策略,只需在验证阶段替换 attention 与 FFN 模块为 folded 版本,不改变训练/微调流程。部署时需针对不同 GPU 与 batch size 调优 Triton 算子;开源仓库 SpecFold 当前 star 较少,生产环境需补充稳定性与多硬件验证。

局限

• **通用性受限**：SpecFold 针对扩散语言模型（DLLMs）的多分支推测解码设计，核心假设是草稿分支继承父分支大部分 token、仅新增少量解掩码位置，导致跨分支隐藏状态高度相似。对于自回归 LLM 的树形推测解码或其他生成范式，冗余模式不同，方法难以直接迁移。实验仅覆盖两个 DLLM 家族、五个模型，缺乏对其他架构（如纯解码器扩散模型或混合模型）的评估，适用范围可能较窄。
• **与时间缓存未充分联合优化**：论文强调 SpecFold 与时间缓存正交，但未在实验中展示两者同时启用时的端到端加速上限。实际部署中，多分支折叠的稀疏计算与时间缓存的显存占用、调度策略可能存在冲突或边际收益递减，需要额外系统级协调。附录 C 提到的数值差异（与 vanilla decoding 的微小偏差）也可能在严格精度要求或长序列生成中累积放大，影响任务一致性。
• **工程落地依赖定制内核与硬件**：系统加速源于 Triton 稀疏多分支内核，其性能受 GPU 架构、内存带宽和 SM 利用率影响。论文未报告在消费级 GPU、非 NVIDIA 硬件或低精度（如 INT8）下的表现，也未讨论 batch size 扩展性。GitHub 仓库 stars 仅 3，社区验证不足，代码成熟度和易用性有待观察，可能阻碍快速集成到现有推理框架。