EPS-MoE：用于成本高效MoE推理的专家流水线调度器
作者/机构: Yulei Qian, Fengcun Li, Xiangyang Ji, Xiaoyu Zhao, Jianchao Tan, Kefeng Zhang, Xunliang Cai (美团)

A1 主要贡献

核心问题与研究目标
大型语言模型（LLM）领域中，专家混合（MoE）模型已成为一种重要架构，它在模型性能和计算效率之间提供了更好的平衡。然而，MoE模型中的通用矩阵乘法（GEMM）操作和庞大的参数量带来了计算效率和通信开销的挑战，这些在推理过程中成为吞吐量的瓶颈。现有的并行策略，如专家并行（EP）、数据并行（DP）、张量并行（TP）或它们的直接组合，通常只能实现次优的推理吞吐量。

为了解决MoE架构的这些推理挑战，并超越现有的次优解决方案，本文提出了EPS-MoE，一种新颖的专家流水线调度器，用于高效服务MoE架构。

创新点与主要贡献
本文的核心贡献可总结如下：
1. 新颖的专家流水线并行调度器：引入了一种新颖的专家流水线并行调度器，用于高效的MoE模型推理，该调度器在核函数（kernel）级别实现了计算与通信的细粒度重叠。
2. GEMM效率分析与优化：对GEMM效率进行了深入分析，并为MoE的输入设计了水平切分，为权重设计了按专家切分的方法。基于这些切分方法，实现了一种在特定负载下从GroupGemm切换到DenseGemm的策略，以提高计算效率。
3. 并发模式探索：探索了Attention和MoE模块的不同并发模式，包括TP+TP、DP+EP和TP+EP，并分析了这些模式的效率。
4. 广泛的实验验证：将该方法应用于多种不同的MoE模型进行测试。实验表明，该方法最多可将预填充（prefill）吞吐量提高52.4%。例如，在高度优化的DeepSeekV2模型上，将官方宣称的每秒10万个令牌的吞吐量提升至至少每秒12万个令牌。

Table 1: EPS-MoE的预填充吞吐量增益。

1 DeepSeekV2和DBRX在8xH800-80GB SXM上测试，Mixtral(8x7B)在4xH800-80GB SXM上测试。

图1: DP、TP和EP在两个设备和两个专家上的权重分区。

A3 背景知识与关键观察

现代GPU架构

测试环境与基础计算单元。我们的测试使用了NVIDIA A100-80GB SXM和H800-80GB SXM GPU。使用的矩阵计算库主要是cutlass【3, Cutlass. https://github.com/NVIDIA/cutlass/tree/main. Accessed: 2024-11-05】和cublas【1, Cublas. https://docs.nvidia.com/cuda/cublas/. Accessed: 2024-11-05】。NVIDIA的矩阵计算主要采用分块（tiling）策略。输入矩阵和权重矩阵沿行或列被划分为不同的块（tile）。每个线程块（block）处理一个或多个块，并且每个线程块运行在单个流式多处理器（SM）上，如图3所示。SM构成了基本的计算单元。所有的硬件资源，如张量核心（Tensor Cores）或CUDA核心、内存带宽和通信带宽，都通过SM进行访问。

图3: Nvidia GPU架构的资源视图

计算

MoE推理的负载敏感性。MoE模型的推理性能受负载场景的显著影响。设E为MoE模型中的专家数量，k为单次前向传播中每个令牌选择的专家数量。当单次前向传播的令牌数为m时，MoE模型中激活的专家数量满足以下关系：

当m较大时（如预填充阶段），MoE中的所有专家都会被激活，此时MoE模型推理是计算密集型（compute-bound）的，可以充分利用其每令牌FLOPs相对较小的优势来加速预填充。当m较小时（如解码阶段），MoE推理则变为内存密集型（memory-bound），随着m在一定范围内增加，模型中激活的专家数量也增加，导致需要加载的权重参数量更高。因此，MoE模型的推理需要考虑其应用场景的负载条件。

GroupGemm的性能分析。MoE通常使用GroupGemm进行计算，如图4所示【20, Who says elephants can’t run: Bringing large scale moe models into cloud scale production, 2022, by Young Jin Kim et al.】。通常情况下，GroupGemm是一个相对高效的核函数。然而，我们发现随着计算负载的增加，GroupGemm核的性能并非保持不变。我们基于真实模型设置对GroupGemm和DenseGemm进行了一些测试，如图5所示，并得出了几个重要结论。

图4: GroupGemm演示。改编自【20, Who says elephants can’t run: Bringing large scale moe models into cloud scale production, 2022, by Young Jin Kim et al.】。所有矩阵乘法操作都通过单个核函数启动执行。

结论1：GroupGemm与DenseGemm的效率随输入规模变化而不同。如图5(a)所示，GroupGemm和DenseGemm的计算效率随着输入m的增加而逐渐提高。此外，当m ≥ 4096时，GroupGemm的计算效率将低于具有相同计算负载的DenseGemm。当m < 2048时，GroupGemm的效率高于DenseGemm。这表明最优的核函数实现随输入规模而变化，这被称为针对不同问题形状的核函数调优。本文还将GEMM类型纳入了考量。在MoE模型推理过程中，输入负载的变化很常见，例如预填充阶段输入通常落在DenseGemm更快的范围内，而解码阶段输入则通常落在GroupGemm更快的范围内。

结论2：对于GroupGemm，当输入大小达到一定值后，增加组数不会带来更高的吞吐量。增加GroupGemm的组数一直是提高计算吞吐量的方法，尤其是在细粒度的MoE模型中。当输入规模较小时，这种方法总是有效的。然而，当输入规模较大时，如图5(b)所示，增加组数并不会进一步提高计算吞吞吐量。因此，在输入规模较大时，将一个较大的GroupGemm拆分成一些较小的GroupGemm并不会导致吞吐量下降。

结论3：对于GroupGemm，当输入大小达到一定值后，使用更多的SMs不会带来更高的吞吐量。Nanoflow【39, Nanoflow: Towards optimal large language model serving throughput, 2024, by Kan Zhu et al.】已经指出，对于DenseGemm，可以在一定范围内减少其占用的SM数量而不影响GEMM的计算吞吐量。对于GroupGemm，我们从图5(c)中可以观察到类似的趋势。此外，测试结果还表明，对于不同组数的GroupGemm，在一定范围内减少GEMM占用的SM数量同样不影响其执行效率。

图5: GEMM性能分析数据。(a) 测试了不同GEMM在不同输入大小下的吞吐量，重点关注MoE块中的Gate和Up矩阵。对于GroupGemm，使用了16个专家，矩阵维度为[1536,5120]，确保与DenseGemm的总计算量相等。(b) 测试了GroupGemm [Pingpong]在不同组数和问题规模下的相对吞吐量，将当前组的吞吐量与最大值进行比较。(c) 在输入大小m=6144的情况下测试了GroupGemm [Pingpong]，评估了在不同SM计数和组数下的吞吐量。

内存访问

内存密集型算子的SM资源需求。在MoE推理过程中，存在许多内存密集型算子，如layernorm、residual、activation、topKGating等，并有许多传统优化方法，如核函数融合、权重 quantization 和向量化内存访问【2, Cuda pro tip: Increase performance with vectorized memory access. https://developer.nvidia.com/blog/cuda-pro-tip-increase-performance-with-vectorized-memory-access/. Accessed: 2024-11-05】。另一个内存密集型算子是Attention KVCache加载，许多论文已经研究了这个问题。原则上，内存密集型核函数也需要更多的SM来实现最佳性能。然而，对silu_activation算子的实际测试表明，虽然执行时间随SM数量的增加而减少，但一旦SM数量超过某个阈值，性能增益就变得微不足道。如图6所示，这个阈值与算子的计算负载密切相关。当令牌数为64时，40个SM足以实现良好的核函数性能。当令牌数为128、256或2048时，60个SM表现良好。当输入行数较小时，MoE GEMM将变为内存密集型，在这种情况下我们也能发现类似的趋势。

图6: 内存密集型silu_activation核函数在不同负载和SMs下的延迟：(a) 张量[64, 8192] - 性能下降在超过40个SM后可忽略不计。(b) 张量[2048, 8192] - 性能下降在超过60个SM后可忽略不计。

通信

通信开销与SM资源。通信是LLM推理中的一个普遍问题，因为模型规模通常过大，无法由单个GPU部署。然而，多GPU并行计算也引入了GPU之间额外的通信开销。从图7可以看出，随着GPU数量的增加，通信所占时间的比例也在增加。如果我们调整核函数使GroupGemm更高效，通信时间的比例将变得更加显著。为了减少通信时间，我们可以使用FP8量化来减少通信量，在精度和延迟之间取得平衡。但这并不能完全解决问题；随着批处理大小的增加，通信量将进一步增加。此外，随着模型规模的增长，将引入多机并行策略，届时机器间的带宽将远低于单机内GPU间的带宽。此时，通信所占时间的比例将变得更加显著。通信消耗SM资源，但我们观察到，在一定范围内，调整分配给通信核的SM数量几乎不影响通信吞吐量，即使在预填充等通信量大的场景中也是如此，如图7所示。

图7: (a) 不同模型下通信与计算的时间消耗。(b) all2all核在不同负载和SMs下的延迟。10-20个SMs已足够，超过此数量后延迟无显著改善。测试环境为NVIDIA A100-80GB SXM与NVLink。

A2 方法细节

EPS-MoE概述

工作原理与核心模块。图8展示了EPS-MoE在DP + EP模式下的工作方式，演示了通信与计算重叠的专家级流水线。EPS-MoE由三个核心模块组成：并行策略（Parallel Strategy）、专家流水线调度器（Expert Pipeline Scheduler, EPS）以及计算与通信重叠（Computation and Communication Overlapping）。

并行策略：并行策略决定了我们采用的通信模式，而通信模式又决定了我们如何应用重叠策略。并行策略的选择应基于内存占用、内存访问和计算量的变化。由于分布式专家计算的独特性，MoE模型为并行模式带来了新的变化。
专家流水线调度器：为了应用流水线策略，我们需要一种方法来划分MoE的输入张量和MoE的权重矩阵。传统的TP划分方法是一种选择。在EPS-MoE中，我们根据MoE模型的特性设计了另一种按专家划分权重的方法。我们将看到，这种划分方法对于应用更高效的GEMM计算和流水线策略是必要的。
计算与通信重叠：我们将输入张量按行拆分，每次只将当前专家所需的令牌传输到相应的设备。通过这种方法，我们以流水线方式并行化了计算和通信。本节还将讨论重叠策略的设计及其开销。

图8: 专家流水线调度器图示。 (a) 专家流水线调度器。专家数=6, topk=2, DP=2, EP=2。GPU0和GPU1上Attention块的输出分别是令牌[0,1,2,3,4]和令牌[5,6,7,8,9]。Expert0接收令牌[0,1,5,9]，Expert1接收令牌[0,2,5,6]，依此类推。
(b) 通信与计算的重叠。LocalReduce是内存密集型操作。当一个令牌的所有专家的计算完成后，可以启动该令牌的下一轮all2all通信。

并行策略

MoE推理时间建模与并行策略分析。MoE模型主要由Attention和MoE两部分组成。从这一点出发，MoE模型的推理时间可以描述如下：

Attention部分可以使用TP（张量并行）或DP（数据并行），而MoE部分可以采用TP或EP（专家并行）。因此，我们有多种MoE推理的并行策略。我们将详细讨论Attention TP + MoE TP、Attention DP + MoE EP和Attention TP + MoE EP。使用的符号如下：
* F(·)：某个块的计算FLOPs。
* W(·)：某个块的内存I/O数据量。
* $V_S(x)$：策略S下，单个设备上数据量x的通信量。$V_O$是通信开销。
* $M_{in}, M_{out}$：每个设备的输入/输出激活量；$M_{KVCache}$：总KVCache量（预填充时为0）；$M_{weight}$：参数量。
* k：MoE模型中一个令牌选择的专家数量。
* D：TP、EP或DP的设备数量。
* FP, BW, NW：计算FLOPS、内存带宽和设备间的NVLink带宽。
* P：批处理请求的激活大小。

对于不同的并行策略，单个设备的计算吞吐量是相同的。不同并行策略的核心影响在于：1. 考虑到现代GPU的并行机制，单个令牌的计算资源；2. 不同设备上的内存占用，因为一些并行策略会导致参数在不同设备间的冗余存储；3. 不同设备上的内存访问量。因此，Attention或MoE的计算吞吐量取决于最大内存访问量、计算FLOPs以及通信时间。

我们将重点讨论W(·)和V(·)。

TP + TP 策略。Attention部分和MoE部分都使用TP模式。

可以看出，如果MoE的top-k值较大，激活值I/O带来的影响会很显著。

由于所有令牌都存在于每个设备上，并且使用了ncclAllReduce通信，因此TP+TP模式下单个设备的通信量为：

DP + EP 策略。在Attention部分使用TP的问题是，当Attention采用MLA（Multi-head Latent Attention）结构时，需要引入额外的通信。因此，我们使用DP来加速MLA Attention部分，并对MoE使用EP。

对于DP+EP，通信时间主要取决于专家的部署方式。如果一个令牌选择的所有专家都在同一个设备上，通信量是最佳的。但如果它们分布在不同设备上，则不可避免地会放大通信量。因此，DeepSeekV2设计了一种设备限制的路由机制来约束MoE相关的通信成本【15, Deepseek-v2: A strong, economical, and efficient mixture-of-experts language model】。设g为一个令牌将被路由到的设备数量，g受限于min(k, D)，即g ≤ min(k, D)。因此，DP+EP模式下单设备的总通信数据量如下：

因此，我们有

TP + EP 策略。DP会导致不同设备间参数的冗余存储，因此对于Attention是MHA（Multi Head Attention）、GQA（Grouped Query Attention）或MQA（Multi Query Attention）的模型，最好对Attention部分使用TP。

至于通信，通常我们为TP通信使用ncclAllReduce。然而，为了应用专家流水线，我们将ncclAllReduce分解为两步以获得更好的通信效率。在分发（dispatch）阶段，我们使用ncclReduceScatter + all2all替代ncclAllReduce。在合并（combine）阶段，我们使用all2all + ncclAllGather替代ncclAllReduce。这将大大减少MoE模型的通信量。对于分发阶段，我们有：

我们很容易发现合并阶段的通信量与分发阶段相同。通过这种方式拆分ncclAllReduce算子，我们可以使用专家流水线策略将all2all算子和GEMM并行化。随后，我们还可以应用Flux【14, Flux: Fast software-based communication overlap on gpus through kernel fusion, 2024, by Li-Wen Chang et al.】的ncclReduceScatter和ncclAllGather，与前后GEMM操作集成。

策略总结。不同的并行模式影响内存访问的数据量和通信数据量，但不影响计算时间。对于Attention块的内存访问数据量，哪种最好取决于Attention权重和Attention激活值之间的关系。对于MoE块的内存访问数据量，我们有

对于通信数据量，我们有

DP不改变整体计算吞吐量，但会导致三个问题：1. 单个令牌的计算时间增加；2. 参数的冗余存储；3. 负载均衡问题。对于问题1和2，我们用它们来换取吞吐量的提升。对于负载均衡问题，我们将使用分块预填充（chunked prefill）来解决。至于EP的平衡策略，诸如平衡损失或令牌丢弃等方法【15, Deepseek-v2: A strong, economical, and efficient mixture-of-experts language model】【22, Gshard: Scaling giant models with conditional computation and automatic sharding, 2020, by Dmitry Lepikhin et al.】【35, M6-t: Exploring sparse expert models and beyond, 2021, by An Yang et al.】已经被研究过，这里不再赘述。

专家流水线调度器

数据划分方式对比。有两种数据划分方法可以实现MoE模型中计算和通信的流水线并行。如图9所示，水平切分（horizontal split）通过对输入张量进行行划分来组织输入数据，而垂直切分（vertical split）则通过列划分来组织输入数据。传统的水平切分方法会导致参数矩阵的重复内存I/O。在应用水平切分时，我们利用了MoE参数的稀疏性特点，并根据专家对参数进行切分，这可以有效避免参数矩阵的重复内存I/O。

图9: 两种不同的输入张量切分方式，水平切分和垂直切分。(a) 水平切分。切分输入张量的行。权重按专家切分。(b) 垂直切分。切分输入张量的列。权重按列切分。

切分方法的效率分析。基于此，让我们考虑这两种切分方法在内存I/O效率上的差异。设样本数为m，单个输入和输出激活的大小分别为$P_0$和$P_1$，每个专家参数的大小为W，专家总数为E，流水线数量为N。总I/O数据量V由下式给出：

方法(a)的I/O数据量远小于方法(b)。设单个专家的计算工作量为C，组大小为G的GroupGemm的计算吞吐量为R(FLOPS|G)。两种方法之间的计算时间差异为：

最终的差异体现在不同组大小的GroupGemm的计算效率上。当E = N时，GroupGemm可以被DenseGemm替代，以在输入m较大时获得更好的性能。我们可以看到，使用水平切分的核心优势是更高效的GEMM实现和更低的内存I/O数据量。因此，专家流水线调度器通过水平切分输入张来实现流水线调度。同时，权重根据专家进行分区，每次只传输特定专家组所需的令牌。从之前的分析数据中，我们可以看到GroupGemm和DenseGemm的相对优势在不同计算负载下是变化的。基于此，我们设计了一种负载感知的自适应调度策略，根据负载类型动态选择不同的高效实现。通过这种并行策略，我们可以相对容易地建立令牌级传输和计算的并行性，如图8a所示。

计算与通信重叠

流水线并行的基本前提。流水线并行是一种在任务需要不同类型硬件资源时提高资源利用率的常用策略【37, Distributed hierarchical gpu parameter server for massive scale deep learning ads systems, 2020, by Weijie Zhao et al.】。通常，当多个硬件资源协同工作时，可以使用流水线并行来提高整体吞吐量。在LLM推理的背景下，一些核函数是内存密集型的，一些是计算密集型的，还有一些是通信密集型的。然而，每种类型的核函数使用的硬件资源都需要SM资源的支持。由于一些核函数的执行可能同时需要内存访问和计算，而另一些可能同时需要内存访问和通信，因此没有单一类型的核函数能够将其使用的所有类型硬件资源利用到理论最大值。从前面的分析可以推断，通过控制一些核函数占用的SM数量，我们可以在保持该核函数执行时间不变或变化不大的前提下，为其他核函数留出一部分SM。这是实现流水线并行的基本前提。在EPS-MoE中，我们并行化了GroupGemm和DenseGemm与all2all通信，以提高硬件资源利用率，如图8b所示。在后续的测试中，我们将展示，如果控制计算和通信占用的SM数量，流水线可以实现更大的收益。

A4 实验环境与结果

实验环境

硬件配置:
- DeepSeekV2、DBRX、Snowflake Arctic测试平台：8x H800-80GB SXM GPU，通过NVLink连接。
- Mixtral 8x7B测试平台：4x H800-80GB SXM GPU，通过NVLink连接。
- 消融实验平台：4x H800-80GB SXM GPU。
- 性能分析也使用了NVIDIA A100-80GB SXM GPU。
模型:
- 细粒度专家模型：DeepSeekV2
- 专家较少的模型：Mixtral 8x7B, DBRX
- 具有并行结构的模型：Snowflake Arctic
软件配置:
- 基线系统：vLLM，EPS-MoE与vLLM集成进行对比测试。
- 并行策略：DeepSeekV2采用Attention DP + MoE EP；Mixtral和DBRX采用Attention TP + MoE EP。

实验结果

在DeepSeekV2上的实验
我们在8xH800-80GB SXM上测试了专家流水线对DeepSeekV2的效果。如表2所示，专家流水线对不同流水线数量（PN）的DeepSeekV2预填充吞吐量带来了显著提升，最高可达21%。考虑到DeepSeekV2本身已提供了非常高的基线（报告的预填充吞吐量为100,000 tokens/s），这一提升是非常可观的。
Table 2: DeepSeekV2的预填充吞吐量增益。

1 PN=1表示未使用流水线，但GroupedGemm会切换到DenseGemm。

在Mixtral 8x7B和DBRX上的实验
我们将Mixtral 8x7B部署在4xH800上，DBRX部署在8xH800上，与vLLM的TP+TP实现作为基线进行对比。如表3所示，EPS-MoE在Mixtral 8x7B上最多可将TTFT（Time To First Token）降低24.3%。与DeepSeekV2相比，EPS-MoE在Mixtral 8x7B上的增益随序列长度增加而增加，因为更大的批处理大小对专家流水线更有利。此外，通过对比TP+EP（PN=1）和基线，仅将并行模式从TP+TP切换到TP+EP并替换通信原语，就带来了8%~13%的TTFT降低。进一步分析发现，GEMM切换和并发模式的改变贡献了最多19.5%的降低，而流水线重叠贡献了4.8%的降低。DBRX模型拥有更多的专家和更大的topk，这使得EPS-MoE能获得更大的收益。
Table 3: Mixtral8x7B和DBRX的TTFT降低。

1 基础数据单位为秒。2 B为批处理大小，S为序列长度。

在Snowflake Arctic上的实验
我们还在具有并行结构的模型Snowflake Arctic上测试了专家流水线方法。如表4所示，在特定批处理大小下，专家流水线方法能取得一定收益。分析发现，这些收益主要来自于将矩阵运算方法从GroupGemm切换到DenseGemm，而流水线本身的收益相对较小。
Table 4: Snowfake Arctic的TTFT降低。

1 基础数据单位为毫秒(ms)。2 Snowflake Arctic使用3个层进行测试。

A7 补充细节：消融实验

我们采用DeepSeekV2的模型结构，使用一个包含Gate、Up、Down GEMM的层作为测试对象，在4xH800-80GB SXM上进行了消融实验，分析了以下问题：1. 从GroupedGEMM切换到DenseGemm的收益；2. FP8通信的影响；3. 控制SM的必要性；4. 不同输出维度的影响；5. 不同流水线数量（PN）的影响。测试数据统一呈现在表5中。

Table 5: 不同流水线数量下GroupGemm和DenseGemm的性能。

1 FP8表示通信数据类型是否为FP8。2 SM表示GEMM的SM计数。我们不控制DenseGemm的SM计数。3 所有数据单位均为毫秒(ms)。

1. 重叠与GEMM切换的贡献
两种优化策略均有收益，但适用场景不同。切换GEMM和重叠都带来了收益，尤其是在输入令牌数量较大时。从表5可以看出，当输入较大时（如m=3072，GEMM尺寸[5120, 1536]），将GEMM从GroupGemm切换到DenseGemm能带来12%的收益（ID=3的4.837ms到ID=7的4.299ms），而重叠的收益为16%（ID=0的5.698ms到ID=3的4.837ms）。当输入较小时，切换GEMM的收益较小甚至可能为负，主要收益来自流水线策略。值得注意的是，当输入较小且GEMM为内存密集型时，不进行重叠的EPS-MoE（仅GEMM切换）比重叠设置略有优势，因为通信开销更重。

2. FP8通信是否会降低专家流水线的收益？
FP8通信与流水线策略互补。在大多数情况下，FP8通信通过减少通信时间比例来影响重叠策略的收益。然而，FP8通信也面临量化带来的精度问题。通过同时应用FP8通信和重叠策略，我们发现这两种方法可以进一步结合产生更大的收益。从表5可以看出，FP8通信和专家流水线策略并不冲突，在大多数情况下，同时使用两者可以获得进一步的收益。例如，在m=3072、GEMM尺寸[5120, 1536]的条件下，比较数据[ID=2]和[ID=8]（通信均为FP8），后者仍然比前者有28.2%的提升。

3. 控制SM数量的必要性
控制SM是实现完美重叠的关键。如表5和表6所示，在应用专家流水线时，控制计算核使用的SM数量至关重要。控制SM主要是管理通信与计算重叠时计算所占用的SM数量，目标是在不影响计算效率的情况下，将更多SM分配给通信算子。分配的SM数量对计算和通信的时间消耗都至关重要。当GEMM算子占用所有SM时，通信算子无法被调度，导致其执行变慢；同时，通信算子的调度也会影响GEMM计算，增加其执行时间。在这种情况下，重叠的收益会减少。然而，通过控制GEMM算子占用的SM数量，通信和计算算子可以互不干扰地并行执行，实现完美的重叠。在表6中，我们将通信核的SM数设为16，因此这些测试中计算的最佳SM数均为116。
Table 6: 不同SM下的时间成本。

1 SM：GEMM的SM计数。m：每个专家的GEMM输入计数。2 所有数据单位均为毫秒(ms)。

4. 何时应该启用专家流水线？
专家流水线并非在所有场景下都有效。如果令牌数量较小，专家流水线将会失效。一般来说，只要划分后的通信和计算时间小于划分前，专家流水线就会有益。然而，由于在某些区域，计算和通信的划分并非线性可变，量化分析相对困难。我们可以计算一个相对宽松的下界。定义通信时间函数为S(m; k,n)，MoE GroupGemm时间函数为G(m,E; k,n)，其中m是输入令牌数，k,n是输入输出维度，E是专家数，N是流水线数。只要对于某个$\epsilon_0, \epsilon_1 \in \mathbb{R}$，满足以下条件，专家流水线就会起作用：

在DeepSeekV2的测试中，我们发现当m ≥ 1700时，对于$\epsilon_0 = 0.32, \epsilon_1 = 0.03$满足条件。从表7可以看出，1700是一个临界点，能带来一些边际收益。如图10所示，当m < $m_0$时，由于通信操作的$\epsilon_0$过大，EPS-MoE无法工作。当$m_0$ < m < $m_1$时，GEMM操作是内存密集型的，EPS-MoE从内存访问和通信重叠中获益。当m > $m_1$时，EPS-MoE从计算和通信重叠中获益。
Table 7: 基线与流水线的比较。

1 Base和Pipeline是在8xH800 SXM上测试单层前向传播时间。单位均为毫秒(ms)。Delta等于(Pipeline - Base)/Base。2 $\epsilon_0$表示当令牌数较小时，通信时间变化很大。

图10: EPS-MoE起作用的范围。

5. 如何找到最佳的PN（流水线数量）？
存在最优流水线数量以最大化收益。流水线数量对EPS-MoE的加速至关重要，如图11所示。不同的流水线数量可能导致完全不同的结果。我们发现：
* 给定特定的推理任务负载和模型计算成本，存在一个最优的流水线数量，可以最大化流水线并行的收益。
* 当令牌数量相对较小时，无论如何划分流水线都难以获得收益。

图11: 包含Gate/Up/Down GEMM的MoE核在不同流水线数量下的延迟。在4x H800 SXM NVLink GPU上测试。Standard表示使用DenseGemm的基线。这里忽略了从GroupGemm切换到DenseGemm的收益。Serial模式将通信和计算分块并串行处理，显示了分区开销。对于(a)和(b)，输入维度为5120，输出维度为1536（与DeepSeekV2相同）。对于(c)，输入维度为5120，输出维度为14336（类似于Llama模型，接近3倍）。

寻找最优PN的分析方法。我们可以通过分析建立一个找到最优流水线数量的方法。设θ为模型参数，N为流水线数，E为单个设备上的专家数。设函数L(θ;N)表示重叠计算和通信后的理论最大收益，R(N)表示流水线带来的开销。我们从测试结果中确定开销与流水线数量呈线性关系，即$R(N) \approx kN + b$。设$T_{comm}$为划分前的总通信时间，$T_{comp}$为MoE层的总计算时间。寻找最优流水线数量可以表述为求解以下方程：

其中，

设$C = \min\{T_{comm}, T_{comp}\}$，G为重叠带来的总延迟降低。

当且仅当$N = \sqrt{C/k}$时，可实现最大延迟收益。开销的存在意味着存在一个最优的流水线数量来实现重叠策略的最大效果；否则，流水线越多，重叠效果应越好。

A5 结论

本文为不同类型的MoE模型设计了不同的并行方案，并实现了专家流水线调度器（Expert Pipeline Scheduler），以实现MoE模型中计算与通信的重叠。除了重叠策略，在特定令牌数下从GroupedGemm切换到DenseGemm也是提高吞吐量的一个重要特性。我们在DeepSeekV2、Mixtral 8x7B和Snowflake Arctic等不同类型的MoE模型上测试了我们的方法。实验结果证明了所提方案的有效性，并为进一步优化MoE推理提供了一种新途径。

EPS-MoE: Expert Pipeline Scheduler for Cost-Efficient MoE Inference