← 返回 2026-06-02

Domino:在投机解码中将因果建模与自回归草稿解耦 Domino: Decoupling Causal Modeling from Autoregressive Drafting in Speculative Decoding

Jianuo Huang, Yaojie Zhang, Qituan Zhang, Hao Lin, Hanlin Xu, Linfeng Zhang 📅 2026-05-28 👍 152 2026-07-13 08:36
LLM推理加速 因果建模 并行草稿 投机解码

通过轻量级因果校正块,在保持低开销的同时恢复块并行草稿的因果依赖建模,实现最高5.49×的推理加速

前置知识

投机解码(Speculative Decoding)

投机解码是一种LLM推理加速技术,使用轻量级草稿模型提议多个未来token,然后由目标模型在单次前向传播中并行验证这些候选。验证过程接受最长有效前缀,并根据标准投机验证规则生成一个bonus token。平均每token延迟为 $L_{spec} = (T_{draft} + T_{verify}) / \tau$,其中 $T_{draft}$ 是生成草稿token的时间,$T_{verify}$ 是验证时间,$\tau$ 是每周期前进的预期token数。加速比为 $\eta = L_{target} / L_{spec} = \tau L_{target} / (T_{draft} + T_{verify})$。

投机解码是当前LLM推理加速的主流方法,理解其工作原理和加速比公式是理解本文权衡分析的基础。

自回归草稿(Autoregressive Drafting)

自回归草稿器按顺序生成候选token,每个token依赖于之前生成的所有token。给定前缀 $x_{\leq t}$,草稿分布被分解为 $q_{AR}(x_{t+1:t+\gamma} | x_{\leq t}) = \prod_{i=1}^{\gamma} q(x_{t+i} | x_{<t+i})$,其中 $x_{<t+i}$ 包括当前前缀和之前草稿的token。这种因子化镜像了目标模型的自回归预测过程,因此可以显式使用之前草稿的token来预测后续草稿位置,从而产生更高的草稿质量和更长的接受长度。然而,生成 $\gamma$ 个草稿token需要 $\gamma$ 个草稿步骤,总草稿成本近似为 $T_{draft}^{AR} \approx \gamma \cdot (t_{net} + t_{head})$,其中 $t_{net}$ 和 $t_{head}$ 分别是模型前向计算和LM头投影的平均延迟。

自回归草稿是EAGLE系列等现有方法的基础,理解其优势(高质量)和劣势(高成本)是理解本文动机的关键。

块并行草稿(Block-Parallel Drafting)

与自回归草稿不同,并行草稿器直接预测块级条件分布 $q_{PAR}(x_{t+1:t+\gamma} | x_{\leq t})$,从而并行生成多个token。其草稿成本近似为 $T_{draft}^{PAR} \approx t_{blocknet} + t_{blockhead}$,其中 $t_{blocknet}$ 和 $t_{blockhead}$ 分别是块级草稿模型前向传播和相应LM头投影的平均延迟。与自回归草稿不同,这些成本对整个草稿块只产生一次,而不是重复 $\gamma$ 次,从而通过并行计算实现更好的GPU利用率。然而,完全并行草稿削弱了草稿token之间的块内因果依赖,可能导致草稿分布与目标模型分布的错配。

块并行草稿是DFlash等并行方法的核心,理解其效率优势和质量劣势是理解本文创新的出发点。

接受长度(Acceptance Length)

接受长度 $\tau$ 表示在每次解码周期中被目标模型接受的平均token数,包括目标模型生成的bonus token。在投机解码中,目标模型在单次前向传播中评估草稿模型提议的 $\gamma$ 个候选token,并接受最长有效前缀。接受长度是影响加速比的关键因素之一:$\eta = \tau L_{target} / (T_{draft} + T_{verify})$。更高的接受长度意味着每次目标模型调用可以推进更多token,从而提高加速比。然而,获得高接受长度通常需要更高质量的草稿,这可能增加草稿成本。

接受长度是评估草稿质量的核心指标,理解其在加速比公式中的作用是理解本文权衡分析的基础。

研究动机

投机解码的实际加速受到草稿质量和草稿成本之间权衡的限制。具体来说,自回归草稿器(如EAGLE系列)通过顺序生成草稿token来建模草稿token之间的因果依赖,这与目标模型的自回归分布更好地对齐,产生长接受长度。例如,在Figure 1中,EAGLE-3在Qwen3-8B上达到了4.86的接受长度。然而,生成 $k$ 个草稿token需要 $k$ 个顺序草稿步骤,每个步骤涉及草稿模型前向传播和全词汇表LM头投影。这种成本随草稿长度线性增长,当草稿模型变深或词汇量变大时会进一步放大,可能抵消更长接受前缀带来的收益,限制通过扩展草稿长度或草稿容量获得的加速。另一方面,并行草稿器(如DFlash)通过单次前向传播生成整个草稿块,避免了重复的每token草稿模型和LM头调用,降低了草稿开销。DFlash在Figure 1中将加速从EAGLE-3的3.28×提高到3.42×。然而,去除顺序草稿依赖削弱了块内因果建模,可能降低草稿分布与目标模型分布的对齐,DFlash的接受长度从EAGLE-3的4.86降低到4.03。

本文的目标是本文的目标是提出一个轻量级因果校正框架,将因果依赖建模从昂贵的自回归草稿执行中解耦。具体而言,作者希望保留块并行生成的高效性,同时重新引入因果信息来恢复块内因果依赖,从而在不引入显著开销的情况下提高草稿质量和接受长度。这一目标体现在Figure 1的Domino结果中:相比DFlash,Domino只增加56M参数(+5.3%)和2.8%的总草稿-验证延迟,但将平均接受长度提高16.6%,端到端加速提高12.3%。

与已有工作不同的是,本文的独特切入角度是:可以保留因果依赖建模的草稿质量优势,同时保持块并行生成的低草稿成本。现有方法要么选择高质量的自回归草稿(但成本高),要么选择低成本的并行草稿(但质量受损)。Domino的关键洞察是:因果依赖建模不一定需要昂贵的自回归执行,可以通过轻量级校正分支在logit空间注入因果信息,而主草稿计算保持并行。这种设计避免了重复的草稿模型执行和昂贵的全LM头计算,同时恢复了有用的块内因果依赖。与在隐藏空间应用校正的方法不同,Domino的logit空间校正保持基础LM头计算并行,并将因果分支限制为低秩残差更新,从而在最小化开销的同时最大化草稿质量提升。

核心方法

Domino的整体架构包含两个组件:一个并行草稿主干和一个Domino头。并行草稿主干在单次并行计算中为草稿块中的所有位置生成初步分布,这保持了块并行生成的高效性。然后,Domino头通过在草稿位置之间传播因果信息来细化这些初步分布。这种设计允许Domino在保持并行主干效率的同时,通过轻量级因果校正恢复块内因果依赖,从而解决草稿质量和草稿成本之间的权衡。在推理时,Domino首先使用并行主干生成基础logits,然后使用Domino头基于之前采样的草稿token生成校正logits,最终从细化后的logits中采样每个草稿token。

Domino的核心创新在于将因果依赖建模从昂贵的自回归草稿执行中解耦,通过轻量级校正分支在logit空间注入因果信息,而主草稿计算保持并行。具体而言,Domino头包含一个因果编码器和一个低秩校正头。因果编码器使用轻量级GRU总结之前草稿token的嵌入,对于每个草稿位置 $i$,因果状态为 $S_{i-1} = \text{GRU}(E_{\leq i-1})$,其中 $E_j$ 是位置 $j$ 的token嵌入,$S_{i-1} \in \mathbb{R}^{d_s}$ 表示位置 $i$ 可用的前缀依赖信息。然后,低秩校正头产生logit空间的残差校正:$\Delta L_i = W_2 \sigma(W_1[H_i; S_{i-1}])$,其中 $W_1$ 将拼接的表示投影到低秩隐藏空间(维度 $r = 256$),$W_2$ 将低秩表示映射到词汇表空间,$\sigma$ 是SiLU激活函数。最终草稿logits为 $L_i = L_{base}^i + \Delta L_i$。这种设计的关键优势是:由于校正从低秩隐藏空间计算,它比重复合用完整LM头便宜得多;此外,在logit空间而不是隐藏空间应用校正避免了在每次因果更新后重新应用完整LM头,避免了昂贵计算回到顺序分支。

方法步骤详情

Domino的训练和推理过程可以分为以下步骤。训练阶段:首先,并行草稿主干(基于DFlash架构)从目标上下文特征和掩码块输入生成块级草稿表示。给定验证的前缀 $x_{\leq t}$,使用最后一个验证token $x_t$ 作为锚点构建掩码草稿块 $\tilde{x}_{t:t+B-1} = [x_t, \text{MASK}, \ldots, \text{MASK}]$,其中剩余的 $B-1$ 个位置对应未来草稿token。主干提取目标模型上下文特征 $C_t$ 并将掩码草稿块的嵌入作为输入,产生块级隐藏状态 $H_{t:t+B-1} = \text{Backbone}(C_t, \text{Embed}(\tilde{x}_{t:t+B-1}))$。然后,使用冻结的目标LM头计算未来位置的初步logits:$L_{base}^i = \text{LMHead}(H_i)$,其中 $i = t+1, \ldots, t+B-1$。接下来,Domino头使用教师强制训练策略:因果编码器接收ground-truth token嵌入而不是自生成前缀,因为只有之前token已被验证为正确的草稿位置才能贡献于最终接受长度。为了防止校正分支短路并行主干,作者采用基础锚定课程:$\mathcal{L} = (1 - \lambda_t) \mathcal{L}_{final} + \lambda_t \mathcal{L}_{base}$,其中 $\mathcal{L}_{final}$ 和 $\mathcal{L}_{base}$ 分别是从 $L_i$ 和 $L_{base}^i$ 计算的交叉熵损失。$\lambda_t$ 在训练过程中从1线性退火到0,因此目标最初锚定在基础logits上,随着Domino头接管残差校正而逐渐转向最终logits。推理阶段:并行主干生成基础logits $L_{base}^i$,Domino头顺序更新从之前采样的草稿token的因果状态 $S_{i-1}$,并生成校正logits $c_i$,从最终logits $l_i + c_i$ 采样每个草稿token $d_i$。为了最小化Domino头的开销,作者使用融合的Triton内核和CUDA图实现校正循环,将Domino头延迟从2.64ms降低到1.20ms。

技术新颖性

Domino的技术新颖性体现在三个方面。第一,训练策略创新:教师强制(Teacher Forcing)与训练时测试(TTT)不同,作者选择在训练时向因果编码器提供ground-truth token嵌入,因为自生成前缀可能嘈杂且经常不正确,尤其是训练早期。监督模型将此类损坏前缀映射到ground-truth下一个token创建了不存在于底层数据分布中的输入-输出映射,这可能降低GRU学习的因果表示质量。第二,基础锚定课程:防止校正分支在训练早期(接收干净前缀时)短路并行主干,避免基础logits变弱和最终预测过度依赖校正分支的问题。实验显示,没有课程时并行主干损失保持高位,而基础锚定课程保持主干损失稳定下降。第三,高效实现:使用融合的Triton内核和CUDA图实现校正循环,显著降低内核启动和Python级别的开销,使轻量级因果校正在实践中可行。此外,在logit空间应用校正而不是隐藏空间,避免了在每次因果更新后重新应用完整LM头,保持了并行计算的效率优势。

Overview of Domino.
Figure 3: Overview of Domino.
Left: parallel backbone loss with and without the base-anchored curriculum. Right: average acceptance length under TTT, TF, and TF+Curr.
Figure 4: Left: parallel backbone loss with and without the base-anchored curriculum. Right: average acceptance length under TTT, TF, and TF+Curr.

实验结果

Domino在多个基准测试中一致优于自回归草稿方法(如EAGLE-3)和并行草稿基线(如DART和DFlash)。在低并发场景下,Table 1显示在Qwen3-4B上,Domino在贪婪解码(T=0)下平均加速从DFlash的4.70×提高到5.47×,在Qwen3-8B上从4.66×提高到5.49×。在采样解码(T=1)下,平均加速从DFlash的4.03×提高到4.61×(Qwen3-4B)和从3.96×提高到4.46×(Qwen3-8B)。具体到任务,在Qwen3-8B上,Domino在GSM8K上达到7.92×加速和10.03接受长度,在MATH-500上达到7.38×加速和9.43接受长度,在HumanEval上达到5.89×加速和7.39接受长度。在高并发场景下,Table 2显示使用SGLang评估服务吞吐量,Domino在两个模型上都达到比EAGLE-3和DFlash更高的吞吐量。例如,在Qwen3-8B上GSM8K任务并发32时,Domino达到3650 TPS,而DFlash为2801 TPS,EAGLE-3(16)为1398 TPS。消融实验验证了设计选择的有效性。Table 3显示在相同训练数据(ShareGPT)下,Domino达到4.65接受长度,优于EAGLE-3的5.01(但吞吐量受限)和DFlash的3.90。Figure 4显示训练策略的影响:训练时测试(TTT)达到3.80接受长度,教师强制(TF)提高到3.96,TF+Curriculum进一步提高到4.19。Table 4显示Domino头的影响:启用Domino头将平均接受长度从3.49提高到4.19,平均加速从2.84×提高到3.31×。Table 6提供完整基准级结果,显示在所有测试任务上Domino头都带来一致的接受长度和加速提升。

Decoding speedup over vanilla autoregressive decoding and average acceptance length (τ) on Qwen3 models with a maximum of 2048 generated tokens.
Table 1: Decoding speedup over vanilla autoregressive decoding and average acceptance length (τ) on Qwen3 models with a maximum of 2048 generated tokens.
High-concurrency throughput on Qwen3 models.
Table 2: High-concurrency throughput on Qwen3 models.
Same-data comparison on Qwen3-8B under greedy decoding.
Table 3: Same-data comparison on Qwen3-8B under greedy decoding.
Effect of the Domino head.
Table 4: Effect of the Domino head.
Baseline draft model checkpoints used in our experiments.
Table 5: Baseline draft model checkpoints used in our experiments.
Full benchmark-level results for the Domino head ablation.
Table 6: Full benchmark-level results for the Domino head ablation.
Speedup comparison of Domino, DFlash, and EAGLE-3 relative to autoregressive decoding on Qwen3-8B using the Transformers backend.
Figure 2: Speedup comparison of Domino, DFlash, and EAGLE-3 relative to autoregressive decoding on Qwen3-8B using the Transformers backend.
查看结构化数据
任务指标本文基线提升
GSM8K (数学推理) 加速比/接受长度 7.92× / 10.03 (Qwen3-8B, T=0) DFlash: 5.21× / 6.59; EAGLE-3 (16): 2.21× / 3.27 相比DFlash加速提升52%,接受长度提升52%
MATH-500 (数学推理) 加速比/接受长度 7.38× / 9.43 (Qwen3-8B, T=0) DFlash: 6.18× / 7.87; EAGLE-3 (16): 2.09× / 3.10 相比DFlash加速提升19%,接受长度提升20%
HumanEval (代码生成) 加速比/接受长度 5.89× / 7.39 (Qwen3-8B, T=0) DFlash: 5.21× / 6.52; EAGLE-3 (16): 2.17× / 3.21 相比DFlash加速提升13%,接受长度提升13%
MT-Bench (对话) 加速比/接受长度 3.29× / 5.18 (Qwen3-8B, T=0) DFlash: 2.73× / 4.18; EAGLE-3 (16): 1.82× / 2.75 相比DFlash加速提升20%,接受长度提升24%
SGLang吞吐量 (GSM8K, 并发32) TPS (tokens/second) 3650 TPS (Qwen3-8B) DFlash: 2801 TPS; EAGLE-3 (16): 1398 TPS 相比DFlash吞吐量提升30%
SGLang吞吐量 (MBPP, 并发32) TPS (tokens/second) 2802 TPS (Qwen3-8B) DFlash: 2487 TPS; EAGLE-3 (16): 1291 TPS 相比DFlash吞吐量提升13%

局限与改进

作者承认了几个局限性。首先,本文专注于推理加速,而不是减少模型训练或微调的成本。其次,当前实现主要适配SGLang,与其他服务框架的兼容性仍有待系统评估。第三,实际加速可能因硬件平台而异,由于内存带宽、计算能力和内核效率的差异,不同环境中的部署可能需要平台特定的优化。此外,本文的评估主要集中在数学推理、代码生成和开放对话任务上,对于其他类型任务(如长文本生成、多模态任务)的有效性尚未充分探索。训练方面,Domino需要冻结目标模型并训练草稿模块,这可能需要针对不同目标模型进行专门训练,增加了部署的复杂性。

独立分析的弱点

Domino的几个潜在 weakness 值得考虑。第一,框架依赖:当前实现主要针对SGLang优化,与其他服务框架(如vLLM、TensorRT-LLM)的兼容性未知,这限制了其在生产环境中的广泛应用。改进方向是扩展到其他框架,可能需要重新设计内核融合和CUDA图实现以适应不同的框架架构。第二,训练成本:需要为每个目标模型专门训练Domino草稿模块,使用1.42M样本的指令调优数据集在8个A100 GPU上训练3个epoch,这可能对资源有限的团队构成障碍。改进方向是探索少样本或零样本适应技术,或者开发预训练的通用Domino头。第三,任务泛化性:评估集中在数学、代码和对话任务,对于其他任务类型(如创意写作、长文档生成)的有效性不确定。改进方向是在更广泛的任务上进行评估,并可能针对特定任务类型调整架构或训练策略。第四,硬件敏感性:实际加速可能因硬件而异,特别是内存带宽和内核效率差异大的情况下。改进方向是开发平台特定的优化,并发布针对不同GPU架构的优化版本。第五,延迟vs吞吐量:虽然Domino在高吞吐量场景下表现出色,但在极低延迟场景下(如单请求),额外的Domino头开销可能相对更大,需要更细致的权衡分析。

未来方向

作者提出或本文成果可延伸的未来研究方向包括:首先,框架兼容性扩展:系统评估Domino与vLLM、TensorRT-LLM等其他主流服务框架的兼容性,并开发相应的优化实现。其次,训练效率提升:探索参数高效微调技术(如LoRA)来适应Domino头到新目标模型,减少训练成本;或者开发通用预训练Domino头,只需轻量微调即可适应不同模型。第三,架构变体探索:尝试不同的因果编码器架构(如Transformer编码器、RNN变体),评估不同校正头设计(如不同低秩维度、不同激活函数);探索多尺度校正或层次校正策略。第四,任务泛化研究:将Domino应用到更多任务类型,如长文本生成、多模态生成、代码补全等,并针对特定任务调整设计。第五,理论分析:更深入分析轻量级因果校正的理论基础,理解在什么条件下logit空间校正最优,如何最优选择低秩维度和课程调度。第六,硬件优化:针对不同GPU架构(如NVIDIA H100、AMD GPU)开发平台特定优化,探索其他硬件(如TPU、专用推理芯片)的适配。第七,自适应策略:开发自适应草稿策略,根据输入复杂度或任务类型动态调整草稿块大小或校正强度。

复现评估

Domino的复现性评估如下。开源情况:论文提到代码和模型可在GitHub和Hugging Face获取(Links: Code (GitHub) | Models (Hugging Face)),但未提供具体链接。数据:使用mlabonne/open-perfectblend数据集,包含1.42M样本覆盖聊天、数学、代码和通用指令遵循任务。所有响应使用相应目标模型重新生成,而不是使用原始数据集响应。算力:除非另有说明,所有实验在NVIDIA A100-SXM4-80GB GPU上进行。训练使用8个A100 GPU,每GPU批大小2,全局批大小16无需梯度累积,训练3个epoch。实现细节:草稿块大小16,5层并行草稿主干,GRU因果编码器隐藏维度1024,低秩校正头隐藏维度256。使用AdamW优化器,学习率 $6 \times 10^{-4}$,零权重衰减,梯度裁剪最大范数1.0,余弦学习率调度预热比0.04,使用FSDP和梯度分片在bfloat16精度下训练。评估:遵循DFlash的评估协议,考虑数学推理(GSM8K、MATH、AIME25)、代码生成(HumanEval、MBPP、LiveCodeBench)和开放对话(MT-Bench、Alpaca)三类任务。报告平均接受长度 $\tau$ 和相对于自回归基线的端到端解码加速。难度评估:总体复现难度中等偏上。需要(1)A100 GPU集群用于训练,(2)DFlash基础实现,(3)Triton内核和CUDA图编程知识,(4)SGLang集成用于吞吐量评估。不过,如果作者提供了预训练检查点和完整代码,仅评估复现难度应较低。