MoESD: Unveil Speculative Decoding’s Potential for Accelerating Sparse MoE
MoESD: Unveil Speculative Decoding’s Potential for Accelerating Sparse MoE
作者/机构: Zongle Huang(1), Lei Zhu(2), Zongyuan Zhan(2), Ting Hu(2), Weikai Mao(2), Xianzhi Yu(2), Yongpan Liu(1,3†), Tianyu Zhang(2†∗)
1. 清华大学 (Tsinghua University)
2. 华为诺亚方舟实验室 (Huawei Noah’s Ark Lab)
3. 北京信息科学与技术国家研究中心 (BNRist)
A1 主要贡献
本文挑战了传统观念,即推测解码(Speculative Decoding, SD)无法有效加速专家混合(Mixture of Experts, MoE)模型。研究的核心洞见是,在中等批量大小(batch size)下,当所有专家都已被激活时,验证多个草稿令牌(draft tokens)不会带来额外的专家参数加载成本。此外,随着MoE模型变得更稀疏,每个专家在每次参数加载后处理的令牌更少,导致算术单元利用率降低,从而为SD创造了更大的加速机会。
核心贡献如下:
- 修正传统观念: 本文证明,在中等批量大小下,对于更稀疏的MoE模型,推测解码(SD)实际上比在密集模型中更有效,且其有效加速的批量大小范围更广。
- 建立可靠的性能模型: 基于理论分析,开发了一个可靠的SD加速比模型,使得加速过程透明且可解释。
- 引入新的系统性指标: 现有指标(如接受率)仅评估算法优化效率,无法完全解释SD的加速效果。本文引入了一个新的系统性指标“目标效率”(target efficiency),它揭示了目标模型固有的加速机会,有助于识别系统瓶颈。
- 发掘新的应用场景: 本研究为无损MoE加速提供了新视角,特别适用于私有化部署(private serving)等场景。在这些场景中,通常处理中等批量的请求。
- 实验验证: 在多种GPU上对Qwen2-57B-14A-Instruct模型进行的实验表明,SD在中等批量大小下可实现高达2.29倍的加速。实验结果也验证了SD对更稀疏的MoE模型更有利的理论预测。
A3 背景知识
预备知识。 LLM(大语言模型)的推理时间由计算和内存访问共同决定。当一个算子在GPU上处理时,内存访问和计算操作会进行流水线化和重叠,其中耗时更长的操作会成为瓶颈,决定了总处理时间,这由roofline模型【35, LLM inference unveiled: Survey and roofline model insights by Zhihang Yuan et al., 2024】、【36, Applying the roofline model by Georg Ofenbeck et al., 2014, IEEE International Symposium on Performance Analysis of Systems and Software (ISPASS)】所描述。硬件的roofline模型脊点(Ridge Point, RP)和软件的算术强度(Arithmetic Intensity, AI)定义如公式1所示。当AI < RP时,系统是内存受限的(memory-bound),增加更多计算不会显著增加处理时间。当AI > RP时,系统是计算受限的(compute-bound),计算量的增加会直接反映在处理时间上。在本文中,当我们描述一个系统“更受内存限制”时,我们指的是$AI/RP$更小。
$$ \mathbf{RP} = \frac{\text{peak computation power } (\textit{unit: Flops})}{\text{peak memory bandwidth } (\textit{unit: bytes/second})} \quad \mathbf{AI} = \frac{\text{computation operation } (\textit{unit: times})}{\text{memory access volume } (\textit{unit: bytes})} $$A2 方法细节
3.1 推测解码加速比与目标效率的公式化
推测解码处理时间的公式化。 我们首先将推测解码的处理时间形式化,记为$T_{SD}$。为了生成长度为S的序列,SD会经历R轮,每轮包含三个阶段:① 草稿模型根据推测策略提出$\gamma$个令牌;② 目标模型验证这些令牌;③ 拒绝采样【9, Accelerating large language model decoding with speculative sampling by Charlie Chen et al., 2023】根据目标模型和草稿模型的logits丢弃错误预测的令牌。我们使用$T_T(b, s)$和$T_D(b, s)$分别表示目标模型和草稿模型单次前向传播的时间,其中b和s是批量大小和待处理令牌数的形参。因此,处理包含B个请求的批次所需的时间由公式2给出:
$$T_{SD} = R \times (T_{propose} + T_{verify} + T_{reject}) = R \times \left( \gamma \cdot T_{D}(B, 1) + T_{T}(B, \gamma) + T_{reject} \right)$$推测解码加速比的公式化。 于是,SD相对于正常自回归解码$T_{AR}$的加速比如公式4所示:
$$\begin{aligned} \begin{aligned} Speedup & = \frac{T_{AR}}{T_{SD}} = \frac{S \cdot T_{T}(B, 1)}{R \cdot\left(\gamma \cdot T_{D}(B, 1) + T_{T}(B, \gamma) + T_{reject}\right)} \\ & = \frac{S}{R} \cdot \frac{1}{\gamma \cdot \frac{T_{D}(B, 1)}{T_{T}(B, 1)} + \frac{T_{T}(B, \gamma)}{T_{T}(B, 1)} + \frac{T_{reject}}{T_{T}(B, 1)}} \end{aligned} \end{aligned}$$接受令牌长度与接受率的关系。 $S/R$代表每轮SD平均接受的令牌长度,可以进一步表示为$\sigma \times (\gamma + 1)$。这里,$\sigma$是实际生成的令牌数与所有草稿令牌都被接受情况下的理论最大值之比。我们注意到,$\sigma$不同于先前工作【10, Fast inference from transformers via speculative decoding by Yaniv Leviathan et al., 2023, International Conference on Machine Learning】、【9, Accelerating large language model decoding with speculative sampling by Charlie Chen et al., 2023】、【7, Eagle: Speculative sampling requires rethinking feature uncertainty by Yuhui Li et al., 2024】中常引用的接受率$\alpha$,$\alpha$表示给定前缀下目标模型接受一个新草稿令牌的概率。$\sigma$可以由$\alpha$计算得出,如公式5所示。分子遵循【10, Fast inference from transformers via speculative decoding by Yaniv Leviathan et al., 2023, International Conference on Machine Learning】的推导,分母则考虑了所有$\gamma$个草稿令牌被接受,外加验证前向传播过程中生成的一个额外令牌。
$$\sigma = \frac{expected\ generated\ tokens}{maximal\ possible\ accepted\ tokens} = \frac{\frac{1-\alpha^{\gamma+1}}{1-\alpha}}{\gamma+1}$$加速比公式的分析。 公式4的分母由三项组成。$T_D(B,1)/T_T(B,1)$是草稿模型前向时间与目标模型前向时间之比,反映了草稿模型和目标模型的相对体量。这个值通常保持很小(通常小于1/10【27, Specinfer: Accelerating large language model serving with tree-based speculative inference and verification by Xupeng Miao et al., 2024, Proceedings of the 29th ACM International Conference on Architectural Support for Programming Languages and Operating Systems】、【7, Eagle: Speculative sampling requires rethinking feature uncertainty by Yuhui Li et al., 2024】、【10, Fast inference from transformers via speculative decoding by Yaniv Leviathan et al., 2023, International Conference on Machine Learning】)以确保推测是高效的。$T_{reject}/T_T(B,1)$更小,因为$T_{reject}$只涉及采样而非模型推理。$T_T(B,\gamma)/T_T(B,1)$是这三项中值最大的,它显著影响最终的加速比。如公式4所示,它的增加会导致加速比下降。有两个不同因素导致其增加,分别解释了SD在(1)大批量(对密集模型和MoE均适用)和(2)小批量MoE下效率低下的原因:
(1) 计算受限性。 当系统更受内存限制时(批量大小B较小),模型的$T_T(B,\gamma)/T_T(B,1)$接近1;但当系统更受计算限制时(批量大小B较大),该比值会增加到$\gamma$。
(2) 额外的内存加载。 对于较小的B,由于激活了更多的专家需要加载,$T_T(B, \gamma)$显著大于$T_T(B, 1)$。由于在小批量下系统仍是内存受限的,内存加载深刻地决定了处理时间。
目标效率的定义。 因此,我们将目标效率定义为$T_T(B,1)/T_T(B,\gamma)$,这个指标有助于理解上述SD加速性能下降的系统性原因。我们在图2中的实验表明,目标效率始终能反映SD加速比的变化趋势。尽管这个值很重要,但之前的工作很少注意到它,这主要是由于研究重点的不同。
研究问题的转变。 以前的SD研究主要通过提高接受率来解决以下问题:
给定目标模型,哪种草稿模型或算法能实现更好的加速?
相比之下,我们的工作通过考察目标效率来关注以下问题:
在相同水平的算法优化下,哪种类型的目标模型或工作负载对SD更有利?
我们相信目标效率有助于研究人员更全面地理解SD。即使目标-草稿对具有相同的接受率$\alpha$,目标模型架构和工作负载的变化也会显著影响整体加速比。通过引入目标效率,我们可以将算法优化与系统优化解耦,从而帮助识别系统性加速瓶颈并评估潜在的加速比。
3.2 中等批量大小为MoE的推测解码加速提供了可能
核心论点。 尽管SD在小批量下对MoE无效,但本小节我们证明,在中等批量大小——一个先前研究忽略的区间——SD的加速比会增加,并且稀疏度越高的MoE获益越多。本质上,当批量大小处于所有专家都已激活但尚未分配足够工作负载的范围内时,FFN(前馈网络)会变得内存受限,这为通过SD几乎免费地利用计算能力提供了机会。为了证明这一点,我们首先公式化激活专家的期望数量,然后展示随着模型变得更稀疏,MoE的FFN会变得更加内存受限。
激活专家数量的期望值。 我们使用伯努利随机变量$X_i$来表示专家的激活状态:如果专家i被激活,则$X_i = 1$,否则为0。为简化起见,我们假设$X_i$是独立同分布的(i.i.d.)。那么,激活专家的期望数量N可以表示为公式6,其中E是专家总数,$Pr(X_i)$表示第i个专家被激活的概率。
$$N = \sum_{i} \mathbb{E}[X_{i}] = \sum_{i} Pr(X_{i}) = E \cdot Pr(X)$$专家激活概率与令牌数的关系。 给定t个令牌通过MoE门,那么$Pr(X)$可以表示为公式7。K表示每个令牌激活的专家数量,是MoE的一个架构超参数:
$Pr(X) = 1 - Pr(\text{t个令牌都未激活该专家}) = 1 - (\frac{E - K}{E})^t$
因此,$N(t)$的总体表达式由公式8给出。我们的推导假设专家被均匀激活,这对于训练良好的MoE模型是合理的。专家间的负载不均衡会导致路由崩溃【5, Outrageously large neural networks: The sparsely-gated mixture-of-experts layer by Noam Shazeer et al., 2017】并降低专家并行部署的计算效率【6, Gshard: Scaling giant models with conditional computation and automatic sharding by Dmitry Lepikhin et al., 2020】,因此,最先进的MoE模型通常采用如加入辅助损失【37, Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity by William Fedus et al., 2022】、【6, Gshard: Scaling giant models with conditional computation and automatic sharding by Dmitry Lepikhin et al., 2020】等方法进行训练,以确保专家负载均衡。实验结果也验证了我们对$N(t)$的理论推导,如图1a和1b所示。
全激活的令牌阈值。 接着我们求解需要多少令牌才能导致完全激活。由于当t趋于无穷时,$N(t)$渐近地接近E,而在实践中$N(t)$应为有限整数,我们将$N(t) > \tau E$视为几乎完全激活,其中$\tau$通常是一个较大的比例,如0.95。我们进一步表示$K = \rho E$,其中$\rho$是MoE的稀疏度,那么令牌阈值$T_{thres}$可以通过公式9求解:
$$N(T_{thres}) = E \cdot \left( 1 - (1 - \rho)^{T_{thres}} \right) \geq \tau E \Rightarrow T_{thres} = \lceil \log_{(1-\rho)} (1 - \tau) \rceil$$稀疏MoE的内存受限特性。 因此,当B超过$T_{thres}$时,激活的专家数量饱和,导致验证阶段的$B\gamma$个令牌只会引起边际上更大的内存访问。解决了3.1节中分析的$T_T(B,\gamma)/T_T(B,1)$增加的第二个因素(即额外的内存加载)后,我们现在转向由第一个因素——计算受限性——引起的潜在限制。如果这样的B值使得系统进入计算受限状态,SD同样无法有效加速MoE。我们对这个问题的回答是:更稀疏的MoE会延迟随着输入令牌数增加而从内存受限向计算受限的转变。
每个专家的平均工作负载。 我们已经得到,给定t个令牌,会有$N(t)$个专家被激活。由于每个令牌激活K个专家,每个专家平均需要处理的令牌数$\overline{T_{exp}}$可以计算为:
$$\overline{T_{exp}}(t ; \rho)=\frac{t \cdot K}{N}=\frac{t \cdot(\rho E)}{E \cdot\left(1-(1-\rho)^{t}\right)}=\frac{\rho t}{1-(1-\rho)^{t}}$$稀疏度对专家负载和系统性能的影响。 正如附录中证明并在图1c中所示,给定$t = T > 1$,$\overline{T_{exp}}(T; \rho)$随着$\rho$的减小而减小,这表明随着MoE变得更稀疏,每个专家在每次参数加载时处理的令牌更少。因此,运行更稀疏MoE的系统更受内存限制,导致算术单元的利用率较低。验证阶段因此可以利用这些闲置资源,而不会显著增加处理时间。相比之下,密集模型是$\rho = 1$的极端情况,其FFN始终接近最大可能的算术强度T,随着T的增加,系统会迅速过渡到计算受限区域。
假设前提。 我们应该注意,我们的结论是基于MoE FFN部分在整个模型中占比较大的前提,这对于当前大部分参数都是专家的大多数MoE模型是成立的。在一个假设的极端情况下,如果Attention层占主导地位而MoE FFN可以忽略不计,那么根据阿姆达尔定律(Amdahl's Law),MoE的稀疏性对整体系统性能的影响将非常有限。
3.3 推测解码加速比的建模方法
建模的必要性。 鉴于影响最终加速比的因素众多,定量理解每个因素的影响具有挑战性。因此,我们开发了一种建模方法,使SD加速比的结果更具可解释性和透明性。如公式4所示,建模SD加速比的核心在于刻画模型前向传播的时间。基于前几节的理论分析,我们确定了影响前向执行的三个主要因素:(1) roofline模型效应,(2) 激活专家的数量,以及(3) 专家负载。由于GPU的实际执行是动态的,且并非所有算子都优化到了其理论极限,我们引入了几个松弛参数。这些参数的值随后通过拟合GPU测量数据自动确定。这些因素及其对执行时间的影响分析如下。
(1) Roofline模型效应。 该效应表现为执行时间随着令牌数t的增加而增加,其增长率先慢后快,最后趋于稳定。根本原因如下。当t较小时,参数加载时间超过计算时间,形成内存瓶颈。因此,给定参数量,内存访问时间是稳定的(内存受限区域)。随着t的增加,计算时间超过参数加载时间并成为瓶颈。由于硬件中的算术单元是固定的,计算时间与计算负载成线性关系(计算受限区域)。为了刻画这一趋势,我们设计了$G(t; \lambda RP, s)$,如公式11所示,其中$\lambda RP$表示区域之间的转换点,s控制执行时间的增长率。这里,RP遵循公式1,$\lambda$是一个小于1的常数,用于解释内存带宽利用率的实际限制。$G(t)$在转换点前表现出逐渐增加的斜率,之后转变为线性函数,并在转换点保持一阶梯度连续性。
$$\begin{aligned} G(t ; \lambda R P, s)=\begin{cases}s^{t}, & t \leq \lambda R P \\ s^{\lambda R P}+\left(\left.\frac{\mathbf{d}\left(s^{t}\right)}{\mathbf{d} t}\right|_{t=\lambda R P}\right)(t-\lambda R P)=s^{\lambda R P}(1+l n(s) \cdot(t-\lambda R P)), & t>\lambda R P\end{cases} \end{aligned}$$(2) 激活专家的数量。 当激活的专家数量增加时,内存访问量增加,从而增加了最终处理时间。我们使用推导出的公式8中的N来刻画工作负载和模型架构如何影响激活专家的数量。
(3) 专家负载。 这指的是在通过MoE门进行令牌分配后,每个专家只处理令牌的子集$\overline{T_{exp}}(t; \rho)$,而不是全部输入令牌t。因此,在将roofline模型应用于MoE专家时,我们应该使用$G(\overline{T_{exp}})$而不是$G(t)$。这证实了我们的理论结论,即更稀疏的MoE在输入令牌数增加时会延迟从内存受限到计算受限的转变。
模型综合表达式。 对于MoE目标模型,因素(1)、(2)和(3)都参与其中。我们以一阶方式将它们组合起来,并引入参数bias、$k_1$、$k_2$和$k_3$来调整实际GPU执行中的非理想因素,完整表达式见算法1的第6和第8行。这些参数具有明确的实际意义:bias表示加载固定参数所需的时间;$k_2 \cdot N$表示加载N个激活专家所需的时间;$k_1 \cdot G(t)$和$k_3 \cdot G(\overline{T_{exp}})$描述了随着令牌数增加,执行时间的增量趋势。对于草稿模型,由于它通常是密集的,只涉及因素(1),其建模形式见算法1的第9行。
1: 测量输入: 共m个测量值,记为M。每个Mi, i = 1, 2, ..., m包含属性:批量大小B,草稿长度γ,每个令牌激活的专家数K,专家总数E,接受令牌数与最大可能接受令牌数之比σ,实际达到的加速比Speedup。
2: 输出: 最优拟合参数params*。
3: def ComputeSpeedup(params, B, γ, K, E, σ): ▷ 计算SD加速比
4: bias, k1, k2, k3, draft_bias, draft_k, reject_bias, reject_k, λ, s = params ▷ 解包参数
5: Nar = E · (1 − ((E − K)/E)^B), Tar_bar = B · K/Nar ▷ 计算AR前向时间
6: ar_time = bias + k1 · G(B; λRP, s) + k2 · Nar + k3 · G(Tar_bar; λRP, s)
7: Nsd = E · (1 − ((E − K)/E)^(Bγ)), Tsd_bar = B · γ · K/Nsd ▷ 计算SD前向时间
8: verify_time = bias + k1 · G(Bγ; λRP, s) + k2 · Nsd + k3 · G(Tsd_bar; λRP, s)
9: draft_time = draft_bias + draft_k · G(B; λRP, s) ▷ 计算草稿模型前向时间
10: reject_time = reject_bias + reject_k · B ▷ 计算拒绝采样时间
11: Speedup = σ · (γ + 1) · ar_time / (draft_time + ar_time + verify_time + reject_time) ▷ 按公式4计算加速比
12: return Speedup
13: params* = argmin Σ (ComputeSpeedup(params, Mi.B, Mi.γ, Mi.K, Mi.E, Mi.σ) − Mi.Speedup)^2 ▷ 通过最小二乘法拟合模型到测量输入,确定最优参数params*。
参数拟合。 在确定了SD加速比的表达式后,我们拟合测量输入以自动确定松弛参数的值,优化标准是最小化模型输出与真实值之间的均方误差(MSE),如算法1的第13行所示。通过在我们的模型中应用这些优化后的参数(即第3行的ComputeSpeedup函数),我们获得了完整的SD加速比模型。此过程的示意图和更多拟合细节在附录C中提供。
模型的可靠性与价值。 由于我们的理论分析捕捉了主要的权衡关系并为建模提供了坚实的基础,拟合过程非常高效。图4中展示了使用21个测量值进行拟合的结果,该结果在各种情况下都与GPU结果的趋势一致。这些结果验证了我们建模的可靠性,从而使其成为一个分析模型前向传播各组成部分、定量理解不同因素之间权衡的有效工具。如4.2节所示,我们借助该模型解释了一些意想不到的结果。
3.4 理论发现的实际价值
理论发现的总结。 尽管前面的章节侧重于理论分析,本节将展示这些发现如何转化为实际的加速效果。我们的理论分析已经揭示,SD对MoE的加速效果在中等批量大小时最为显著,其趋势是先增后减。我们从基础部署和扩展配置两个方面讨论其实际价值。
基础部署的价值。
(1) 私有化服务场景: 中等批量大小在私有化服务中很常见,这类服务因数据安全而日益普及,代表性应用如企业内部聊天机器人。
(2) 延迟敏感场景: 当延迟要求严格时,大批量通常不可行。LLM服务必须满足多个服务水平目标(SLOs)【38, Revisiting slo and goodput metrics in llm serving by Zhibin Wang et al., 2024】,包括首个令牌时间(TTFT)和每输出令牌时间(TPOT)。大批量会减少每个请求的计算资源,导致延迟违规。在这种情况下,中等批量大小是常见的。
(3) 放宽延迟-吞吐量权衡: 我们的工作实际上揭示了SD在MoE上放宽了传统的延迟-吞吐量权衡。具体来说,MoE模型存在一个区间,其中SD加速比增加(延迟更低)的同时批量大小也在增加(吞吐量更高)。
从模型的角度看,中等批量大小代表了MoE的“效率差距”。在这个规模下,所有参数都必须加载(不同于小批量时的选择性专家激活),但GPU的FLOPs并未被充分利用(不同于大批量)。我们的发现为在不牺牲模型质量的情况下解决这一效率挑战提供了新颖的视角。
扩展配置下的价值。 我们考虑对MoE的典型系统优化,如卸载(offloading)和专家并行(expert parallelism, EP)。
* 卸载(Offloading): 当MoE模型超出GPU内存容量时,FFN参数会被卸载到CPU内存【39, Ktransfromers: A flexible framework for experiencing cutting-edge llm inference optimizations by kvcache ai, 2025】。这会将参数加载带宽从GPU内存带宽降低到低得多的PCIe带宽,使系统更加内存受限。因此,额外的计算不会显著增加处理时间,为SD创造了有利条件。值得注意的是,现有的优化如专家预取【23, Moe-infinity: Activation-aware expert offloading for efficient moe serving by Leyang Xue et al., 2024】、【24, Adapmoe: Adaptive sensitivity-based expert gating and management for efficient moe inference by Shuzhang Zhong et al., 2024】和缓存【25, Expertflow: Optimized expert activation and token allocation for efficient mixture-of-experts inference by Xin He et al., 2024】、【26, Hobbit: A mixed precision expert offloading system for fast moe inference by Peng Tang et al., 2024】在中等批量大小下效率会降低,因为几乎所有专家都被激活了。
* 专家并行(EP): 我们的发现也与EP兼容。在EP中,专家分布在多个GPU上,这既不影响$N(t)$也不影响$\overline{T_{exp}}$,使我们之前的分析仍然有效。由于MoE FFN之外的组件也进行了并行化,MoE FFN仍然构成处理时间的很大一部分,使得内存受限效应在端到端性能中仍然可以观察到。值得注意的是,在广泛的EP配置下,由于大量EP GPU提供了额外的内存带宽,SD在小批量MoE上的低效率可能会消失。
A4 实验环境与结果
实验环境
-
模型与数据集:
-
MoE模型对:
- 目标模型:Qwen2-57B-A14B-Instruct,草稿模型:Qwen2-0.5B-Instruct【40, Qwen2 technical report by An Yang et al., 2024】。
- 目标模型:Mixtral-8x7B-Instructv0.1【41, Mixtral of experts by Albert Q. Jiang et al., 2024】,草稿模型:Eagle speculation head【7, Eagle: Speculative sampling requires rethinking feature uncertainty by Yuhui Li et al., 2024】。
-
密集模型对比: 目标模型:Opt-30b,草稿模型:Opt-350m【42, Opt: Open pre-trained transformer language models by Susan Zhang et al., 2022】。
- 稀疏度调整: 通过修改模型
config.json文件中的num_experts_per_token参数来改变稀疏度。 - 数据集:
- HumanEval【43, Evaluating large language models trained on code by Mark Chen et al., 2021】(代码生成任务),提示长度范围为38到391个令牌。
- MT-bench【44, Judging llm-as-a-judge with mt-bench and chatbot arena by Lianmin Zheng et al., 2023】(对话任务),提示长度范围为5到356个令牌。
-
-
硬件配置:
- 2xGPU-A, 2xGPU-B, 4xGPU-A, 4xGPU-C。
-
软件配置:
- 框架: 使用现有的vllm【46, Efficient memory management for large language model serving with pagedattention by Woosuk Kwon et al., 2023】框架进行实验。
- 数据采集: 所有数据均为10次运行中最后5次结果的平均值,以避免初始性能不稳定的影响。
实验结果
MoE的推测解码加速趋势
-
实验内容: 评估MoE模型在不同设置下(硬件、模型、数据集、温度、$\gamma$值)的端到端SD加速比随批量大小的变化趋势。
-
实验结果与分析:
- 加速比趋势(图2): 在各种设置下,SD加速比都呈现出随批量大小增长而先增后减的趋势,验证了理论预测。
-
峰值加速比(表1):
- Qwen2在HumanEval(temp=0.0, $\gamma=4$)上达到2.18倍的峰值加速。
- Mixtral在HumanEval(temp=0.0, $\gamma=4$)上达到1.79倍的峰值加速。
- 对于更具可预测模式的任务(如代码生成)或更少随机性(如更低的温度),使用更长的$\gamma$可以实现更高的加速,这与先前研究一致。
-
硬件影响(表2):
- 具有更高脊点(ridge point)的GPU(如GPU-A vs GPU-B,GPU-A vs GPU-C)能产生更大的SD加速比,因为它们为验证提供了更多的算术单元。Qwen2在2xGPU-B上达到2.29倍的峰值加速。
- 从2xGPU-A扩展到4xGPU-A虽然降低了绝对运行时间,但SD加速比略有下降。这是因为大模型受益于GPU间的并行化,而小的草稿模型仍然在单个GPU上运行,使其相对前向成本更高。
-
目标效率指标的有效性(图2): “目标效率”(右y轴)的变化趋势与端到端加速比(左y轴)高度一致,验证了该指标的有效性。相比之下,接受率在不同批量大小下仅在小范围内波动,无法有效解释加速比的剧烈变化。
- 与密集模型的对比(图3):
- MoE模型的目标效率随批量大小先增后减,而密集模型的目标效率则持续下降。
- 这表明,尽管SD在小批量下对MoE效果较差,但在更广泛的较大批量大小范围内,SD对MoE具有更强的潜力。
MoE稀疏度的影响与建模方法验证
-
实验内容: 通过改变Qwen2-57B-A14B-Instruct中每个令牌激活的专家数量(K)来评估MoE稀疏度对SD加速的影响。为了公平比较,对原始加速比进行了校正(乘以$\sigma_{K=8} / \sigma$)。同时,将实验结果与我们提出的性能模型进行对比。
-
实验结果与分析(图4):
- 模型的可靠性: 性能模型的结果与在不同稀疏度($\rho$)和草稿长度($\gamma$)下的实验结果始终保持一致,验证了我们建模方法的可靠性。
- 极稀疏MoE的特殊情况: 对于非常稀疏的MoE(K=1, 2),SD加速比呈现持续下降的趋势。通过模型分析发现,这是因为这些模型中FFN的比例过低,使得Attention层成为主导,MoE FFN的内存受限特性难以在系统层面体现(阿姆达尔定律)。这是一种人为构造的情况;在实践中,更稀疏的MoE通常会配置更多的FFN参数以维持平衡。
- 稀疏度对加速范围的影响: 随着MoE模型变得更稀疏($\rho$更小),系统从内存受限到计算受限的转换被延迟。这体现在两个现象上:(1) 达到最大加速比(x)所需的批量大小更大;(2) 维持较高加速比(如图中棕色虚线所示的$x/\sqrt{2}$)的批量大小范围更宽。这验证了我们的理论分析,并表明SD在更稀疏的MoE中具有更广泛的适用性。
A5 结论
本研究挑战了推测解码(SD)不能有效加速MoE模型的传统观念,并指出在中等批量大小下,由于FFN变得更加内存受限,更稀疏的MoE模型实际上能从SD中获得更大的益处。我们通过理论分析和实验验证支持了这一结论。考虑到影响最终加速比的多个因素之间复杂的相互作用,我们开发了一个可靠的SD性能模型,使我们能够以透明和可解释的方式理解加速过程。我们还引入了“目标效率”这一指标,以帮助研究人员全面理解目标模型架构和工作负载如何影响SD加速。我们的工作为MoE加速提供了一个新的视角,尤其适用于具有中等批量大小的私有化服务或内存受限的场景。我们的大部分分析假设KV缓存的体积相对小于模型参数,而当KV缓存占主导地位时SD的行为已经被MagicDec【34, Magicdec: Breaking the latencythroughput tradeoff for long context generation with speculative decoding by Ranajoy Sadhukhan et al., 2024】分析过。这两项工作可以结合起来,为不同批量大小下的SD提供一个更全面的视图。
A6 附录
A 补充实验结果
A.1 更多配置下的SD加速趋势
在更多配置下的SD加速比趋势。 图5展示了在不同数据集、温度和模型类型下SD加速比的更多趋势,作为图2的补充。结果表明,SD加速比呈现出一致的先增后减模式,这与我们的理论分析非常吻合。
统计显著性。 为了确认我们发现的统计显著性,我们还展示了构成图5中平均值的五次独立运行。不同运行之间的方差很小,这是预期的,因为所有运行的随机种子都是固定的,以确保工作负载相同。
局部波动现象的解释。 尽管总体趋势遵循先增后减的模式,但在曲线上可以观察到局部波动。例如,图5(c)呈现出锯齿状的下降趋势。这种现象可归因于GPU的量化效应,正如NVIDIA的文档【47, Matrix multiplication background user’s guide by NVIDIA Docs Hub】中所记录的。当维度不能被GPU的原生tile大小整除时,计算性能会下降。自回归(AR)解码比SD对这种效应更敏感,使得AR与SD的时间比(即SD加速比)发生波动。尽管存在这些局部变化,但总体加速趋势遵循我们的理论预测,证实了我们结论的有效性。
A.2 MoE与密集模型的端到端加速比比较
端到端加速比的补充比较。 为了隔离接受率变化的影响,并更清晰地关注系统瓶颈,我们在4.1节中使用了目标效率来比较MoE和密集模型。本节中,我们进一步在图6中比较了它们在各种设置下的端到端加速比作为补充。
实验观察。 实验结果揭示了两个关键观察点。首先,MoE模型的SD加速比先增后减,而密集模型的SD加速比则持续下降。因此,在中等批量大小下,SD为MoE模型实现了更显著的端到端加速,这与4.1节图3的趋势一致。其次,SD对MoE的偏好程度在不同配置下有所不同。例如,在温度=1(第二行)时,与温度=0(第一行)相比,SD对MoE显示出更大的相对优势。这种差异源于不同设置下不同的接受率,这可能会掩盖对系统瓶颈的观察。总之,目标效率在控制算法优化等混杂因素的同时,可以作为一个可靠的比较指标。
B $\overline{T_{exp}}(T; \rho)$随$\rho$变化的趋势证明
证明目标。 在3.2节中,图1(c)表明:给定输入令牌数$t = T > 1$,每个专家平均处理的令牌数$\overline{T_{exp}}(T; \rho) = \frac{\rho T}{1-(1-\rho)^T}$会随着$\rho$的减小而减小。我们通过证明当$T > 1$时,$\frac{d(\overline{T_{exp}}(T; \rho))}{d\rho} > 0$来证明这一点。
求导过程。
等价命题。 由于$\rho$表示MoE的稀疏度,属于(0, 1),原命题等价于证明:
证明思路。 注意到当$\rho \rightarrow 0$时,$F(\rho; T) \rightarrow 1$。因此,如果我们能证明当$\rho$从0增加到1时,$F(\rho; T)$是减函数,那么原命题就得证。我们通过计算$\frac{d(F(\rho;T))}{d\rho}$来证明这一点:
结论。 当$T > 1$时,$\frac{d(F(\rho;T))}{d\rho} < 0$。这证实了$F(\rho; T)$随着$\rho$的增加而减小,从而证明了我们的原始命题:当$T > 1$时,$\overline{T_{exp}}(T; \rho)$随着$\rho$的减小而减小。
C 关于建模方法的更多细节
建模方法的价值。 我们的建模方法的价值体现在两方面。一方面,它仅用少量简单的参数就实现了与真实测量值的一致,从而验证了我们理论分析的正确性。另一方面,它提供了端到端结果中各种因素的分解,使得整个SD加速过程可解释且透明。
C.1 建模过程的描述与概述
建模流程。 图7展示了我们建模方法的整体流程图。基于3.1和3.2节的理论分析,我们推导出SD加速比与工作负载的函数表达式。这个表达式包含几个待定的松弛参数。我们通过经验性剖析(profiling)来确定这些参数。首先,我们收集一小组包含各种工作负载及其相应SD加速比的真实测量数据。然后,我们使用这些测量数据,在最小二乘准则下进行参数拟合,以获得最优参数值。由于我们理论推导的SD加速比表达式已经捕捉了基本的性能权衡,拟合过程计算开销小且鲁棒。一旦获得最优参数,得到的表达式就可以预测任意工作负载下的SD加速比。
剖析的必要性。 我们现在解释为什么剖析是必要的,以及为什么我们不能纯粹通过理论分析推导出完整的SD加速比表达式,这涉及软件和硬件两方面的考虑。
- 软件考虑: 对于具有不同实现的复杂算子,实际执行时间可能与理论预测有很大偏差。在GPU上,GEMM操作因其规则结构和高度优化的实现而确实可预测。然而,对于像Attention这样的算子,预测变得具有挑战性,因为它涉及定制的核优化(如FlashAttention1/2、eager attention、SDPA attention)和算子融合策略。例如,Qwen2-57B-A14B(隐藏层大小3584)和Mixtral-8x7B(隐藏层大小4096)的剖析结果显示,FFN耗时与隐藏层大小一致,但Attention耗时却与理论预期相反。
- 硬件考虑: 不同系列GPU的微架构各不相同,这会极大地影响执行时间。例如,Attention的效率在很大程度上依赖于硬件感知的编程优化,而不同GPU在缓存配置和内存层次结构中的线程-内存交互模式上有所不同。FlashAttention-3【48, Flashattention-3: Fast and accurate attention with asynchrony and low-precision by Jay Shah et al., 2024】正是利用了现代硬件的新功能,成功地将H100 GPU的利用率从35%提高到75%。此外,许多关键的硬件细节GPU厂商并未公开,使得理论预测不切实际。
因此,我们的超参数方法提供了一个更通用且易于使用的范式,在有效性和实用性之间取得了平衡。这种方法也被其他系统优化框架如NanoFlow【49, Nanoflow: Towards optimal large language model serving throughput by Kan Zhu et al., 2025】所采用。
C.2 图4中建模的拟合细节
测量数据选择。 我们从总共228个GPU测量值(跨越6个不同的K值,2个$\gamma$值和19个B值)中,通过固定步长(stride=11)进行均匀采样,选择了21个测量值。这种采样策略确保了所选测量值包含不同的设置,使建模更具鲁棒性。
拟合方法与参数边界。 我们使用了scipy.optimize.least_squares函数和信赖域反射(TRR)算法来优化非线性的SD加速比函数。拟合过程高效,约0.114秒完成。我们的模型包含10个需要松弛的参数,其搜索边界设定如下:
* bias: 目标模型密集参数的加载时间。下界由理论最小加载时间确定,上界设为下界的5倍。
* k1: 密集组件roofline效应的强度。下界为0,无确定上界。
* k2: 加载一个专家的时间。下界由理论最小加载时间确定,上界设为下界的5倍。
* k3: 稀疏组件roofline效应的强度。下界为0,无确定上界。
* draft_bias: 加载密集草稿模型的时间。下界由理论最小加载时间确定,上界设为下界的5倍。
* draft_k: 草稿模型roofline效应的强度。下界为0,无确定上界。
* reject_bias: 拒绝采样的固定开销。边界设为[0, $T_{rej}$],其中$T_{rej}$是测量的最大拒绝采样时间。
* reject_k: 拒绝采样中随输入令牌数增加的增量处理时间。边界设为[0, $T_{rej}$]。
* λ: 经验脊点与理论脊点之比。边界设为[0.2, 1]。
* s: 调整执行时间随输入令牌数增加的增长率。边界设为[1, 2]。
C.3 备选测量选择的探索
不同测量数量的影响。 本节我们展示了用于拟合的测量数量(m)对建模结果的影响。由于模型有10个参数,至少需要10个剖析数据点(m ≥ 10)。我们展示了m从10到228的建模拟合结果。
结果分析。 如表3及附图8-28所示,通常情况下,模型与真实测量值拟合得很好,除了m=10、12、13。原因如下:
* 当m=10时,测量数量等于参数数量,数据不足以进行鲁棒拟合。
* 当m=12和m=13时,测量数据的分布存在偏差。基于步长的选择导致m=12的测量中缺少大于40的批量大小,而m=13的测量中缺少1-24范围内的批量大小。尽管m=11的数据点更少,但其MSE值更小。
基于此分析,我们建议在选择测量数据时优先考虑数据的均匀分布,因为这种方法即使使用较小的数据集也能开发出更可靠的模型。
... (后续图12至28的引用被省略以保持简洁)
💬 评论讨论
欢迎在这里分享您的想法和见解!