DominoTree：条件树草稿生成加速大模型推测解码

背景

大规模语言模型（LLM）的推理效率是关键瓶颈，自回归解码每次生成一个 token，计算成本随序列长度线性增长。推测性解码（Speculative Decoding）通过使用一个轻量级草稿模型（draft model）快速生成多个候选 token，再通过目标模型并行验证，从而在保持生成质量的同时显著加速推理。现有方法主要分为两类：块扩散（block-diffusion）起草器和最优优先树（best-first tree）方法。

• 块扩散起草器（如 DFlash）在一次前向传播中生成一个完整的草稿块（draft block），但每个位置只建模边缘分布（per-position marginal），忽略了 token 间的路径依赖。
• 最优优先树方法（如 DDTree）从这些边缘分布出发，通过扩展候选树来搜索更优的草稿路径，但树的结构基于分解（factorized）公式，无法表示 draft token 之间的条件依赖。

最近发布的 Domino 起草器引入了一种基于 GRU 的因果修正（causal correction），使得每个 draft token 的分布依赖于路径历史（path-dependent）。然而，DDTree 的分解公式无法捕捉这种结构。因此，需要一种新的树构建方法，既能利用 Domino 的条件修正能力，又保持高效构建。

核心内容

本文提出 DominoTree，一种无训练（training-free）的最优优先草稿树方法，用于推测性解码。DominoTree 的核心思想是：使用 Domino 的条件、非分解（non-factorized）修正沿着每一条根到节点的路径来评分候选节点，并通过限制每个节点的修正计算到候选 top-M 来保持可行性。

具体而言，DominoTree 的构建过程如下：

• 从目标模型（target model）的当前上下文出发，使用 Domino 的 GRU 修正器计算每个可能草稿 token 的条件概率分布（即路径依赖的分布）。
• 基于这些分布，以一种最优优先的方式扩展草稿树：每次选择当前树中评分最高的叶子节点，然后为该节点生成候选子节点（top-M token），并沿着该路径重新计算 Domino 的修正概率。
• 由于 Domino 的分布计算本身依赖于路径，因此不能像 DDTree 那样简单地复用边缘分布；DominoTree 通过在扩展时仅对 top-M 候选进行节点级修正计算，避免了全部枚举的高成本，使树构建在 CUDA 上高效实现。

DominoTree 的树构建使用 GPU 原生、CUDA-graph 的构建器，该构建器与参考 Python 实现保持比特一致（bit-identical），因此树的接受率（acceptance）与 Python 版本相同，而每轮树构建的开销极低。

实验评估在 Qwen3-4B 和 Qwen3-8B 模型上进行，覆盖八个基准测试。主要结果：

• Qwen3-4B：相比自回归解码，DominoTree 达到最高 6.6 倍加速；在所有测试温度下，平均接受长度（mean accept length）达到最高 10.7 tokens/轮，超过所有对比方法。
• 吞吐量对比：在默认构建器下，DominoTree 在所有温度下都优于 Domino 解码器：Qwen3-4B 上整体吞吐量提升 9-10%，Alpaca 数据集上高达 +22%；同时优于 DDTree 和 CaDDTree。
• Qwen3-8B：DominoTree 在所有温度下保持最高的接受长度；在温度 T=0 时，吞吐量比 DDTree 高出 +24%；在更高温度下，与 DDTree/CaDDTree 的优势缩小为平局或小幅损失，但整体仍优于 DFlash 和 Domino。

关键要点

1. 无训练、即插即用：DominoTree 不修改目标模型或草稿模型的权重，仅通过树结构优化草稿生成，可直接应用于现有推理框架。
2. 路径依赖建模：利用 Domino 的 GRU 因果修正，使每个 draft token 的分布取决于路径历史，突破了 DDTree 等分解方法的表达能力限制。
3. 高效的 top-M 修正策略：通过将每个节点的修正限制在候选 top-M，平衡了计算成本与树质量，使 GPU 原生实现可行。
4. 比特一致的 CUDA 图构建器：同时保证准确性和高性能，树构建的开销极低，且与 Python 参考实现接受率一致。
5. 显著的加速与吞吐量提升：在 Qwen3-4B 上获得高达 6.6 倍加速（相对自回归），在多个温度和模型尺寸下均超越当前最先进的草稿树方法。
6. 鲁棒的温度泛化性：在不同温度（低到高）下均能保持最优或接近最优的接受长度，尤其在低温（T=0）下吞吐优势明显。

意义与影响

DominoTree 是推测性解码领域的一项重要进展。它首次将路径依赖的条件分布（通过 Domino 的 GRU 修正）与最优优先树构建结合起来，验证了非分解分布的树草稿可以带来更高的接受长度和吞吐量。其意义体现在以下方面：

• 理论与方法的桥梁：DominoTree 解决了 DDTree 无法表示路径依赖的缺陷，将块扩散起草器的边缘分布与 Domino 的条件修正统一到树搜索框架中，为后续工作提供了新的范式。
• 实用高效：无训练属性和 GPU 原生实现使其易于集成到现有 LLM 推理系统（如 vLLM、TensorRT-LLM），有望降低实际部署中的推理延迟和成本。
• 基准性能标杆：在 Qwen3 系列上的全面评估表明，DominoTree 在多种场景下均优于现有的 DFlash、DDTree、CaDDTree 和 Domino 解码器，为社区提供了强有力的基线。
• 未来方向：该方法可能扩展到更大规模模型（如 70B+）以及多模态 LLM；同时，探索不同于 top-M 的剪枝策略或自适应树深度，可能进一步提升效率。

总体而言，DominoTree 在推测性解码中实现了 质量（接受长度）与效率（吞吐量）的双重提升，为 LLM 推理加速开辟了新的实用路径。