Fast-dLLM: Training-free Acceleration of Diffusion LLM by Enabling KV Cache and Parallel Decoding
Fast-dLLM: Training-free Acceleration of Diffusion LLM by Enabling KV Cache and Parallel Decoding
Fast-dLLM:通过启用 KV 缓存和并行解码实现扩散大语言模型的免训练加速
作者/机构: Chengyue Wu (香港大学, NVIDIA), Hao Zhang (NVIDIA), Shuchen Xue (独立研究员), Zhijian Liu (NVIDIA), Shizhe Diao (NVIDIA), Ligeng Zhu (NVIDIA), Ping Luo (香港大学), Song Han (NVIDIA, MIT), Enze Xie (NVIDIA)
A1 主要贡献
本文旨在解决开源的基于扩散的大语言模型(Diffusion LLM)在实际推理速度上落后于自回归(AR)模型的核心问题。其性能瓶颈主要源于两个方面:首先,扩散大语言模型由于其双向注意力机制,天然不支持自回归模型中用于加速推理的关键组件——键值(Key-Value, KV)缓存;其次,在尝试并行解码多个词元(token)时,其生成质量会显著下降。
为了弥合这一性能差距,本文提出了 Fast-dLLM,一个无需训练的加速框架,主要包含以下创新点:
-
分块近似KV缓存机制 (Block-Wise Approximate KV Cache):针对扩散模型中的双向注意力特性,设计了一种新颖的近似KV缓存机制。该方法采用分块生成(block-wise generation)的策略,在生成一个块之前,计算并缓存其他块的KV值以供重用。在生成完一个块之后,再重新计算所有块的KV缓存。实验可视化证明,在块内相邻的推理步骤中,KV激活值的相似性非常高,因此这种近似方法能够在几乎不损失模型性能的前提下,显著减少冗余计算。本文还提出了一个双缓存(DualCache)版本,同时缓存前缀(prefix)和后缀(suffix)词元的键值,进一步提升计算重用率。
-
置信度感知的并行解码策略 (Confidence-Aware Parallel Decoding):本文通过理论分析和实证研究,指出并行解码导致质量下降的根本原因在于,在条件独立性假设下同时对相互依赖的词元进行采样,破坏了词元间的关键依赖关系。为解决此问题,Fast-dLLM 提出了一种基于置信度阈值的并行解码策略。该策略不再是像先前工作那样解码固定数量(top K)的词元,而是动态选择置信度超过某个全局阈值的词元进行解码。这种方法选择性地解码高置信度的词元,减轻了对词元依赖性的破坏,从而在保持高质量生成的同时,实现了高达13.3倍的推理加速。
-
SOTA的加速效果与泛化性验证:本文在多个开源扩散大语言模型(LLaDA, Dream)和四个主流基准测试(GSM8K, MATH, HumanEval, MBPP)上进行了全面实验。结果表明,Fast-dLLM 能够稳定地带来数量级的速度提升,同时对模型准确率的负面影响极小,甚至没有影响。特别是在处理长序列(例如,生成长度为1024)的场景下,与原始 LLaDA 基线相比,实现了高达 27.6倍 的端到端加速。这证明了该方法的通用性和在实际部署中的价值。
(c) 相对于原始LLaDA基线的端到端加速倍数。
图 1 | Fast-dLLM各组件在不同方法中的有效性。 我们使用NVIDIA A100 GPU,批大小为1,无其他推理加速框架。(a) 在最大生成长度为256的条件下,不同设计和模型在推理吞吐量(词元/秒)和GSM8K(5-shot)准确率上的表现。缓存机制和并行解码能显著加速推理,二者结合最多可带来8.1倍的吞吐量提升,且准确率下降可忽略不计。(b) 我们通过展示每步生成的词元数(折线图)和总吞吐量(柱状图)来分解每种方法的贡献。(c) 在长前缀填充(8-shot)和最大生成长度为1024的条件下,我们的组合方法相比原始LLaDA基线实现了高达27.6倍的端到端加速。
A3 背景知识
2.1. 掩码扩散模型 (Masked Diffusion Model)
离散扩散模型的演进。用于离散数据的扩散模型最早由【29, Jascha Sohl-Dickstein et al. Deep unsupervised learning using nonequilibrium thermodynamics. ICML 2015】和【11, Emiel Hoogeboom et al. Argmax flows and multinomial diffusion: Learning categorical distributions. NeurIPS 2021】探索。随后,D3PM【2, Jacob Austin et al. Structured denoising diffusion models in discrete state-spaces. NeurIPS 2021】提出了一个更通用的框架,通过特定的转移矩阵 $Q_{t|t-1}$ 定义离散状态马尔可夫链作为前向加噪过程,并通过最大化证据下界(ELBO)来学习反向过程的参数化模型 $p_{\theta}(x_0|x_t)$。CTMC【3, Andrew Campbell et al. A continuous time framework for discrete denoising models. NeurIPS 2022】进一步将D3PM扩展到连续时间,将其形式化为一个连续时间马尔可夫链(CTMC)框架。另一种不同的方法是SEDD【17, Aaron Lou et al. Discrete diffusion language modeling by estimating the ratios of the data distribution. arXiv 2023】,它通过参数化似然比 $p_t(x)/p_0(x)$ 来学习反向过程,并使用去噪分数熵(Denoising Score Entropy)来训练该比率。
掩码扩散模型(MDM)的定义。在各种离散扩散的加噪过程中,掩码扩散模型(MDMs),也被称为吸收状态离散扩散模型,获得了广泛关注。MDMs采用一种前向加噪过程,其中词元被逐步替换为一个特殊的 [MASK] 词元。该过程由转移概率定义:
这里,$t \in [0, 1]$ 表示扩散时间(或掩码级别),控制着从原始数据 $x_0$(在 $t=0$ 时)到完全被掩码的序列(在 $t=1$ 时)之间的插值。
MDM的简化训练目标。最近,MDLM【28, Jiaxin Shi et al. Simplified and generalized masked diffusion for discrete data. arXiv 2024; 27, Subham Sekhar Sahoo et al. Simple and effective masked diffusion language models. arXiv 2024; 41, Kaiwen Zheng et al. Masked diffusion models are secretly time-agnostic masked models and exploit inaccurate categorical sampling. arXiv 2024】和 RADD【22, Jingyang Ou et al. Your absorbing discrete diffusion secretly models the conditional distributions of clean data. arXiv 2024】的研究表明,对于MDMs,不同的参数化是等价的。此外,他们证明了MDMs的训练目标可以被简化或直接从数据似然中推导出来。这得到了以下的目标函数,即 $\log p_{\theta}(x)$ 的证据下界(ELBO):
2.2. MDMs的生成过程
生成过程的挑战与加速。直接解析地逆转公式1中定义的前向过程在生成时计算效率低下,因为它通常每步只修改一个词元【3, Andrew Campbell et al. A continuous time framework for discrete denoising models. NeurIPS 2022; 17, Aaron Lou et al. Discrete diffusion language modeling by estimating the ratios of the data distribution. arXiv 2023】。一个常见的加速策略是采用 $\gamma$-leaping【6, Daniel T Gillespie. Approximate accelerated stochastic simulation of chemically reacting systems. The Journal of chemical physics 2001】近似来处理反向过程。在MDMs的背景下,这允许一个迭代的生成过程,其中多个被掩码的词元可以从一个噪声水平 $\tau$ 近似地在单一步骤中恢复到一个更早的水平 $t < \tau$。
这里,$p_{0|t}(x_{0i}|x_t)$(当 $x_{ti} = \text{[MASK]}$ 时)代表模型为预测一个非 [MASK] 词元而提供的词汇表上的分布。在涉及条件数据的场景中,例如根据提示 $c$ 生成回应 $x_0$,MDM的反向过程(如公式3所定义)需要调整。具体来说,模型用于揭示词元 $x_{ti}$ 的预测分布 $p_{0|t}(x_{0i}|x_t)$ 现在也需要以提示 $c$ 为条件,即 $p_{0|t}(x_{0i}|x_t, c)$。
并行解码的诅咒。直接逆转公式1的前向过程进行生成速度很慢,通常每步只改变一个词元【3, Andrew Campbell et al. A continuous time framework for discrete denoising models. NeurIPS 2022; 17, Aaron Lou et al. Discrete diffusion language modeling by estimating the ratios of the data distribution. arXiv 2023】。一个常用的加速策略是采用 $\gamma$-leaping【6, Daniel T Gillespie. Approximate accelerated stochastic simulation of chemically reacting systems. The Journal of chemical physics 2001】近似来处理反向过程。对于MDMs,这意味着多个被掩码的词元将在单一步骤中并行生成。然而,由于条件独立性假设,多词元预测中出现了一个重大挑战。考虑一个来自【30, Jiaming Song and Linqi Zhou. Ideas in inference-time scaling can benefit generative pre-training algorithms. arXiv 2025】的例子:由两个英文单词组成的扑克牌手牌列表是: 。接下来的两个词可以是“high card”、“two pair”、“full house”或“straight flush”。值得注意的是,这两个词之间存在相关性。然而,MDMs中的多词元预测过程首先为每个词元生成一个概率分布,然后独立地从这些分布中采样。这种独立采样可能导致不希望的组合,例如“high house”。
条件独立性假设的形式化。为了形式化这个问题,考虑揭示两个词元位置 $i$ 和 $j$。由于条件独立性假设,MDMs从 $p(x_{0i}|x_t) \cdot p(x_{0j}|x_t)$ 中对它们进行采样。然而,真实的联合概率需要考虑它们之间的依赖关系:
$p(x_{0i}, x_{0j}|x_t) = p(x_{0i}|x_t) \cdot p(x_{0j}|x_t, x_{0i})$(或者对称地,通过将 $i$ 条件化于 $j$)。这种假设的独立生成与真实依赖的数据分布之间的差异会降低生成序列的质量和连贯性。当在单一步骤中同时揭示大量词元时,这个问题会变得更加严重。
图 2 | 我们的分块解码键值缓存机制示意图。 (a) 在仅前缀缓存期间,KV缓存为提示计算一次,并在每个块内的多个解码步骤中重复使用。在完成一个块后更新缓存以保持一致性,开销可忽略不计。(b) 双缓存(DualCache)通过缓存前缀和被掩码的后缀词元来扩展此方法,进一步加速解码。跨步骤的KV激活值的高度相似性使得在最小化近似误差的情况下能够有效重用。
A2 方法细节
3.1. 流程概述
Fast-dLLM框架。我们的方法 Fast-dLLM 建立在掩码扩散模型(MDM)架构之上,以实现高效和高质量的序列生成。为了加速推理,整个流程整合了两个关键策略:通过键值(KV)缓存实现高效的注意力计算,以及一个由预测置信度引导的并行解码方案。
两大核心策略。具体来说,我们采用了分块解码的键值缓存(Key-Value Cache for Block-Wise Decoding),这允许在不同步骤之间重用注意力激活值,从而显著减少了冗余计算。在每个块内部,我们进一步提出了置信度感知的并行解码(Confidence-Aware Parallel Decoding),根据置信度分数选择性地更新词元,以在保持输出质量的同时提高效率。
组合效果。通过结合这些策略,Fast-dLLM 在对生成性能影响最小的情况下,显著加快了MDMs的推理速度。整个过程在算法1中进行了总结。
3.2. 分块解码的键值缓存
分块解码与KV缓存。如图2所示,我们采用分块解码策略来支持键值(KV)缓存的使用。最初,我们为提示(prompt)计算并存储KV缓存,该缓存在整个块0(Block 0)中被重用。在每个块内部,相同的缓存被用于多个解码步骤。在完成一个块的解码后,我们会更新所有词元(不仅是新生成的词元)的缓存。这个缓存更新可以与解码步骤联合执行,因此与不使用缓存相比,没有额外的计算开销。由于掩码扩散模型中使用的是完全注意力机制【21, Shen Nie et al. Large language diffusion models. 2025; 36, Jiacheng Ye et al. Dream 7b. 2025】,这种方法导致了一个近似的解码过程。
近似KV缓存的有效性。我们近似KV缓存方法的有效性源于一个观察:KV激活值在相邻的推理步骤中表现出高度的相似性,如图3所示。图3a中的红色框区域突显了块内的相似性分数,这些分数始终接近于1。这表明在块解码期间,前缀(prefix)的键(keys)和值(values)的差异可以忽略不计,使我们能够在不显著损失准确性的情况下安全地重用缓存。
双缓存(DualCache)。此外,我们实现了一个双向版本的KV缓存机制,名为双缓存(DualCache),它不仅缓存前缀词元,还缓存后缀词元。在我们的分块解码方案下,后缀词元完全由掩码词元组成。如表4所示,双缓存带来了进一步的加速。图3b中的红色框区域进一步证明,在块解码期间,后缀的键和值的差异也是可以忽略的。
图 3 | LLaDA-Instruct中推理步骤间键值激活余弦相似度热力图。 (a) 提示块的余弦相似度热力图,对所有提示词元取平均值。(b) 最后一个块的余弦相似度热力图,对最后一个块中的所有词元取平均值(用于代表后缀词元,因为最后一个块在解码前总是属于后缀)。在(a)和(b)中,对角线附近($t \approx \tau$)观察到高度相似性,表明一个块内相邻推理步骤的键值激活值高度相似。红色框区域突显了这一效应,支持使用近似的分块KV缓存:先前步骤的缓存激活值可以在块解码期间安全地重用,而准确性损失极小。双缓存策略额外缓存了后缀词元,同样显示出在块解码期间激活值的差异可以忽略不计,从而在保持有竞争力的准确性的同时实现更大的加速。
3.3. 置信度感知的并行解码
方法动机。虽然存在使用辅助模型来显式捕捉词元间依赖关系的方法【16, Anji Liu et al. Discrete copula diffusion. arXiv 2024; 34, Minkai Xu et al. Energy-based diffusion language models for text generation. arXiv 2024】,但它们通常会增加整个流程的复杂性。与这些方法相反,我们提出了一个简单而有效的置信度感知解码算法,旨在缓解条件独立性问题。
算法描述。具体来说,在每次迭代中,我们不是用独立的边际概率激进地揭示所有被掩码的词元,而是为每个词元计算一个置信度分数(例如,最大softmax概率)。只有那些置信度超过阈值的词元才在当前步骤被揭示;其余的保持被掩码状态,并在未来的步骤中重新考虑。如果没有词元的置信度超过阈值,我们总是揭示置信度最高的那个词元,以确保进程推进并防止无限循环。该策略在加速生成的同时,减少了由不确定或模糊预测引起的错误。
理论基础。然而,一个关键问题是:在理论上,何时使用独立的边际概率进行并行解码是合理的,尽管真实的联合分布可能包含依赖关系?我们用以下正式结果来回答这个问题,它描述了在高置信度情况下,贪婪并行(边际分布乘积)解码等价于贪婪序贯(真实联合分布)解码的条件,并量化了两种分布之间的差异。
数学符号定义。在陈述定理之前,我们定义其陈述中使用的数学符号。令 $P_{\theta}(\cdot|c)$ 表示由MDM在给定 $c$(包括提示 $c_0$ 和先前生成的词元)的条件下给出的条件概率质量函数(PMF)。假设模型要预测不在 $c$ 中的 $k$ 个位置 $i_1, \ldots, i_k$ 的词元。令 $x = (x_{i_1}, \ldots, x_{i_k})$ 为 $k$ 个词元的向量,其中每个 $x_{i_j}$ 在词汇表 $\mathcal{V}$ 中取值。令 $P(x|c) \equiv P_{\theta}(x_{i_1}, \ldots, x_{i_k}|c)$ 为模型下的联合条件PMF。令 $P_j(x_{i_j}|c) \equiv P_{\theta}(x_{i_j}|c)$ 为位置 $i_j$ 的边际条件PMF。并行解码使用边际概率的乘积来生成词元:$Q(x|c) = \prod_{j=1}^{k} P_j(x_{i_j}|c)$。定理1的证明及相关讨论见附录A。
定理1 (高置信度下的并行解码)。假设存在一个特定的词元序列 $x^* = (x_{i_1}^*, \ldots, x_{i_k}^*)$,使得对于每个 $j \in \{1, \ldots, k\}$,模型对 $x_{i_j}^*$ 具有高置信度:$P_j(x_{i_j} = x_{i_j}^*|c) > 1 - \epsilon$,其中 $\epsilon > 0$ 是一个小值。那么,以下结果成立:
1. 贪婪解码的等价性:如果 $(k+1)\epsilon \le 1$(即 $\epsilon \le \frac{1}{k+1}$),则
这意味着贪婪并行解码(选择 argmax $Q$)与贪婪序贯解码(选择 argmax $P$)产生相同的结果。
这个界是紧的:如果 $\epsilon > \frac{1}{k+1}$,则存在满足高置信度边际假设的分布 $P(x|c)$,使得 $\text{argmax}_x P(x|c) \neq \text{argmax}_x Q(x|c)$。
2. 距离和散度界:为简洁起见,将 $P(\cdot|c)$ 和 $Q(\cdot|c)$ 分别表示为 $P$ 和 $Q$。
- $L_p$ 距离 ($p \ge 1$):对于 $p > 1$,$L_p(P, Q) < ((k-1)\epsilon^p + 2\epsilon)^{1/p}$。特别地,对于全变差距离($TV(P, Q) = \frac{1}{2} L_1(P, Q)$):$TV(P, Q) < \frac{3k-1}{2}\epsilon$。
- 前向KL散度:对于 $k > 1$,$D_{\text{KL}}(P\|Q) < (k-1)(H_b(\epsilon) + \epsilon \ln(|\mathcal{V}|-1))$,其中 $H_b(\epsilon) = -\epsilon \ln \epsilon - (1-\epsilon) \ln(1-\epsilon)$ 是二元熵函数,$|\mathcal{V}|$ 是词汇表的大小。
基于因子的并行解码策略。基于该定理,我们提出了一个实用的基于因子的并行解码策略,作为阈值策略的扩展,它根据置信度水平自适应地选择并行解码的词元数量。具体来说,给定模型对一个块中 $m$ 个词元的边际置信度估计,我们对这些置信度进行排序,并选择最大的 $k$ 个,使得 $(k+1)(1 - c_{(k)}) < \gamma$,其中 $\gamma$ 是一个固定的解码因子超参数,$c_{(k)}$ 是第 $k$ 高的置信度。在每一步中,前 $k$ 个词元被并行解码。这个公式反映了定理1中的界,并确保只有当边际置信度足够高以可靠地近似联合解码时才进行解码。与静态的基于阈值的策略相比,基于因子的解码以一种有理论依据的方式动态地控制并行度。
算法 1 分块置信度感知的并行解码与(双)KV缓存
需要: 模型 $\theta$, 提示 $c_0$, 回答长度 $L$, 块数 $B$, 块大小 $S$, 每块步数 $T$, 阈值 $\tau$, 是否使用双缓存 use_DualCache, 策略 strategy $\in \{\text{threshold, factor}\}$, 因子 $\gamma$
1: $x \leftarrow [c_0; \text{[MASK]}, ..., \text{[MASK]}]$
2: 为 $x$ 初始化KV缓存(单向或双向)(与解码融合) // KV缓存初始化
3: for $b = 1$ to $B$ do
4: $start \leftarrow |c_0| + (b-1)S$, $end \leftarrow |c_0| + bS$
5: for $t = 1$ to $T$ do
6: 使用缓存,在 $x[start, end)$ 上运行 $\theta$(如果 use_DualCache 则在 $x[0, :)$ 上运行) // 缓存重用
7: 对于被掩码的 $x_i$,计算置信度 $c_i = \max_v p_{\theta}(v|·)$ // 置信度评分
8: if strategy == threshold then
9: 揭示 $[start, end)$ 中所有 $c_i \ge \tau$ 的词元,并始终揭示 $c_i$ 最大的词元
10: else if strategy == factor then
11: 将 $c_i$ 降序排序为 $(c_{(1)}, c_{(2)}, ..., c_{(m)})$
12: 找到最大的 $k$ 使得 $(k+1)(1 - c_{(k)}) < \gamma$
13: 揭示前 $k$ 个词元,并始终揭示 $c_i$ 最大的词元
14: end if
15: if 所有 $x[start, end)$ 都被揭示 then
16: break
17: end if
18: end for
19: 更新KV缓存:如果 use_DualCache,更新前缀和后缀;否则,只更新前缀。 // 缓存更新
20: end for
21: return $x$
A4 实验
实验环境
- 硬件配置:所有实验均在单张 NVIDIA A100 80GB GPU 上进行。
- 模型:评估了两个近期的基于扩散的语言模型:LLaDA【21, Shen Nie et al. Large language diffusion models. 2025】、LLaDA-1.5【44, Fengqi Zhu et al. Llada 1.5: Variance-reduced preference optimization for large language diffusion models. 2025】和 Dream【36, Jiacheng Ye et al. Dream 7b. 2025】。此外,还评估了LLaDA的多模态变体 LLaDA-V【38, Zebin You et al. Llada-v: Large language diffusion models with visual instruction tuning. 2025】。
- 数据集与基准:
- 语言任务:使用了四个广泛使用的数据集——GSM8K、MATH、HumanEval 和 MBPP,以评估在多样化的推理和代码生成任务上的性能。
- 多模态任务:针对 LLaDA-V,使用了两个具有挑战性的多模态基准:MathVista 和 MathVerse,它们要求解决植根于复杂视觉场景的数学问题。
- 软件配置:所有评估均使用标准化的
lm-eval库进行,以确保一致性和可复现性。 - 超参数:
- KV 缓存:缓存块大小在 4 到 32 之间变化。
- 并行解码:置信度阈值在 0.5 到 1.0 的范围内探索。
- 默认配置:除非另有说明,实验使用前缀缓存(PrefixCache),块大小为32,置信度阈值为0.9。
- 评估指标:推理吞吐量以每秒生成的平均输出词元数量来衡量,计算从序列开始到结束(
<eos>)的完整过程,反映了真实的端到端解码速度。
实验结果
表 1 | LLaDA-Instruct套件的综合基准测试结果。 每个单元格显示准确率和解码吞吐量(词元/秒),以及相对于LLaDA基线(底行,蓝色:词元/秒/橙色:相对加速倍数)的相对加速。每个配置的最高吞吐量和加速倍数已高亮显示。
表 2 | Dream-Base变体在四个任务上不同生成长度(256和512)的综合基准测试结果。 每个单元格显示准确率(顶行)和解码吞吐量(词元/秒),以及相对于Dream-Base基线(底行,蓝色:词元/秒/橙色:相对加速倍数)的相对加速。黄色数字表示每个配置的最高吞吐量和加速倍数。
主要结果总结:
- 综合性能提升:如表1和表2所示,Fast-dLLM 在 LLaDA-Instruct 和 Dream-Base 模型上均取得了显著的速度和性能提升。
- 单独使用 KV缓存 机制,通常能带来 2倍至3.6倍 的速度提升。
- 单独使用 并行解码 策略,加速效果更明显,通常能达到 4倍至6倍。
- 当 两者结合 时,性能提升最为显著。在 LLaDA 上,组合方法在 GSM8K (长度512) 上最高可达 11倍 加速,在 MBPP (长度512) 上可达 9.2倍。在 Dream-Base 上,MBPP (长度512) 任务上最高可达 7.8倍 加速。这表明两种方法不仅各自有效,而且高度互补。
- 准确率保持稳定:所有加速方法带来的性能提升几乎没有以准确率为代价。在所有基准测试和设置中,加速后模型的准确率与基线相比,差距在1-2个百分点以内,在某些情况下甚至略有提高。
- 对长序列更有效:在少样本学习和代码生成等需要生成较长序列的场景中,由于缓存重用和批量计算的机会更多,Fast-dLLM 的加速效果更为明显。LLaDA-1.5 的评估(表12)也显示了在保持更强准确率的同时,实现了相当或更高的吞吐量。
- 多模态任务上的有效性:Fast-dLLM 在多模态模型 LLaDA-V 上也表现出色。如表3所示,通过调整缓存策略(采用基于刷新的更新而非小块缓存),在 MathVista 数据集上实现了高达 9.9倍 的加速,且准确率下降极小。在 MathVerse 上,准确率甚至略有提升,证明了该方法对多模态推理任务的广泛适用性。
表 3 | LLaDA-V在MathVista和MathVerse上的性能与加速比较。 每个基准测试包括完整步骤、半数步骤和Fast-dLLM的结果。Fast-dLLM显著提高了吞吐量(高亮显示),且准确率损失极小。
消融实验与分析:
- 上下文与生成长度的影响:如表4和表5所示,更长的前缀填充(prefill length)和生成长度能带来更大的加速比。例如,在生成长度为1024时,将上下文从5-shot增加到8-shot,DualCache的加速比从19.6倍提升到27.6倍。这符合预期,即在更长的序列上分摊计算开销能带来更显著的效率增益。
- PrefixCache vs. DualCache:如表5所示,DualCache 通常比 PrefixCache 实现更高的加速比,特别是在长序列生成任务中。在8-shot、生成长度1024的设置下,DualCache 实现了27.6倍的加速,优于PrefixCache的18.6倍,同时保持了有竞争力的准确率。
- 缓存块大小的影响:图4显示,块大小为32时在吞吐量和准确率之间取得了最佳平衡。较小的块大小准确率更高但缓存更新开销大,而较大的块大小则可能因上下文不匹配而降低准确率。
图 4 | 缓存块大小对准确率和吞吐量的影响。 橙色线表示不同缓存块大小对吞吐量的影响,而蓝色线表示准确率。
- 动态阈值 vs. 固定步长策略:如图5所示,置信度感知的并行解码策略(动态阈值)在准确率和效率上均优于固定每步解码词元数的基线方法。它能以更少的函数评估次数(NFE)达到更高的准确率,并在接近单步解码(1-token)准确率的同时,实现远超后者的吞吐量。
图 5 | (a) GSM8K(5-shot)准确率在不同置信度阈值下的变化。 红色线上的数字表示每步解码的平均词元数。三条虚线代表基线方法在每步选择前2、4或8个词元时的准确率。(b) 不同置信度阈值下所需的推理步数。(c) 我们的方法与基线在GSM8K(5-shot)准确率上的比较,以每步平均词元数为横轴。我们的方法始终优于基线。
- 因子解码策略的优势:与固定步长基线(图8)和阈值策略(表11)相比,基于因子的解码策略在准确率和吞吐量之间提供了更好的权衡。它能自适应地控制并行度,以更少的推理步骤实现有竞争力的准确率,并获得比阈值策略更高的吞吐量。
- 解码效率分析与局限性:在计算密集型场景中,PrefixCache 可为 LLaDA 带来高达5倍的吞吐量提升。在小批量下,其吞吐量可与自回归模型 LLaMA 相媲美甚至超越。然而,随着批量大小增加,LLaMA 从内存密集型转向计算密集型,性能优势扩大,而 LLaDA 等扩散模型因完全注意力机制带来的高计算开销,扩展性受限。
表 4 | 在生成长度为1024时,LLaDA在5-shot和8-shot设置下性能与加速的比较。 此表比较了在5-shot和8-shot配置下,使用1024生成长度时不同解码策略的准确率和吞吐量加速。结果表明,增加前缀填充长度增强了缓存策略的有效性,特别是对于双缓存(DualCache)。
表 5 | 在8-shot设置下,生成长度对LLaDA准确率和加速的影响。 此表展示了在8-shot设置下,不同生成长度(256、512和1024)对不同缓存策略解码性能和效率的影响。更长的生成长度带来了更高的吞吐量增益,特别是对于双缓存(DualCache),验证了我们方法的可扩展性。
A7 补充细节
5. 相关工作
5.1. 扩散大语言模型 (Diffusion LLM)
扩散模型的演进。扩散模型已成为生成模型领域的一大变革性范式,最初在图像【25, Robin Rombach et al. High-resolution image synthesis with latent diffusion models. 2022; 19, Alex Nichol et al. Glide: Towards photorealistic image generation and editing with text-guided diffusion models. 2022; 23, Aditya Ramesh et al. Zero-shot text-to-image generation. 2021; 26, Chitwan Saharia et al. Photorealistic text-to-image diffusion models with deep language understanding. 2022】和音频合成【35, Dongchao Yang et al. Diffsound: Discrete diffusion model for text-to-sound generation. 2023; 12, Rongjie Huang et al. Make-an-audio: Text-to-audio generation with prompt-enhanced diffusion models. 2023】等连续领域取得了巨大成功,之后扩展到自然语言处理领域。近期,离散扩散模型【2, 20, 21, 11, 3, 10, 18, 24, 31, 14, 42, 4, 37, 27, 28, 41, 5, 40, 39】的进展重塑了文本生成的格局,为大型语言模型(LLMs)中的自回归(AR)范式提供了可行的替代方案。这些模型通过创新的数学公式重新定义了噪声注入和去噪过程,解决了离散数据的内在挑战。
离散扩散的理论基础。离散数据的扩散模型最早由【29, Jascha Sohl-Dickstein et al. Deep unsupervised learning using nonequilibrium thermodynamics. ICML 2015】和【11, Emiel Hoogeboom et al. Argmax flows and multinomial diffusion: Learning categorical distributions. NeurIPS 2021】探索。随后,D3PM【2, Jacob Austin et al. Structured denoising diffusion models in discrete state-spaces. NeurIPS 2021】提供了一个更通用的框架。该框架将前向加噪过程建模为使用特定转移矩阵的离散状态马尔可夫链。对于反向过程,D3PM通过最大化证据下界(ELBO)学习了给定噪声版本下原始数据的条件概率的参数化模型。CTMC【3, Andrew Campbell et al. A continuous time framework for discrete denoising models. NeurIPS 2022】进一步将D3PM扩展到连续时间设置,将其形式化为连续时间马尔可夫链(CTMC)。SEDD【17, Aaron Lou et al. Discrete diffusion language modeling by estimating the ratios of the data distribution. arXiv 2023】则采用了一种独特的方法,通过参数化给定噪声时间步下不同数据实例的边际似然比来学习反向过程,该比率模型使用去噪分数熵目标进行训练。最近,MDLM【28, 27, 41】和 RADD【22, Jingyang Ou et al. Your absorbing discrete diffusion secretly models the conditional distributions of clean data. arXiv 2024】对掩码扩散模型(MDMs)的研究带来了重要的澄清。这些研究表明,MDMs的不同参数化是等价的。
与预训练语言模型的集成。一个关键的突破是将离散扩散与现有的大语言模型架构相结合。Diffusion-NAT【43, Kun Zhou et al. Diffusion-nat: Self-prompting discrete diffusion for non-autoregressive text generation. 2023】将离散扩散的去噪过程与BART【15, Mike Lewis et al. Bart: Denoising sequence-to-sequence pre-training for natural language generation, translation, and comprehension. 2019】的非自回归解码统一起来,实现了对掩码词元的迭代细化。通过将BART的推理与扩散步骤对齐,该方法利用了预训练知识,同时保持了比同类自回归Transformer快20倍的生成速度。类似地,LLaDA【21, Shen Nie et al. Large language diffusion models. 2025】和DiffuLLaMA【7, Shansan Gong et al. Scaling diffusion language models via adaptation from autoregressive models. arXiv 2024】框架使用掩码去噪将扩散模型扩展到70亿参数,而LLaDA和Dream【36, Jiacheng Ye et al. Dream 7b. 2025】通过在扩散时间步上进行递归词元预测,展现了与LLaMA3【9, Aaron Grattafiori et al. The llama 3 herd of models. 2024】等自回归基线相媲美的性能。
5.2. 大语言模型加速 (LLM Acceleration)
键值缓存 (Key-Value Cache)。键值(KV)缓存是现代大型语言模型(LLM)在使用Transformer架构【32, Ashish Vaswani. Attention is all you need. arXiv 2017】进行推理时的一项基础优化技术。它通过存储和重用先前计算的注意力状态,实现了高效的自回归文本生成。然而,由于完全注意力机制的存在,将KV缓存应用于像LLaDA这样的扩散语言模型并非易事。块扩散(Block diffusion)【1, Marianne Arriola et al. Block diffusion: Interpolating between autoregressive and diffusion language models. 2025】通过逐块生成的方式克服了先前扩散语言模型的一个关键限制,使得先前解码块的键和值可以被存储和重用。
非自回归生成 (Non-Autoregressive Generation)。非自回归(NAR)生成标志着从顺序词元生成向同时生成多个词元的根本性转变,从而显著加速了推理过程【33, Yisheng Xiao et al. A survey on non-autoregressive generation for neural machine translation and beyond. 2023】。NAR方法最初被引入神经机器翻译领域,此后已扩展到各种任务,包括语法错误纠正、文本摘要、对话系统和自动语音识别。尽管NAR生成相比自回归方法具有显著的速度优势,但它通常会牺牲生成质量。扩散LLMs代表了非自回归文本生成的一种新范式;然而,先前的工作【21, Shen Nie et al. Large language diffusion models. 2025】由于输出质量的显著下降,未能实现预期的加速效果。
A5 结论
本文解决了基于扩散的大语言模型(Diffusion LLMs)在推理效率方面的主要局限性,这些模型历史上缺乏对KV缓存的支持,并且在并行解码时表现出性能下降。为了弥合与自回归模型的差距,我们提出了 Fast-dLLM,一个基于扩散的框架,它引入了一种为扩散LLM的双向注意力特性量身定制的近似KV缓存机制,并通过分块生成方案得以实现。此外,我们确定了有效并行解码的主要障碍是由于条件独立性假设导致的词元依赖关系被破坏。为了解决这个问题,Fast-dLLM 采用了一种置信度感知的并行解码策略,促进了安全高效的多词元生成。在多个基准测试和模型基线(LLaDA和Dream)上的广泛实验表明,Fast-dLLM 实现了高达 27.6倍 的加速,而准确率损失极小。这些发现为将扩散LLM部署为实际应用中自回归模型的有力替代方案提供了实用的解决方案。
A6 附录
A. 证明
本节将给出定理1的详细证明和讨论。
证明。步骤1:证明 $x^*$ 是 $Q(x)$ 的唯一最大化者。
令 $p_{j}^* = P_j(x_{i_j} = x_{i_j}^*|c)$。我们已知 $p_{j}^* > 1 - \epsilon$。令 $p'_{j} = 1 - p_{j}^* = P_j(x_{i_j} \neq x_{i_j}^*|c)$。因此,$p'_{j} < \epsilon$。边际概率乘积的概率质量函数(PMF)为
要最大化 $Q(x|c)$,我们必须独立地最大化每一项 $P_j(x_{i_j} = x_j|c)$。条件 $(k+1)\epsilon \le 1$ 意味着 $\epsilon \le 1/(k+1)$。由于 $k \ge 1$,因此 $1/(k+1) \le 1/2$。所以 $\epsilon \le 1/2$。因此,对于选择的 $x_{i_j}^*$:
这意味着 $x_{i_j}^*$ 是 $P_j(\cdot|c)$ 的唯一最大化者。所以,
步骤2:证明 $x^*$ 是 $P(x)$ 的唯一最大化者。
我们想要证明对于所有 $x \neq x^*$,都有 $P(x^*|c) > P(x|c)$。使用 Bonferroni 不等式:
因为对所有 $j$ 都有 $p'_j < \epsilon$,所以我们有 $\sum_{j=1}^{k} p'_j < k\epsilon$。所以,
现在考虑任何 $x = (x_1, \ldots, x_k)$ 使得 $x \neq x^*$。这意味着至少有一个索引 $j$ 使得 $x_j \neq x_{i_j}^*$。事件 $\{X=x\}$ 是事件 $\{X_{i_j} = x_j\}$ 的子事件。所以,
由于 $x_j \neq x_{i_j}^*$,
因此,
要使 $P(x^*|c) > P(x|c)$ 成立,只需满足
这简化为 $1 \ge (k+1)\epsilon$,即 $\epsilon \le \frac{1}{k+1}$。定理假设 $(k+1)\epsilon < 1$,这正是这个条件。严格不等式 $P(x^*|c) \ge 1 - \sum p' > 1 - k\epsilon$ 和 $P(x|c) \le p' < \epsilon$ 确保了 $P(x^*|c) > P(x|c)$。因此,
结合 $Q$ 的 argmax,这证明了第1部分的主要论点:
步骤3:界 $\frac{1}{k+1}$ 的紧性。
界 $\epsilon \le \frac{1}{k+1}$ 是紧的。这意味着如果 $\epsilon > \frac{1}{k+1}$,可以构造一个场景,其中边际条件 $P_j(x_{i_j} = x_{i_j}^*|c) > 1 - \epsilon$ 成立,但 $\text{argmax}_x P(x|c) \neq x^*$(只要 $\epsilon \le 1/2$,$x^*$ 就是 $\text{argmax}_x Q(x|c)$)。
考虑一个词汇表 $\mathcal{V} = \{0, 1\}$,并让所有 $j$ 的 $x_{i_j}^* = 0$,所以 $x^* = (0, \ldots, 0)$。对于每个 $j \in \{1, \ldots, k\}$,令 $e_j$ 是在位置 $j$ 为1,其他位置为0的向量。令 $\delta = \epsilon - \frac{1}{k+1}(1 - \frac{1}{k+1}) > 0$。设置 $P(e_j|c) = \frac{1}{k+1} + \frac{1}{k}\delta, \forall 1 \le j \le k$ 并且 $P(x^*|c) = \frac{1}{k+1} - \delta$,那么 $x^* \notin \text{argmax}_x P(x|c)$。边际概率是:
因为
所以,边际条件 $P_j(x_{i_j} = x_{i_j}^*|c) > 1 - \epsilon$ (其中 $x_{i_j}^* = 0$) 成立。如上所示,$\text{argmax}_x P(x|c)$ 可以被构造成与 $x^*$ 不同。因此,如果 $\epsilon > \frac{1}{k+1}$,则 $P$ 和 $Q$ 的 argmax 可能不相同。
步骤4:界定 $L_p$ 距离。 令 $A_j$ 为事件 $\{X_{i_j} = x_{i_j}^*\}$。
项 $|P(\cap_{j=1}^k A_j|c) - \prod_{j=1}^k P(A_j|c)|$(使用 $P(A_j|c)$ 表示 $P_j(x_{i_j} = x_{i_j}^*|c)$)可以被界定。由于
因此,
对于 $x \neq x^*$:$P(x|c) < \epsilon$ 且 $Q(x|c) < \epsilon$。所以,
和 $\sum_{x \neq x^*} |P(x|c) - Q(x|c)|$ 可以被界定:
所以,
那么,
因此,
所以,
对于 $p=1$,
对于全变差距离,
步骤4:界定前向KL散度。
条件总相关性可以使用链式法则展开:
每一项都受条件熵的限制:
条件熵 $H(X_{i_j}|c)$ 是有界的。由于 $P_j(x_{i_j} = x_{i_j}^*|c) > 1 - \epsilon$,这意味着 $P_j(x_{i_j} \neq x_{i_j}^*|c) = p'_j < \epsilon$。当剩余概率 $p'_j$ 均匀分布时,熵最大化,导致:
将 $(k-1)$ 个这样的项相加(对于 $j = 2, \ldots, k$):
评注1。
- 关于定义明确的联合PMF的假设:该定理和证明依赖于 $P_{\theta}(x_{i_1}, \ldots, x_{i_k}|c)$ 是一个定义明确的联合概率质量函数,边际概率 $P_{\theta}(x_{i_j}|c)$ 是从中一致推导出来的。这意味着联合PMF是连贯的,并且其定义不依赖于超出条件 $c$ 所捕获的特定分解顺序。在实践中,虽然MDM可能不严格满足此属性,但其行为通常提供了很好的近似。该定理适用于具有这些性质的理想化 $P_{\theta}$。随着MDM变得越来越大、越来越强大,它们学习到的分布可能会更好地逼近这种一致性。
- 最坏情况分析:定理中提供的条件和界(例如 $(k+1)\epsilon \le 1$)是从最坏情况分析中得出的。这意味着如果满足条件,无论 $P(x|c)$ 的具体结构如何(除了高置信度边际属性),这些界都保证成立。实际上,真实情况可能比最坏情况“表现得更好”。例如,$x_{i_j}$ 和 $x_{i_l}$ 之间的依赖关系(给定 $c$)可能比最坏情况构造所假设的要弱。因此,即使 $(k+1)\epsilon$ 略大于1(但不是大很多),argmax等价性(结果1)可能仍然频繁成立。该条件确定了一个阈值,超过该阈值后,在最坏情况下保证会失效,但实际性能可能更具鲁棒性。同样,如果真实的联合 $P(x|c)$ 比最坏情况配置更接近边际乘积 $Q(x|c)$,那么实际的 $L_p$ 距离或KL散度可能比上界所暗示的要小。
B. 案例研究
表 6 | 不同方法响应的定性比较。
表 7 | 双缓存(DualCache)在不同块大小下响应的定性比较。
表 8 | 不同阈值设置下响应的定性比较。
B.1. 缓存策略对响应质量的影响
表6定性地比较了原始方法、前缀缓存(PrefixCache)和双缓存(DualCache)方法对算术提示的回答。所有方法都正确计算出答案(3个螺栓),并遵循相似的逐步推理,仅在措辞上有细微差异。这表明缓存策略在提高效率的同时,保持了答案的准确性和逻辑清晰度;语义保真度和可解释性未受影响。
B.2. 双缓存(DualCache)中块大小的影响
表7检验了双缓存中不同的块大小(8、16、32)。对于这个算术提示,所有设置都得出了正确、解释清晰的答案,输出没有有意义的差异。因此,对于此类问题,双缓存对块大小具有鲁棒性,允许在不牺牲质量的情况下提高效率。
B.3. 动态阈值设置的影响
表8研究了动态阈值(0.7、0.8、0.9)的影响。无论阈值如何,模型始终产生正确的答案和清晰的解释。虽然较高的阈值增加了计算量(函数评估次数NFE从9增加到20),但答案质量保持稳定,表明对于直接的算术问题,调整阈值主要影响效率,而非正确性。
B.4. 使用LLAda-V进行多模态生成
为了定性分析我们的Fast-dLLM框架在多模态场景中的有效性,我们进行了一个视觉案例研究,任务是让模型生成一个高度详细的图像描述。如图6所示,基线模型和我们的Fast-dLLM都接收了相同的视觉输入和用户提示:“请详细描述这张图片。”
用户: 请详细描述这张图片。
基线 (63.0 秒): 这张图片捕捉了一片宁静的乡村风光。主导场景的是一个木制谷仓,其风化的外观见证了时间的流逝。谷仓坐落在一片绿草地上,草地延伸至地平线。远处,一座山脉耸立,其山峰被云层笼罩。上方的天空是柔和的粉红色,暗示着黎明或黄昏。图片的视角是从低角度拍摄,使谷仓显得宏伟而壮观。图片的整体构图营造出一种深度和宁静感,邀请观众想象自己置身于这片宁静的乡村环境中。
Fast-dLLM (6.8 秒): 这张图片捕捉了一片宁静的乡村风光。主导场景的是一个木制谷仓,其屋顶上的灰色瓦片增强了其乡村魅力。谷仓坐落在一片高草丛中,草地向地平线延伸。这片田野被一座山脉环绕,山峰上覆盖着积雪,为景色增添了一丝宁静。上方的天空是柔和的粉红色,暗示着清晨或傍晚。图片的视角是从低角度拍摄,使谷仓显得宏伟而壮观。图片的整体构图创造了人造结构与自然世界的和谐融合。
图 6 | 基线与Fast-dLLM在视觉描述任务上的比较。 Fast-dLLM在极短的解码时间内生成了质量相当且忠实于图像的描述。
基线模型需要63.0秒来完成生成,产生了一个对乡村景观的详细而富有诗意的描述。它突出了风化的木制谷仓、柔和的粉红色天空和宁静的氛围等元素。
相比之下,我们的Fast-dLLM仅用6.8秒就完成了任务——速度提升了近10倍——同时保持了丰富的视觉细节。它进一步增强了描述,增加了额外的细节(例如,“屋顶上的灰色瓦片”,“一丝宁静”),反映了与图像外观和氛围的高度契合。值得注意的是,生成的描述保留了构图深度和文体流畅性,说明即使在基于扩散的并行解码下,模型也能平衡流畅性和事实性。
这个案例突显了LLAda-V与Fast-dLLM解码相结合,如何以显著提高的效率实现高质量的视觉-语言生成,为更快、更具交互性的多模态应用铺平了道路。
C. 实验细节
C.1. LLaDA-V的进一步实验
表 9 | 块长度对性能的影响(MathVista,48步)
表 10 | 在不同刷新间隔下,Fast-dLLM在MathVista上的性能(块长度=96)
在表9中,我们研究了在固定解码长度为48步的情况下,块长度的选择如何影响LLaDA-V在MathVista上的性能。结果显示,当块长度为96时,模型达到最高准确率。然而,当块大小减小到8或4时,准确率显著下降超过8%。
鉴于这种对块长度的敏感性,我们选择不将输出分成小块来单独更新缓存。相反,我们将块长度固定为96,并采用基于刷新的策略:缓存仅每隔 $k$ 个解码步骤使用最新的完整块进行更新。如表10所示,增加刷新间隔会导致吞吐量持续增加——从间隔为2时的15.9词元/秒增加到间隔为32时的28.2词元/秒。虽然准确率随间隔增大而略有下降,但仍保持在56.6%以上,这表明激进的刷新调度可以在仅有轻微性能下降的情况下带来可观的加速。
C.2. 阈值策略与因子策略的性能比较
表 11 | 在GSM8K和MATH基准测试上,生成长度为256和512时,阈值和因子置信度感知解码的性能比较。 每个块显示准确率(顶行)和吞吐量及加速倍数(底行)。在大多数设置中,因子解码提供了有利的权衡。
我们比较了基于阈值和基于因子的置信度感知并行解码策略在GSM8K和MATH基准测试上的性能(表11)。虽然阈值策略在大多数设置中实现了略高的准确率(例如,在256词元的GSM8K上为78.5%对77.5%),但因子策略展示了明显更优的吞吐量性能。
具体来说,在所有设置中,因子解码的吞吐量比阈值解码高出1.4-1.5倍。在256词元的GSM8K上,因子解码达到78.5词元/秒(11.7倍加速),而阈值解码为54.4词元/秒(8.1倍加速)。这种吞吐量优势在更长的生成任务中变得更加明显——对于512词元的GSM8K,因子解码达到47.1词元/秒,而阈值仅为35.3词元/秒。
结果表明,因子解码提供了一个引人注目的权衡:它牺牲了极小的准确率(通常为1-3%),换取了显著的吞吐量提升(高出40-50%)。这使得因子解码对于延迟敏感的应用特别有吸引力,在这些应用中,轻微的准确率降低是可以接受的。在两个基准和生成长度上的一致模式验证了因子策略理论基础的鲁棒性,该理论基础根据置信度界 $(k+1)\epsilon < \gamma$ 自适应地控制并行度。
C.3. LLaDA与LLaDA-1.5的比较
我们在GSM8K(5-shot)和MATH(4-shot)两个基准上,以及两种生成长度设置(256和512词元)下,比较了LLaDA及其增强版LLaDA-1.5的性能,如表12所示。每个单元格报告了准确率和解码吞吐量(词元/秒),以及相对于贪婪基线的相对加速。
在GSM8K的设置中,LLaDA-1.5始终比原始LLaDA提高了准确率,在256词元生成时实现了显著的+2.2%绝对增益,在512词元生成时为+3.2%。此外,它保持了强大的解码效率,在256词元时吞吐量达到59.4词元/秒,优于相同设置下LLaDA的54.1词元/秒。
在MATH基准上,两个版本之间的准确率相当。然而,LLaDA-1.5在256词元时略微提高了吞吐量(53.7对51.7),而在512词元设置下效率略有下降(41.1对47.1)。这表明,虽然LLaDA-1.5引入的增强功能对较短或中等解码上下文有利,但更长的序列可能需要进一步优化。
总的来说,LLaDA-1.5在不同设置下始终提供更高的准确率或更好的解码速度,展示了更好的性能-效率权衡,并突出了在基础LLaDA架构之上整合自适应改进的好处。
表 12 | LLaDA与LLaDA-1.5的性能比较。 每个单元格显示准确率和解码吞吐量(词元/秒),以及相对于LLaDA基线(底行,蓝色:词元/秒/橙色:相对加速倍数)的相对加速。
C.4. 解码步骤中并行词元数量分析
为了更好地理解基于因子的并行生成行为,我们分析了每个解码步骤中生成的平均词元数。具体来说,我们收集了采样过程中所有中间步骤的统计数据,并计算了每步并行生成的平均词元数。结果在图7中可视化,并附有95%置信区间,以显示样本间的变异性。
图 7 | 每个解码步骤生成的平均词元数。 蓝线显示平均词元数,阴影区域表示95%置信区间。
如图7所示,并行生成的平均词元数在解码的早期到中期阶段逐渐增加,大约在第30步到第60步之间达到峰值。此峰值后,并行度在生成结束时趋于略微下降。这表明模型在解码中期阶段对生成输出变得更有信心,使其能够同时产生更多词元。而在最后几步,解码过程趋于变得更加保守,减少了每步产生的词元数。
阴影置信区间揭示了在后期解码步骤中存在更大的方差,表明样本间的生成行为不稳定且不一致。这是预料之中的,因为末端解码步骤通常只处理完成输出所需的少数剩余词元,而不同样本中剩余词元数量可能差异很大(例如,由于提前完成或填充)。
这些观察对于理解如何优化解码效率非常重要:在高置信度阶段(中间步骤)增加并行度可以节省计算资源,而在边界附近采取保守行为则可以保持质量。
图 8 | (a) 使用我们的基于因子的解码策略,在不同因子值下的GSM8K(5-shot)准确率。 每个点上方的数字表示每步解码的平均词元数。虚线显示了基线方法每步2或4个词元,以及非并行(1词元/步)基线的准确率。(b) 每种因子设置下所需的相应推理步数。我们的方法通常比固定步长基线需要少得多的步数。(c) 在GSM8K(5-shot)上,准确率与每步解码的平均词元数的关系。我们的基于因子的解码相比基线实现了更好的准确率-效率权衡。红色的“Selected”点代表我们主要结果中选择的设置。
C.5. 不同批量大小下的吞吐量比较
所有实验均在NVIDIA A100 GPU上进行,前缀填充长度固定为256个词元。生成长度在16、32和64个词元之间变化,批量大小范围从1到32。此设置反映了实际部署场景,允许在不同条件下评估解码效率。
需要注意的是,并行解码允许同时生成多个受虚拟输入词元影响的词元。为确保公平性,我们仅关注缓存技术提供的加速。
PrefixCache被设计为LLaDA(一种基于扩散的LLM)的加速机制,并成功地显著提升了吞吐量。图9显示,PrefixCache在所有批量大小和生成长度上都实现了持续的改进,特别适用于生成长度较小和批量较大的场景。例如,当生成长度为16,批量大小为32时,PrefixCache的吞吐量超过211词元/秒,远超原生LLaDA仅达到的43词元/秒,展示了近5倍的提升。
图 9 | 在不同生成长度和批量大小下,PrefixCache、LLaDA和LLaMA的吞吐量比较。 所有模型均在NVIDIA A100 GPU上评估,前缀填充长度固定为256。
尽管LLaDA随着批量大小的增加其可扩展性有限——其吞吐量在批量大小为8后趋于平稳——这个限制是基于扩散的LLM固有的,它们本质上是计算密集型的。相比之下,LLaMA,一个自回归(AR)模型,从大批量中获益匪浅。随着批量大小的增加,LLaMA从内存密集型转为计算密集型,使其在更大的批量设置下能实现很高的绝对吞吐量。
这些结果突显了PrefixCache在加速像LLaDA这样的计算密集型扩散模型方面的实际优势,尤其对于延迟关键和高吞吐量应用。此外,PrefixCache提供的可扩展性和效率缩小了基于扩散的LLM与像LLaMA这样的AR模型之间的差距,展示了其在规模化部署设置中的重要性。
💬 评论讨论
欢迎在这里分享您的想法和见解!