MoE-Lightning: High-Throughput MoE Inference on Memory-constrained GPUs

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.

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 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.

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.

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.

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.

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.

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:

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:

where $P$ is the achievable performance.

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:

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.

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.

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).

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}$..

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:

• •

  Compute Roof at level $i$:

  (4)

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

  This represents the maximum computational capability at level $i$, independent of the operational intensity.

• •

  Memory Roof at level $i$:

  (5)

  $$P_{x}^{i}\leq B_{\text{peak}}^{i}\times I_{x}^{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$:

• •

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

  (6)

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

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):

If operator $x$ is executed on level $i$ without fetching data from other levels, it reduces to the traditional roofline model and can achieve:

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:

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.

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.

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.

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.

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.

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}$.

We then build the performance model based on Eq. 7 and Eq. 8 in HRM to estimate per-layer decode latency $T$ by:

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:

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:

and similarly for $T_{attn}^{g}$, $T_{attn}^{c}$ and $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.

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.

• •

  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.

• •

  FlexGen(c) is FlexGen enabling CPU attention.

• •

  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.

• •

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

• •

  MoE-Lightning (p) represents our system running with requests padded to the maximum prompt length in the batch to compare with 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.

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.

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.

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.

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.

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..

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.

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.

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.

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.

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..

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.