MoE-Lightning: High-Throughput MoE Inference on Memory-constrained GPUs
MoE-Lightning：在内存受限 GPU 上实现高吞吐量 MoE 推理

Large Language Models (LLMs) have significantly improved in performance due to the advancements in architecture and scalable training methods. In particular, Mixture of Experts (MoE) models have shown remarkable improvements in model capacity, training time, and model quality (shazeer2017outrageously, ; glam, ; deepseek, ; dbrx, ; mixtral, ; lepikhin2020gshard, ), revitalizing an idea that dates back to the early 1990s (JacobsJNH91, ; jordan1994hierarchical, ) where ensembles of specialized models are used in conjunction with a gating mechanism to dynamically select the appropriate “expert” for a given task.
由于架构的进步和可扩展的训练方法， 大型语言模型（LLMs）在性能上显著提升。特别是，专家混合（MoE）模型在模型容量、训练时间和模型质量方面表现出显著的改进（shazeer2017outrageously,；glam,；deepseek,；dbrx,；mixtral,；lepikhin2020gshard,），复兴了一种可以追溯到 20 世纪 90 年代初的想法（JacobsJNH91,；jordan1994hierarchical,），即使用专门模型的集合并结合门控机制来动态选择适合特定任务的“专家”。

The key idea behind MoE is a gating function that routes inputs to specific experts within a larger neural network. Each expert is specialized in handling particular types of inputs. The gating function selects only a subset of experts to process an input, which allows LLMs to scale the number of parameters without increasing inference operations.
MoE 的关键理念是使用一个门控功能来将输入路由到大规模神经网络中的特定专家。每个专家专门处理特定类型的输入。门控功能仅选择一部分专家来处理输入，这使得 LLM 可以在不增加推理操作数的情况下扩展参数数量。

MoE models adopt a conventional LLM architecture, which uses learned embeddings for tokens and stacked transformer layers. MoE LLMs typically modify the Feed-Forward Network (FFN) within a transformer layer by adding a gating network that selects expert FFNs, usually implemented as multi-layer perceptrons, to process the input token (glam, ; zhou2022mixture, ; chen2023lifelong, ). These designs can surpass traditional dense models (chatbot-arena, ; deepseek, ; mixtral, ) in effectiveness while being more parameter-efficient and cost-effective during training and inference.
MoE 模型采用传统的 LLM 架构，使用学习得到的嵌入来表示标记及堆叠的 Transformer 层。MoE LLM 通常通过在 Transformer 层内的前馈网络（FFN）中添加门控网络来选择专家 FFN（通常实现为多层感知器）来处理输入标记（glam,；zhou2022mixture,；chen2023lifelong,）。这些设计能够在有效性上超越传统的密集模型（chatbot-arena,；deepseek,；mixtral,），同时在训练和推理过程中更具参数效率和成本效益。

Despite their advantages, the widespread use of MoE models faces challenges due to the difficulties in managing and deploying models with extremely high parameter counts that demand substantial memory. Thus, our work aims to make MoE models more accessible to those lacking extensive high-end GPU resources.
尽管有这些优势，MoE 模型的广泛使用面临挑战，因为管理和部署具有极高参数数量的模型需要大量内存。因此，我们的工作旨在使那些缺乏高端 GPU 资源的人也能更容易地使用 MoE 模型。

LLMs are trained to predict the conditional probability distribution for the next token, $P(x_{n+1}\mid x_{1},\ldots,x_{n})$, given a list of input tokens $(x_{1},\ldots,x_{n})$. When deployed as a service, the LLM takes in a list of tokens from a user request and generates an output sequence $(x_{n+1},\ldots,x_{n+T})$. The generation process involves sequentially evaluating the probability and sampling the token at each position for $T$ iterations. The stage where the model generates the first token $x_{n+1}$ given the initial list of tokens $(x_{1},\ldots,x_{n})$, is defined as the prefill stage. In the prefill stage, at each layer, the input hidden states to the attention block will be projected into the query, key, and value vectors. The key and value vectors will be stored in the KV cache. Following the prefill stage is the decode stage, where the model generates the remaining tokens $(x_{n+2},\ldots,x_{n+T})$ sequentially. When generating token $x_{n+2}$, all the KV cache of the previous tokens $(x_{1},\ldots,x_{n+1})$ will be needed, and the token $x_{n+2}$’s key and value at each layer will be appended to the KV cache.
LLM 被训练用于预测在给定一系列输入标记 $(x_{1},\ldots,x_{n})$ 的条件下下一个标记的条件概率分布 $P(x_{n+1}\mid x_{1},\ldots,x_{n})$ 。当作为服务部署时，LLM 接受用户请求中的一系列标记并生成输出序列 $(x_{n+1},\ldots,x_{n+T})$ 。生成过程包括依序评估概率并在每个位置采样标记，进行 $T$ 次迭代。模型在给定初始标记列表 $(x_{1},\ldots,x_{n})$ 时生成第一个标记 $x_{n+1}$ 的阶段定义为预填充阶段。在预填充阶段，每层中输入到注意力模块的隐藏状态将被投射为查询、键和值向量。键和值向量将存储在 KV 缓存中。预填充阶段之后是解码阶段，在此阶段，模型依次生成其余标记 $(x_{n+2},\ldots,x_{n+T})$ 。在生成标记 $x_{n+2}$ 时，所有先前标记的 KV 缓存 $(x_{1},\ldots,x_{n+1})$ 都是必要的，并且每层的标记 $x_{n+2}$ 的键和值将附加到 KV 缓存中。

The auto-regressive nature of LLM generation, where tokens are generated sequentially, can lead to sub-optimal device utilization and decreased serving throughput (pope2023efficiently, ). Batching is a critical strategy for improving GPU utilization: (yu2022orca, ) proposed continuous batching which increases the serving throughput by orders of magnitude. Numerous studies have developed methods to tackle associated challenges such as memory fragmentation (kwon2023efficient, ) and the heavy memory pressure imposed by the KV cache (he2024fastdecode, ; sheng2023flexgen, ; juravsky2024hydragen, ). The scenario of limited GPU memory introduces further challenges, especially for large MoE models, as it requires transferring large amounts of data between the GPU and CPU for various computational tasks with distinct characteristics. Naive scheduling of the computation task and data transfer can result in poor resource utilization. This paper explores how each resource in a heterogeneous system affects LLM inference performance and proposes efficient scheduling strategies and system optimizations to enhance resource utilization.
大型语言模型生成的自回归特性，按顺序生成标记，可能导致设备利用率不足和服务吞吐量降低（pope2023efficiently）。批处理是提高 GPU 利用率的关键策略：（yu2022orca）提出的连续批处理可将服务吞吐量提高若干数量级。许多研究已开发出方法，解决如内存碎片化（kwon2023efficient）和 KV 缓存带来的高内存压力（he2024fastdecode；sheng2023flexgen；juravsky2024hydragen）等相关挑战。有限的 GPU 内存情境带来了进一步的挑战，特别是对于大型 MoE 模型，因为这需要在 GPU 和 CPU 之间传输大量数据以进行不同特性计算任务。计算任务和数据传输的简单调度可能导致资源利用不足。本文探讨了异构系统中每种资源如何影响 LLM 推理性能，并提出高效的调度策略和系统优化，以提升资源利用。

We will start with the original Roofline Model (WilliamsWP09, ), which provides a visual performance model to estimate the performance of a given application by showing inherent hardware limitations and potential opportunities for optimizations. It correlates a system’s peak performance and memory bandwidth with the operational intensity of a given computation, where Operational Intensity ($I$) denotes the ratio of the number of operations in FLOPs performed to the number of bytes accessed from memory, expressed in FLOPs/Bytes.
我们将以原始 Roofline 模型（WilliamsWP09）开始，该模型提供一个视觉性能模型，通过展示固有的硬件限制和潜在的优化机会来估计给定应用程序的性能。它关联了系统的峰值性能和内存带宽与给定计算的操作强度，其中操作强度（ $I$ ）表示以 FLOPs 执行的操作数量与从内存访问的字节数之比，用 FLOPs/字节表示。

The fundamental representation in the Roofline Model is a performance graph, where the x-axis represents operational intensity $I$ in FLOPs/byte and the y-axis represents performance $P$ in FLOPs/sec. The model is graphically depicted by two main components:
Roofline 模型中基础的表示是性能图，其中 x 轴表示操作强度 $I$ （以 FLOPs/字节为单位），y 轴表示性能 $P$ （以 FLOPs/秒为单位）。该模型通过两个主要组成部分图形显示：

Memory Roof: It serves as the upper-performance limit indicated by memory bandwidth. It is determined by the product of the peak memory bandwidth ($B_{\text{peak}}$ in Bytes/sec) and the operational intensity ($I$). Intuitively, if the data needed for the computation is supplied slower than the computation itself, the processor will idly wait for data, making memory bandwidth the primary bottleneck. The memory-bound region (in blue) of the roofline is then represented by:
内存顶：它作为内存带宽指示的最高性能限制。由峰值内存带宽（ $B_{\text{peak}}$ 以字节/秒为单位）与操作强度（ $I$ ）的乘积决定。直观地说，如果计算所需的数据供应速度慢于计算本身，处理器将在数据上无所事事地等待，使内存带宽成为主要瓶颈。那么，Roofline 的内存受限区域（蓝色）被表示为：

where $P$ is the achievable performance.
其中 $P$ 是可实现的性能。

Compute Roof: This represents the maximum performance achievable limited by the machine’s peak computational capability ($P_{\text{peak}}$). It is a horizontal line on the graph (top edge of the yellow region), independent of the operational intensity, indicating that when data transfer is not the bottleneck, the maximum achievable performance is determined by the processor’s computation capability. The compute-bound part (yellow region) is then defined by:
计算顶：这表示受机器的峰值计算能力（ $P_{\text{peak}}$ ）限制的最大可达性能。它在图表上表示为一条水平线（黄色区域的顶边），独立于操作强度，表明当数据传输不是瓶颈时，最大可达性能由处理器的计算能力决定。计算受限部分（黄色区域）定义为：

The turning point is the intersection of the compute and memory roofs, given by the equation:
转折点是计算顶与内存顶的交点，由下式给出：

defines the critical operational intensity $\bar{I}$. Applications with $I\geq\bar{I}$ are typically compute-bound, while those with $I<\bar{I}$ are memory-bound.
定义了关键的操作强度 $\bar{I}$ 。具有 $I\geq\bar{I}$ 的应用程序通常是计算密集型的，而具有 $I<\bar{I}$ 的应用程序则是内存密集型的。

In practice, analyzing an application’s placement on the roofline model helps identify the critical bottleneck for performance improvements. Recent works (yuan2024llm, ) analyze different computations (e.g., softmax and linear projection) in LLM using the Roofline Model.
实际上，分析应用程序在屋顶线模型上的位置有助于识别性能改进的关键瓶颈。最近的研究（yuan2024llm）分析了在 LLM 中使用屋顶线模型的不同计算（例如 softmax 和线性投影）。

While the original Roofline Model demonstrates great power for application performance analysis, it is not enough for analyzing applications such as LLM inference that utilize diverse computing resources (e.g., CPU and GPU) and move data across multiple memory hierarchies (e.g., GPU HBM, CPU DRAM, and Disk storage).
尽管原始的屋顶线模型在应用性能分析中显示出了很大的威力，但它不足以分析利用多种计算资源（例如 CPU 和 GPU）并在多个内存层次（例如 GPU HBM、CPU DRAM 和磁盘存储）之间移动数据的应用程序，如 LLM 推理。

Consider a system with $n$ levels of memory hierarchies. Each level $i$ in this hierarchy is coupled with a computing processor. The peak bandwidth at which the processor at level $i$ can access the memory at the same level is denoted by $B_{\text{peak}}^{i}$. Additionally, the peak performance of the processor is denoted by $P_{\text{peak}}^{i}$¹¹1In this paper we assume when $i<j$, $P_{\text{peak}}^{i}\geq P_{\text{peak}}^{j}$ and $B_{\text{peak}}^{i}\geq B_{\text{peak}}^{j}$..
考虑一个具有 $n$ 个内存层次的系统。这个层次中的每一个层级 $i$ 都与一个计算处理器配对。层级 $i$ 上的处理器访问其所在层级内存的峰值带宽被标记为 $B_{\text{peak}}^{i}$ 。此外，处理器的峰值性能被标记为 $P_{\text{peak}}^{i}$ ¹ 。

For computation $x$ executed at level $i$ in the HRM, we can define its compute and memory roofs similarly as in the original Roofline Model:
对于在 HRM 中第 $i$ 层执行的计算 $x$ ，我们可以类似于原屋顶线模型来定义其计算和内存屋顶：

• •

  Compute Roof at level $i$:

  Report issue for preceding element

  (4)

  $$P_{x}^{i}\leq P_{\text{peak}}^{i}$$

  This represents the maximum computational capability at level $i$, independent of the operational intensity.
  这表示第 $i$ 层的最大计算能力，与操作强度无关。

  Report issue for preceding element
  
  • 第 $i$ 层的计算屋顶：
• •

  Memory Roof at level $i$:

  Report issue for preceding element

  (5)

  $$P_{x}^{i}\leq B_{\text{peak}}^{i}\times I_{x}^{i}$$

  
  • 第 $i$ 层的内存屋顶：

More importantly, in HRM, there is also the memory bandwidth from level $j$ to level $i$, denoted as $B_{\text{peak}}^{j,i}$, which will define another memory roof for computation $x$ that is executed on level $i$ and transfers data from level $j$:
更重要的是，在 HRM 中，还有从第 $j$ 层到第 $i$ 层的内存带宽，标记为 $B_{\text{peak}}^{j,i}$ ，这将定义另一个内存屋顶用于执行在第 $i$ 层并从第 $j$ 层传输数据的计算 $x$ ：

• •

  Memory Roof from level $j$ to $i$:

  Report issue for preceding element

  (6)

  $$P_{x}^{i}\leq B_{\text{peak}}^{j,i}\times I_{x}^{j}$$

  
  • 从第 $j$ 层到第 $i$ 层的内存屋顶：

Therefore, if computation $x$ is executed on level $i$, data from level $j$ needs to be fetched, and the peak performance will be bounded by the three roofs listed above (Eqs. 4, 5 and 6):
因此，如果计算 $x$ 在第 $i$ 层执行，需要从第 $j$ 层提取数据，则峰值性能将被上述三个屋顶（方程式 4、5 和 6）所限：

If operator $x$ is executed on level $i$ without fetching data from other levels, it reduces to the traditional roofline model and can achieve:
如果运算符 $x$ 在第 $i$ 层执行而不从其他层提取数据，它将简化为传统的屋顶线模型，并且可以实现：

To visualize the turning points and balance points discussed in the preceding sections, we conduct a case study with real HRM plots for computations²²2We only discuss the attention and feed-forward blocks since they account for the majority of computation time and represent quite different computation characteristics. in a single layer of the Mixtral 8x7B model on a Google Cloud Platform L4 instance. The hardware setting is as detailed in Fig. 3. Specifically, we let levels $i$ and $j$ represent GPU and CPU, respectively. Then, we define the following:
为了可视化前面部分讨论的转折点和平衡点，我们在 Google Cloud Platform L4 实例上的单层 Mixtral 8x7B 模型上进行了一个实际 HRM 图的案例研究来计算 ² 。硬件设置如图 3 所示。具体来说，我们将级别 $i$ 和 $j$ 分别代表 GPU 和 CPU。然后，我们定义如下：

Pipeline scheduling is a common approach to maximize compute and I/O resource utilization. Yet, the pipeline concerning GPU, CPU, and I/O is not trivial. In traditional pipeline parallelism for deep learning training (huang2019gpipe, ; narayanan2019pipedream, ; fan2021dapple, ), models are divided into stages which are assigned to different devices. Therefore, only output activations are transferred between stages, resulting in a single type of data transfer in each direction at a time. In our scenario, both weights and intermediate results need to be transferred between GPU and CPU. Intermediate results are required immediately after computation to avoid blocking subsequent operations, whereas weights for the next layer are needed only after all micro-batches for the current layer are processed. Additionally, weight transfers typically take significantly longer than intermediate results. Consequently, naive scheduling of I/O events can lead to low I/O utilization, which also hinders computation.
流水线调度是一种常见的方法，用于最大化计算和 I/O 资源的利用。然而，涉及 GPU、CPU 和 I/O 的流水线并不是简单的。在深度学习训练的传统流水线并行中（huang2019gpipe,; narayanan2019pipedream,; fan2021dapple,），模型分为多个阶段并分配到不同的设备上。因此，仅有输出激活在阶段之间传输，导致每次仅有一种方向的数据传输。在我们的场景中，权重和中间结果都需要在 GPU 和 CPU 之间传输。中间结果在计算后立即需要，以避免阻塞后续操作，而下一层的权重仅在处理完当前层的所有微批次后才需要。此外，权重传输通常比中间结果消耗更长的时间。因此，I/O 事件的简单调度可能导致 I/O 利用率低下，这也会妨碍计算。

FlexGen (sheng2023flexgen, ) primarily employs the fourth schedule ($\mathcal{S}_{4}$), where attention is performed on GPU and the KV cache for the next micro-batch is prefetched during the current computation. This approach results in higher KV cache transfer latency than performing attention directly on the CPU (Section 3.3) and consumes I/O bandwidth that could otherwise be used for weight transfers, reducing resource utilization compared to CGOPipe. FlexGen also supports CPU attention and adopts the third schedule ($\mathcal{S}_{3}$), which is the least optimized and may even perform worse than $\mathcal{S}_{4}$ if KV cache transfer latency is less than the sum of pre-attention, post-attention, and CPU attention latencies, as later shown by our evaluation results (Section 5). FastDecode (he2024fastdecode, ) suggests overlapping CPU attention with GPU computation, similar to the second schedule ($\mathcal{S}_{2}$). However, it does not target memory-constrained settings, so weight transfer scheduling is not considered.
FlexGen（sheng2023flexgen）主要采用第四种计划（ $\mathcal{S}_{4}$ ），在该计划中注意力在 GPU 上执行，并且在当前计算期间预取下一个微批次的 KV 缓存。这种方法导致 KV 缓存传输延迟高于直接在 CPU 上执行注意力（第 3.3 节），并消耗了本可以用于权重传输的 I/O 带宽，与 CGOPipe 相比降低了资源利用率。FlexGen 也支持 CPU 注意力，并采用第三种计划（ $\mathcal{S}_{3}$ ），这是一种最不优化的计划，如果 KV 缓存传输延迟小于前注意力、后注意力和 CPU 注意力延迟之和，性能可能甚至比 $\mathcal{S}_{4}$ 差，如我们的评估结果所示（第 5 节）。FastDecode（he2024fastdecode）建议将 CPU 注意力与 GPU 计算重叠，类似于第二种计划（ $\mathcal{S}_{2}$ ），然而，这并不针对内存受限的设置，因此权重传输调度未被考虑。

Due to independent data paths, data transfers in opposite directions can happen simultaneously. Data transfer will be performed sequentially in the same direction. The challenge then mainly lies in the scheduling of $\mathcal{D}_{2}$, $\mathcal{D}_{3}$ and $\mathcal{D}_{4}$, which are all from CPU to GPU. For the case without CPU attention ($\mathcal{S}_{4}$), while $\mathcal{D}_{4}$ usually takes a similar or longer time compared with a layer’s computation, the I/O bandwidth is almost fully utilized, leaving little room for more efficient scheduling for data transfer. As we can see from the diagram of $\mathcal{S}_{2}$ and $\mathcal{S}_{3}$, conducting the weights transfer as a whole will block the next layer’s first $\mathcal{D}_{2}$ for a long time, resulting in poor overall system efficiency. Instead, we can chunk the weights to be transferred into $n$ pages where $n$ equals the number of micro-batches in the pipeline, and the performance model and optimizer (Section 4.2) select the proper micro-batch size, batch size and the proportion of weights to be transferred from CPU to GPU.
由于独立的数据路径，相反方向的数据传输可以同时进行。同一方向的数据传输将顺序执行。挑战主要在于调度 $\mathcal{D}_{2}$ 、 $\mathcal{D}_{3}$ 和 $\mathcal{D}_{4}$ ，这些都从 CPU 传输到 GPU。在没有 CPU 注意力的情况下（ $\mathcal{S}_{4}$ ），尽管 $\mathcal{D}_{4}$ 通常需要与一个层计算相似或更长的时间，但 I/O 带宽几乎被完全利用，几乎没有剩余空间来为数据传输进行更高效的调度。如我们从 $\mathcal{S}_{2}$ 和 $\mathcal{S}_{3}$ 的图中看到，整体进行权重的传输将长时间阻碍下一层的第一次 $\mathcal{D}_{2}$ ，导致整体系统效率低下。相反，我们可以将要传输的权重分块成 $n$ 页，其中 $n$ 等于流水线中微批次的数量，性能模型和优化器（第 4.2 节）选择合适的微批次大小、批次大小和从 CPU 传输到 GPU 的权重量比例。

Algorithm 1 provides the order in which the main CPU task launcher thread launches the tasks to enable CGOPipe. All the tasks are executed asynchronously, and necessary synchronization primitives are added to each task to enforce the correct data dependency.
算法 1 提供了主 CPU 任务启动线程启动任务的顺序，以启用 CGOPipe。所有任务均异步执行，并在每个任务中添加了必要的同步原语以确保正确的数据依赖性。

Given a hardware configuration $\mathcal{H}$, a model configuration $\mathcal{M}$, and a workload configuration $\mathcal{W}$, we search for the optimal policy $\mathcal{P}$ that minimizes per-layer latency $T(\mathcal{M},\mathcal{H},\mathcal{W},\mathcal{P})$ for the pipeline schedule in Section 4.1, without violating the CPU and GPU memory constraints, in order to reach the optimal balance point (Eq. 11). Compared with FlexGen, we exclude disk-related variables from the search space and add two binaries to indicate whether to perform attention or MoE FFN on GPU.
给定硬件配置 $\mathcal{H}$ 、模型配置 $\mathcal{M}$ 和工作负载配置 $\mathcal{W}$ ，我们寻找最佳策略 $\mathcal{P}$ ，以最小化第 4.1 节中的流水线计划的每层延迟 $T(\mathcal{M},\mathcal{H},\mathcal{W},\mathcal{P})$ ，同时不违反 CPU 和 GPU 内存限制，以达到最佳平衡点（方程 11）。与 FlexGen 相比，我们从搜索空间中排除了磁盘相关变量，增加了两个二进制指标以指示是否在 GPU 上执行注意力或 MoE FFN。

The search space (Tab. 1) covers 2 integer values: the micro-batch size ($\mu$) and batch size ($N$), 2 binary indicators $A_{g}$ to indicate whether to perform the attention on GPU and $F_{g}$ to indicate whether to perform the MoE FFN on GPU. When $F_{g}=1$, we also need to decide the percent of weights $r_{w}$ that can be statically stored on GPU and the percent of weights $1-r_{w}$ that need to be transferred to GPU. Similarly, for $A_{g}=1$, we need to decide $r_{c}$. The generated policy will be a 6-tuple $(N,\mu,A_{g},F_{g},r_{w},r_{c})$. For our major setting, we always get $A_{g}=0$ and $F_{g}=1$. However, we discuss in Section 6.3 different policies for various hardware settings. Notably, CGOPipe is primarily designed for $A_{g}=0$ and when $A_{g}=1$, MoE-Lightning adopt $\mathcal{S}_{4}$.
搜索空间（表 1）涵盖了 2 个整数：微批大小（ $\mu$ ）和批大小（ $N$ ），两个二进制指标 $A_{g}$ 指示是否在 GPU 上执行注意力， $F_{g}$ 指示是否在 GPU 上执行 MoE FFN。当 $F_{g}=1$ 时，我们还需要决定可以静态存储在 GPU 上的权重百分比 $r_{w}$ 和需要传输到 GPU 的权重百分比 $1-r_{w}$ 。类似地，对于 $A_{g}=1$ ，我们需要决定 $r_{c}$ 。生成的策略将是一个 6 元组 $(N,\mu,A_{g},F_{g},r_{w},r_{c})$ 。对于我们的主要设置，我们总是得到 $A_{g}=0$ 和 $F_{g}=1$ 。然而，我们在第 6.3 节中讨论了针对各种硬件设置的不同策略。值得注意的是，CGOPipe 主要为 $A_{g}=0$ 而设计，当 $A_{g}=1$ 时，MoE-Lightning 采用 $\mathcal{S}_{4}$ 。

We then build the performance model based on Eq. 7 and Eq. 8 in HRM to estimate per-layer decode latency $T$ by:
然后，我们根据 HRM 中的方程 7 和方程 8 构建性能模型，估计每层解码延迟 $T$ ：

where $comm^{cpu\_to\_gpu}$ can be computed as the number of bytes needed to be transferred from CPU to GPU for a layer’s computation divided by the CPU to GPU memory bandwidth $b_{cg}$. Here, for simplicity, we only consider the attention computation and the MoE FFN computation in a transformer block, and therefore we have:
其中 $comm^{cpu\_to\_gpu}$ 可计算为从 CPU 转移到 GPU 进行一层计算所需字节数除以 CPU 到 GPU 的内存带宽 $b_{cg}$ 。在此，为简单起见，我们仅考虑 Transformer 块中的注意力计算和 MoE FFN 计算，因此我们有：

To estimate the time to perform a computation $x$ on GPU or CPU, we can use $T_{x}=\max(comm_{x},comp_{x})$ according to Eq. 8 in HRM, resulting in:
为估计在 GPU 或 CPU 上执行计算 $x$ 所需的时间，我们可以根据 HRM 中的方程 8 使用 $T_{x}=\max(comm_{x},comp_{x})$ ，所得结果为：

and similarly for $T_{attn}^{g}$, $T_{attn}^{c}$ and $T_{ffn}^{c}$.
同样地，对于 $T_{attn}^{g}$ ， $T_{attn}^{c}$ 和 $T_{ffn}^{c}$ 。

For a given computation $x$, we can calculate their theoretical FLOPS and data transfer based on $\mathcal{M}$ and then we have $comm_{x}^{g}=bytes_{x}/b_{g}$ and $comp_{x}^{g}=flops_{x}/p_{g}$ (same for CPU). While there are discrepancies between the theoretical performance estimation and the kernel’s real performance, such modeling can provide a reasonable estimation of the relative effectiveness of any two policies. In this paper, all the evaluation results of MoE-Lightning follow policies generated by a performance model with theoretically calculated computation flops and bytes with profiled peak performance and memory bandwidth for the hardware.
对于给定的计算 $x$ ，我们可以根据 $\mathcal{M}$ 计算其理论 FLOPS 和数据传输，然后得到 $comm_{x}^{g}=bytes_{x}/b_{g}$ 和 $comp_{x}^{g}=flops_{x}/p_{g}$ （同样适用于 CPU）。虽然理论性能估计与内核的实际性能之间存在差异，但这种建模可以为任何两个策略的相对效果提供合理的估计。在本文中，MoE-Lightning 的所有评估结果均遵循由性能模型生成的策略，这些模型使用理论计算的计算 FLOPS 和字节以及为硬件配置的峰值性能和内存带宽。

In existing works (sheng2023flexgen, ), pipeline parallelism is used for scaling beyond a single GPU, which requires the number of devices to scale with the model depth instead of the layer size. However, according to our analysis for MoE models in Section 3.3, Total GPU memory capacity can decide the upper bound throughput the system can achieve. Therefore, MoE-Lightning implements tensor parallelism (narayanan2021efficient, ) within a single node to get a higher throughput upper bound. In this case, we have $tp\_size$ times more GPU memory capacity and GPU memory bandwidth, we can then search for the policy similarly as for single GPU.
在现有作品中（sheng2023flexgen），管道并行用于超越单个 GPU 的扩展，这需要设备数量随着模型深度而不是层大小扩展。然而，根据我们在第 3.3 节中对 MoE 模型的分析，总 GPU 内存容量可以决定系统可以达到的上限吞吐量。因此，MoE-Lightning 在单节点内实现张量并行（narayanan2021efficient），以获得更高的吞吐量上限。在这种情况下，我们有 $tp\_size$ 倍更多的 GPU 内存容量和 GPU 内存带宽，然后我们可以像单个 GPU 一样搜索策略。

• •

  FlexGen (sheng2023flexgen, ) is the state-of-the-art offloading system that targets high-throughput batch inference for OPT (opt, ) models. It does not support variable prompt length in a batch and needs to pad all the requests to the maximum prompt length in the batch.

  Report issue for preceding element
  
  • FlexGen（sheng2023flexgen）是面向 OPT（opt）模型的最先进的卸载系统，专注于高吞吐量批量推理。它不支持批量中不同的提示长度，需要将所有请求填充到批量中的最大提示长度。
• •

  FlexGen(c) is FlexGen enabling CPU attention.

  Report issue for preceding element
  
  • FlexGen(c) 是启用了 CPU 注意力的 FlexGen。
• •

  DeepSpeed Zero-Inference (aminabadi2022deepspeed, ) is an offloading system that pins model weights to CPU memory and streams them layer-by-layer to GPU for computation. We use version 0.14.3 in the evaluation.

  Report issue for preceding element
  
  • DeepSpeed Zero-Inference（aminabadi2022deepspeed）是一个卸载系统，将模型权重固定到 CPU 内存，然后逐层流到 GPU 进行计算。在评估中，我们使用版本 0.14.3。
• •

  MoE-Lightning represents our system with all the optimizations enabled.

  Report issue for preceding element
  
  • MoE-Lightning 表示我们的系统在启用所有优化情况下的表现。
• •

  MoE-Lightning (p) represents our system running with requests padded to the maximum prompt length in the batch to compare with FlexGen.

  Report issue for preceding element
  
  • MoE-Lightning (p) 表示我们系统在将请求填充到批量中的最大提示长度情况下运行，以便与 FlexGen 进行比较。

Metrics. We measure the generation throughput for each workload, which is calculated as the number of tokens generated divided by total generation time (i.e., prefill time + decode time).
指标。我们测量每个工作负载的生成吞吐量，该吞吐量计算为生成的标记数除以总生成时间（即预填充时间 + 解码时间）。

We evaluate the maximum generation throughput for all baseline systems on three workloads under S1, S2, S6, and S7 settings. As shown in Footnote 9 and Tab. 4, MoE-Lightning (p) outperforms all baselines in all settings, and MoE-Lightning achieves up to $10.3\times$ better throughput compared with the best of the baselines for MTBench and HELM benchmark. In the following sections, we analyze how MoE-Lightning (p) outperforms our baselines by integrating the key methods from Section 4.2.
我们评估了 S1、S2、S6 和 S7 设置下所有基线系统在三个工作负载上的最大生成吞吐量。如脚注 9 和表 4 所示，MoE-Lightning (p) 在所有设置中均优于所有基线，而 MoE-Lightning 相比于 MTBench 和 HELM 基准中最好的基线，其吞吐量提高了 $10.3\times$ 。在接下来的章节中，我们分析了 MoE-Lightning (p) 如何通过整合第 4.2 节中的关键方法来超越我们的基线。

We observe this pattern for FlexGen and FlexGen(c) in all settings. However, MoE-Lightning (p) avoids a decrease in throughput under S1 and S6, which feature similar ratios of GPU to CPU memory. We attribute this performance improvement to CGOPipe, which significantly improves the resource utilization and renders the system GPU memory capacity bound in these settings.
我们在所有设置中观察到了 FlexGen 和 FlexGen(c)的这一模式。然而，MoE-Lightning (p)在 S1 和 S6 下避免了吞吐量的下降，这些设置具有相似的 GPU 到 CPU 内存比例。我们将这种性能改善归因于 CGOPipe，它显著提高了资源利用率，使系统在这些设置下受到 GPU 内存容量的限制。

On a single GPU (S1 and S2), MoE-Lightning (p) achieves up to 3.5$\times$, 5$\times$, and 6.7$\times$ improvement over FlexGen, FlexGen(c), and DeepSpeed, respectively.
在单个 GPU 上（S1 和 S2），MoE-Lightning (p)相对于 FlexGen、FlexGen(c)和 DeepSpeed 的性能提升分别高达 3.5 $\times$ 、5 $\times$ 和 6.7 $\times$ 。

The synthetic reasoning task enables all systems to have a larger micro-batch size due to the shorter prompt length. Under S1, MoE-Lightning (p) achieves a 1.16$\times$, 1.56$\times$, 2.22$\times$ higher throughput than FlexGen, FlexGen(c) and DeepSpeed respectively. Under S2, MoE-Lightning (p) finds a better balance point and uses less batch size than FlexGen, achieving 2.1$\times$ and 5.26$\times$ higher throughput compared to FlexGen and FlexGen(c), demonstrating the efficiency of CGOPipe and HRM.
合成推理任务由于较短的提示长度使所有系统都能使用更大的微批量大小。在 S1 下，MoE-Lightning (p)相比 FlexGen、FlexGen(c)和 DeepSpeed 的吞吐量分别提升 1.16 $\times$ 、1.56 $\times$ 、2.22 $\times$ 。在 S2 下，MoE-Lightning (p)找到一个更好的平衡点，并使用较小的批量大小，相比 FlexGen 分别实现了 2.1 $\times$ 和 5.26 $\times$ 的更高吞吐量，展示了 CGOPipe 和 HRM 的效率。

This section evaluates MoE-Lightning’s ability to run on multiple GPUs with tensor parallelism. As shown in S1 and S2, due to our efficient resource utilization, MoE-Lightning’s throughput is predominantly bounded by GPU memory capacity. This shows that increasing GPU memory can raise the system’s throughput upper bound.
本节评估了 MoE-Lightning 在多 GPU 上使用张量并行运行的能力。如在 S1 和 S2 所示，由于我们高效的资源利用，MoE-Lightning 的吞吐量主要受限于 GPU 内存容量。这表明增加 GPU 内存可以提高系统的吞吐量上限。

S6 and S7 in Footnote 9 show the end-to-end throughput results on Mixtral 8x22B of MoE-Lightning (p), FlexGen, and DeepSpeed on MTBench for multiple T4 GPUs. Notably, MoE-Lightning (p) achieves 2.77-3.38$\times$ higher throughput with 4xT4 GPUs than with 2xT4 GPUs, demonstrating super-linear scaling performance. DeepSpeed demonstrates a linear-scaling performance but uses a small batch size of 32, resulting in low throughput. FlexGen fails to scale under settings S6 and S7, largely due to the pipeline parallelism approach it employs. In this method, when using 4 GPUs, during the saturated phase, four layers are simultaneously active across four GPUs, increasing CPU peak memory consumption. As a result, FlexGen is bottlenecked by the CPU to GPU memory bandwidth and fails to take advantage of the added GPUs. Note that pipeline parallelism is more effective across multiple GPU nodes. In such configurations, doubling the number of GPUs also doubles the CPU to GPU bandwidth, the CPU memory capacity, and the CPU memory bandwidth¹⁰¹⁰10In this paper, we focus on the cases within one node..
页脚 9 中的 S6 和 S7 展示了 MoE-Lightning（p）、FlexGen 和 DeepSpeed 在 MTBench 上针对多个 T4 GPU 的 Mixtral 8x22B 的端到端吞吐量结果。值得注意的是，相较于 2xT4 GPU，MoE-Lightning（p）在 4xT4 GPU 下实现了 2.77-3.38 倍更高的吞吐量，表明具备超线性扩展性能。DeepSpeed 表现为线性扩展性能，但使用了较小的批量大小 32，导致吞吐量低。FlexGen 在设置 S6 和 S7 下无法扩展，主要因为其采用的流水线并行方法。在这种方法中，使用 4 个 GPU 时，在饱和阶段，四个层次在四个 GPU 上同时激活，增加了 CPU 峰值内存消耗。结果，FlexGen 受制于 CPU 至 GPU 的内存带宽，无法利用新增的 GPU。需要注意的是，流水线并行在多个 GPU 节点上更为有效。在此类配置中，GPU 数量翻倍同时意味着 CPU 至 GPU 带宽、CPU 内存容量和 CPU 内存带宽也会翻倍。

Fig. 8 demonstrates MoE-Lightning’s generation throughput results on DBRX to showcase the performance when all optimizations are enabled (CGOPipe, HRM and variable length prompts). For the DBRX model and without request padding (i.e., shorter prompt length), the system becomes less GPU memory capacity bound. We can see 2.1-2.8$\times$ improvement when scaling from 2 GPUs to 4 GPUs.
图 8 展示了 MoE-Lightning 在 DBRX 上的生成吞吐量结果，以展示开启所有优化（CGOPipe、HRM 和变量长度提示）后的性能。对于 DBRX 模型，且不进行请求填充（即，较短的提示长度），系统对 GPU 内存容量的限制较小。从 2 个 GPU 扩展到 4 个 GPU 时，我们可以看到 2.1-2.8 倍的改进。

In this section, we compare MoE-Lightning (p), FlexGen with its policy and FlexGen with our policy. For this experiment, we do not turn on the CPU attention for FlexGen as it is consistently worse than FlexGen w/o CPU attention. We use the workload from MTBench on the S1 setting with a generation length of 128. The results are displayed in Tab. 5. By deploying our policy, we can see a $1.77\times$ improvement in FlexGen. We also increase the batch size to better amortize the weights transfer overhead and it gives a $2.17\times$ speedup. However, it still cannot match MoE-Lightning’s throughput under the same policy, as KV cache swapping becomes the bottleneck for FlexGen in this case.
在本节中，我们比较了 MoE-Lightning（p）、使用其策略的 FlexGen 和使用我们策略的 FlexGen。对于此实验，我们没有打开 FlexGen 的 CPU 注意力功能，因为它一直表现不如不使用 CPU 注意力的 FlexGen。我们使用 MTBench 在 S1 设置下生成长度为 128 的工作负载。结果显示在表 5 中。通过部署我们的策略，我们可以看到 FlexGen 的性能有所提高。我们还增加了批量大小以更好地摊销权重转移开销，并达到了加速效果。然而，在相同策略下，它仍无法匹敌 MoE-Lightning 的吞吐量，因为 KV 缓存交换在这种情况下成为 FlexGen 的瓶颈。

In this section, we study when CPU attention will become the bottleneck in the decode stage. For different batch sizes (from 32 to 256), we test the latency of the MoE FFN kernel on L4 GPU and compare it with the latency of the CPU attention kernel on a 24-core Intel(R) Xeon(R) CPU @ 2.20GHz with various context lengths (from 128 to 2048). Additionally, we also measure the latency for swapping the KV cache needed for the attention from CPU pinned memory to GPU to validate the efficiency of our CPU GQA kernel.
在本节中，我们研究在解码阶段何时 CPU 注意力会成为瓶颈。对于不同的批量大小（从 32 到 256），我们在 L4 GPU 上测试 MoE FFN 内核的延迟，并将其与在不同上下文长度（从 128 到 2048）下在 24 核 Intel(R) Xeon(R) CPU @ 2.20GHz 上的 CPU 注意力内核延迟进行比较。此外，我们还测量了从 CPU 锁页内存到 GPU 为注意力所需的 KV 缓存交换的延迟，以验证我们 CPU GQA 内核的效率。

As shown in Fig. 9, our CPU attention kernel is $3-4\times$ faster than KV cache transfer, which is close to the ratio of CPU memory bandwidth and the CPU to GPU memory bandwidth. The MoE FFN’s latency doesn’t change so much across different micro batch sizes, which is as expected since the kernel is memory-bound for the decode stage. As the micro-batch size and context length increase, the CPU attention will eventually become the bottleneck, which calls for higher CPU memory bandwidth.
如图 9 所示，我们的 CPU 注意力内核比 KV 缓存传输要快，这接近于 CPU 内存带宽与 CPU 到 GPU 内存带宽的比率。MoE FFN 的延迟在不同的微批量大小下变化不大，这是意料之中的，因为该内核对解码阶段来说是受内存限制的。随着微批量大小和上下文长度的增加，CPU 注意力最终会成为瓶颈，这需要更高的 CPU 内存带宽。

In this section, we study how the best policy changes under different hardware settings. As we have shown in the previous ablation study, CPU attention can actually become the bottleneck for large batch size and context length, which means if we have more powerful GPUs, at some point, CPU attention may not be worth it. Moreover, if we have higher CPU to GPU memory bandwidth, the trade-offs will also change. Then the question becomes: when we have enough GPU memory (e.g., 2xA100-80G) to hold the model weights (e.g., Mixtral 8x7B), is it still beneficial to perform CPU computation or to offload weights/KV cache to the CPU? To conduct the analysis, we use 2xA100-80G for the GPU specification and vary the CPU to GPU memory bandwidth from 100 to 500 GB/s alongside different CPU capabilities. We set base CPU specifications at $m_{c}=200$GB/s, $b_{c}=100$GB, and $p_{c}=1.6$TFLOPS/s, scaling these values by multiplying with the CPU scaling ratio for various configurations¹¹¹¹11Note that this setup doesn’t reflect real-world hardware scaling; rather, it simplifies the number of variables to offer a rough idea of how these hardware configurations might impact performance..
在本节中，我们研究在不同硬件设置下最佳策略如何变化。正如我们在先前的消融研究中所展示的，CPU 注意力实际上可能成为大批量大小和上下文长度的瓶颈，这意味着如果我们有更强大的 GPU，在某些情况下，CPU 注意力可能不值得使用。此外，如果我们有更高的 CPU 到 GPU 内存带宽，权衡也会发生变化。那么问题就变成：当我们有足够的 GPU 内存（例如，2 个 A100-80G）来存储模型权重（例如，Mixtral 8x7B）时，执行 CPU 计算或将权重/KV 缓存卸载到 CPU 上是否仍然有益呢？为了进行分析，我们使用 2 个 A100-80G 作为 GPU 规格，并将 CPU 到 GPU 内存带宽从 100 变到 500 GB/s，同时结合不同的 CPU 性能。我们将基本 CPU 规格设置为 $m_{c}=200$ GB/s、 $b_{c}=100$ GB 和 $p_{c}=1.6$ TFLOPS/s，通过与 CPU 缩放比例相乘来扩展这些值，以适应各种配置 ¹¹ 。

We can see that when running Mixtral 8x7B on two A100 GPUs, as CPU-to-GPU memory bandwidth increases, more weight will be offloaded to the CPU. KV cache offloading is highly related to the CPU scaling ratio in this setup: when the CPU scaling ratio is low (i.e., low CPU memory bandwidth), even with the highest CPU to GPU memory bandwidth tested here, it is not beneficial to offload KV cache.
我们可以看到，当在两个 A100 GPU 上运行 Mixtral 8x7B 时，随着 CPU 到 GPU 内存带宽的增加，更多的权重会被卸载到 CPU。在此设置中，KV 缓存卸载与 CPU 缩放比例高度相关：当 CPU 缩放比例较低时（即低 CPU 内存带宽），即使是在此测试中的最高 CPU 到 GPU 内存带宽下，卸载 KV 缓存也是无益的。