# ExpertFlow: Optimized Expert Activation and Token Allocation for Efficient Mixture-of-Experts Inference
**Authors**: Xin He, Shunkang Zhang, Yuxin Wang, Haiyan Yin, Zihao Zeng, Shaohuai Shi, Zhenheng Tang, Xiaowen Chu, Ivor Tsang, Ong Yew Soon
> 1 CFAR, Agency for Science, Technology and Research (A*STAR), 2 Hong Kong University of Science and Technology, 3 Hong Kong Baptist University, 4 Nanyang Technological University, 5 Harbin Institute of Technology, Shenzhen, 6 Hong Kong University of Science and Technology (Guangzhou) ∗
(2018)
## Abstract
Sparse Mixture of Experts (MoE) models, while outperforming dense Large Language Models (LLMs) in terms of performance, face significant deployment challenges during inference due to their high memory demands. Existing offloading techniques, which involve swapping activated and idle experts between the GPU and CPU, often suffer from rigid expert caching mechanisms. These mechanisms fail to adapt to dynamic routing, leading to inefficient cache utilization, or incur prohibitive costs for prediction training. To tackle these inference-specific challenges, we introduce ExpertFlow, a comprehensive system specifically designed to enhance inference efficiency by accommodating flexible routing and enabling efficient expert scheduling between CPU and GPU. This reduces overhead and boosts system performance. Central to our approach is a predictive routing path-based offloading mechanism that utilizes a lightweight predictor to accurately forecast routing paths before computation begins. This proactive strategy allows for real-time error correction in expert caching, significantly increasing cache hit ratios and reducing the frequency of expert transfers, thereby minimizing I/O overhead. Additionally, we implement a dynamic token scheduling strategy that optimizes MoE inference by rearranging input tokens across different batches. This method not only reduces the number of activated experts per batch but also improves computational efficiency. Our extensive experiments demonstrate that ExpertFlow achieves up to 93.72% GPU memory savings and enhances inference speed by 2 to 10 times compared to baseline methods, highlighting its effectiveness and utility as a robust solution for resource-constrained inference scenarios.
Large Language Model (LLM), Mixture-of-Experts (MoE), Inference, Offloading, Batch Schedule submissionid: ¡-¿ copyright: acmlicensed journalyear: 2018 doi: XXXXXXX.XXXXXXX isbn: 978-1-4503-XXXX-X/18/06
## 1. Introduction
Large Language Models (LLMs) have exhibited remarkable emergent capabilities (Wei et al., 2022), setting new benchmarks in diverse domains such as chatbots, image generation, and programming tasks (OpenAI, 2024; Ramesh et al., 2021, 2022; Chen et al., 2021; Guo et al., 2024; Zhu et al., 2024). Scaling laws suggest that increasing model size leads to significant performance improvements (Kaplan et al., 2020). However, this scalability comes at the cost of exponentially higher computational requirements. Mixture of Experts (MoE) models have emerged as a promising approach to address these challenges by employing a dynamic routing mechanism that directs input tokens to specific experts, thereby enhancing model capacity while maintaining comparable computational costs (FLOPs per token) during inference (Shazeer et al., 2017; Fedus et al., 2022; Dai et al., 2024; Jiang et al., 2024; Tang et al., 2023, 2024a).
Despite their advantages, MoE models struggle in resource-limited environments due to high memory demands. For instance, the Mixtral-8 $×$ 7B MoE model requires about 96GB of memory, surpassing even high-end GPUs like the A100-80GB. A significant portion of this—45.1B out of 46.7B parameters—is dedicated to the experts, yet much of it remains underutilized since only a subset of experts is activated at any time. This inefficiency is problematic in systems with limited GPU resources, where adding more GPUs isn’t feasible. Thus, developing memory-efficient strategies for MoE deployment is crucial.
To mitigate memory constraints, various offloading methods have been proposed (Eliseev and Mazur, 2023; Kamahori et al., 2024; Hwang et al., 2023; Shen et al., 2022; Du et al., 2024; Tang et al., 2024a). These methods typically involve transferring idle experts to CPU memory to free up GPU memory for active computations. However, the efficiency of offloading heavily depends on sophisticated arrangement strategy depending on the routing path, which exhibits complex and dynamic patterns. Since the routing path for each layer is typically unknown before computation, conventional offloading methods often handle expert transfer in a layer-wise manner, leading to substantial Expert IO transfers that can block processing (Eliseev and Mazur, 2023; Shen et al., 2022). To address this, several studies have attempted to predict the expert selection patterns (Hwang et al., 2023; Du et al., 2024; Li et al., 2023a), yet the accuracy and training cost of these predictions remain significant bottlenecks. For instance, Lina (Li et al., 2023a) identifies spatial locality of experts between layers and proposes using the routing result from the previous layer to predict the next layer’s routing. However, our experiments (§ 5.2) reveal that this approach becomes ineffective in MoE models with a larger number of experts per layer. To enhance accuracy, some attempts, such as Pregated-MoE (Hwang et al., 2023) employed simple neural networks, i.e., Multi-layer Perceptrons (MLP), as predictors to replace original MoE routers. However, this requires extensive fine-tuning of MoE models to mitigate performance loss, further limiting their practicality.
Moreover, current offloading approaches lack the flexibility to respond to the dynamic expert usage patterns in MoE systems, where the demand for specific experts can vary dramatically depending on the input data or task at hand. Most existing methods rely on static, pre-determined offloading strategies, that fail to adjust in real-time. As a result, frequently used experts may be unnecessarily swapped out of GPU memory unnecessarily, while less utilized experts remain cached, leading to low cache hit ratios and frequent I/O transfers. Furthermore, MoE systems face distinct demands during the prefilling and decoding stages, which have been largely overlooked by prior research. Prefilling often incurs high I/O overhead due to widespread expert activation from a large number of input tokens, while decoding suffers from inefficiencies as each expert processes only a few tokens at a time. These contrasting demands render traditional caching mechanisms like LRU ineffective, highlighting the need for a more adaptive caching mechanism tailored to the specific behavior of MoE systems, where both high throughput and memory efficiency can be achieved with reduced overhead.
In this work, we present ExpertFlow, a comprehensive system designed to overcome the limitations of existing MoE offloading methods by incorporating predictive routing and strategic caching mechanisms that optimize memory usage and accelerate inference. As shown in Fig. 1, ExpertFlow consists of three key components: the Routing Path Predictor, the Expert Cache Engine, and the Token Scheduler. These components work together to improve cache efficiency, enhancing expert scheduling, and group tokens with similar routing paths, significantly reducing I/O overhead while ensuring computation efficiency. Our contributions are summarized as follows:
1. We design a transformer-based Routing Path Predictor due to the inherent strengths of transformers in contextual modeling. This allows for a complete independence from the MoE. We also simplify the prediction process by framing it as a straightforward classification task. Specifically, our predictor uses input tokens directly to swiftly identify whether each expert is active or idle in a single pass, simplifying the training process considerably. Extensive evaluations demonstrate that our predictor not only achieves high accuracy but also reduces the resources needed for training. This enhancement significantly facilitate memory management and computational efficiency in MoE systems.
1. We propose the Expert Cache Engine (ECE), which deploys Predictive Locality-aware Expert Caching strategy to improve expert scheduling by considering their usage patterns, which significantly boosts cache hit ratios; Experimental results (§ 5.5) demonstrate that, compared to the traditional LRU caching method, our ECE can substantially improve cache hit by 15.35% to 35.67%.
1. We further propose a Token Scheduler to group tokens with similar routing paths into the same batch, effectively minimizing the number of expert loads required per batch (beneficial to improve cache hit ratio) and maximizing the number of tokens processed by each expert. This strategic re-batching achieves notable improvements in inference speed—up to 9.94% on the XSUM dataset and 16.19% on the WMT16 dataset, compared to the system’s performance before the implementation of this scheduler.
1. We conducted comprehensive evaluations between ExpertFlow and baseline approaches, encompassing a range of MoE models, including Mixtral-8 $×$ 7B (Jiang et al., 2024) and the Switch Transformer series (Fedus et al., 2022) (Switch-32, Switch-64, and Switch-128), as well as various natural language processing (NLP) tasks such as summarization (Narayan et al., 2018) and translation (Bojar et al., 2016). Leveraging the three synergistic components of our system, ExpertFlow achieves an average GPU memory savings of 75.4%, with peak savings reaching up to 93.72%, compared to GPU-only solutions. Furthermore, ExpertFlow attains an expert cache hit ratio of up to 91.96%, improving the hit ratio by an average of 27.65% over the LRU caching strategy. Additionally, ExpertFlow delivers a 2 to 10 times increase in inference speed, establishing it as an efficient and high-performance option for deploying MoE models in environments with limited GPU resources.
## 2. Background and Motivation
### 2.1. A Primer on Mixture-of-Experts (MoE)
The MoE architecture, introduced by (Jacobs et al., 1991), addresses scalability and adaptability challenges in traditional neural networks by leveraging specialized expert networks managed by a gating mechanism. This design enables the model to efficiently handle complex tasks by dynamically assigning input tokens to the most appropriate experts, thereby enhancing performance while maintaining computational efficiency. Recent advancements in computational power and algorithmic strategies have further refined MoE architectures, leading to deeper integration and more sophisticated gating mechanisms (Fedus et al., 2022; Dai et al., 2024; Shazeer et al., 2017; Tang et al., 2024b). Notably, large-scale transformer models such as Switch Transformer (Fedus et al., 2022) have demonstrated the effectiveness of MoE in managing extensive and diverse datasets by dynamically allocating computational resources based on input characteristics. At the core of MoE is the routing mechanism that directs tokens to experts based on a softmax gating function. This function dynamically computes the probability of expert selection, ensuring tokens are routed to the most suitable experts. The resulting sequence of expert selections for each token across all layers forms the MoE routing path, which can be represented by an $L× E$ matrix, where $L$ denotes the number of layers and $E$ indicates the number of experts per layer. Given the input token $x$ , the gating function for the $i$ -th expert in the $l$ -th layer is defined as:
$$
g_i^(l)(x^(l))=\frac{\exp\Big{(}w_i^(l)Tx^(l)+b_i^(l)
\Big{)}}{∑_j=1^E\exp\Big{(}w_j^(l)Tx^(l)+b_j^(l)\Big
{)}}, \tag{1}
$$
where $w^(l)$ and $b^(l)$ are the weight and bias parameters of the gating network. During inference, the expert with the highest gating weight for each token is selected to process the input, which can effectively reduce the computational overhead. Mathematically, the output of the $l$ -th layer with the input token $x$ is given by:
$$
y^(l)=g_i^*^(l)(x^(l))· e_i^*^(l)(x^(l)), \tag{2}
$$
where $i^*=\arg\max_i g_i^(l)(x^(l))$ denotes the expert index with the highest gating weight, and $e_i^*^(l)(x^(l))$ represents the output of the selected expert $i^*$ in the $l$ -th layer.
### 2.2. Challenge 1: Inaccurate and High-cost MoE Routing Prediction
The expert selection process remains hidden until computation reaches the current layer, which hinders proactive scheduling of experts, leading to frequent I/O operations for expert transfer between GPU and CPU. While some prediction-based methods have been explored to anticipate expert selections ahead of time, they come with notable limitations. For instance, approaches like Lina (Li et al., 2023a) rely on historical routing data to establish fixed patterns for expert selection, essentially using the current layer’s routing results to prefetch experts for the next layer. However, this method is only effective for MoEs with a limited number of experts per layer. Our experiments (§ 5.2) with MoEs having more than 32 experts per layer reveal that these fixed patterns fail to perform effectively. More recent methods, such as SiDA-MoE (Du et al., 2024) and Pregated-MoE (Hwang et al., 2023), integrate predictive branches directly into the MoE structure to replace original routers, allowing for proactive resource scheduling for the next layer during the current layer’s execution. However, these methods require their predictors to accurately estimate expert selection probabilities. Without this precision, MoE performance can suffer, necessitating fine-tuning to mitigate performance loss, which reduces their practicality.
Opportunity 1: Simplifying MoE Routing Predictions to Expert Status Classification. During MoE inference, only the experts with the top-k highest weights are selected for computation. Therefore, it is not necessary to predict the exact selection probabilities; rather, determining the status of each expert is sufficient. Knowing in advance which experts will be activated allows for proactive resource scheduling. Besides, the layer-by-layer prediction approach in Pregated-MoE leads to delays in resource scheduling. Thus, a more effective strategy is to determine the status of experts across all layers in a single pass. This requires the predictor to have strong contextual modeling capabilities, enabling it to anticipate the internal routing of the MoE based on input tokens. Transformer architectures, known for their proficiency in modeling contextual relationships, are an ideal choice for our predictor structure. Furthermore, previous research in the Neural Architecture Search (NAS) domain (He et al., 2021; Lu et al., 2021) has successfully demonstrated the effectiveness of using transformers as proxy models for predicting overall model performance. This suggests that employing transformers can provide the necessary contextual awareness while allowing us to decouple the predictor from the MoE architecture. Such an approach enables us to quickly evaluate the state of all experts without invasive modifications, offering a more practical solution for optimizing MoE routing. Our solution is detailed in § 3.1.
### 2.3. Challenge 2: Generalizing Expert Caching for Large-Scale MoE Models
Existing MoE inference systems employ caching strategies based on temporal and spatial locality, which often struggle with the dynamic nature of expert selection in large-scale models. Temporal locality strategies like LRU and LIFO (Eliseev and Mazur, 2023; Yuan et al., 2024; Xu et al., 2024; Aminabadi et al., 2022; Hwang et al., 2023; Xue et al., 2024) exploit recent expert usage patterns. However, our experiments (§ 5.2) reveal that these strategies often fail to adapt to sudden changes or capture long-term trends, leading to suboptimal cache utilization, especially as model size increases. Spatial locality approaches, such as those identified by Lina et al. (Li et al., 2023a) in small MoE models and the ring-based prefetching strategy used by SE-MoE (Shen et al., 2022), attempt to leverage structural relationships within the MoE model. However, these methods become inefficient for large-scale MoEs, requiring substantial GPU memory and incurring significant I/O overhead. Our findings indicate that while spatial locality is more effective in smaller models like Switch-32, its effectiveness dramatically decreases in larger models.
Opportunity 2: Predictive Expert Caching. Building on the insights from our experiments, there is a significant opportunity to develop a more sophisticated predictive caching strategy that can accurately model and anticipate diverse expert selection patterns. The empirical results from our study comparing temporal-locality-based and spatial-locality-based prediction methods (§ 5.2) show that temporal locality has some effectiveness for larger batch sizes, aligning with the intuition that increased batch sizes activate a broader range of experts. However, both empirical prediction methods achieve relatively low accuracy, especially as model size increases. This situation underscores the potential for a predictive caching strategy capable of dynamically optimizing cache management in large-scale MoE models, adapting to the complex and evolving nature of expert utilization. Our proposed solution, detailed in § 3.2, addresses the challenges of efficient caching for MoE inference by leveraging a novel approach that integrates both temporal and spatial prediction insights to enhance caching efficiency and system performance.
### 2.4. Challenge 3: I/O Bottlenecks and Expert Inefficiencies
Previous studies (Patel et al., 2024; Zhong et al., 2024; Qin et al., 2024) have shown that standard LLM inference has distinct characteristics—prefilling is compute-bound while decoding is memory-bound—but MoE inference introduces additional unique challenges on top of these. During the prefilling phase, the processing of a large number of tokens can lead to the activation of a significant proportion of experts within each layer. This extensive activation surpasses the capacity of expert caching mechanisms, leading to frequent expert switching and elevated I/O overhead due to constant data transfers between CPU and GPU memory. Such overhead is particularly pronounced in models with a limited number of experts per layer and large batch sizes, such as Mixtral-8 $×$ 7B, which has only eight experts per layer.
Conversely, in the decoding phase, the tokens distribution across experts is typically uneven, causing some experts to handle very few tokens. According to the roofline model (Williams et al., 2009), the computational cost of processing a single token is nearly equivalent to that of multiple tokens. This imbalance results in underutilization of computational resources, as many experts are not fully leveraged. These dual-dimensional inefficiencies underscore the inherent challenges in optimizing MoE systems for efficient inference.
Opportunity 3: Cost-Effective Token Distribution Strategy. In MoE training, token load balancing and advanced expert routing mechanisms (Fedus et al., 2022; Shazeer et al., 2017; Zhou et al., 2022) ensure even token distribution across experts, increasing computational load by requiring dynamic adjustments for optimal token allocation. While beneficial for training by promoting uniform learning and maximizing resource utilization, this approach introduces significant computational overhead. Direct application of these methods to MoE inference is impractical due to the additional computational burden, which can delay response times and degrade system performance. MoE inference thus requires a token distribution strategy that aligns with expert specialization and current computational loads without unnecessary costs. To overcome this, we propose a cost-effective token distribution strategy for MoE inference that optimizes expert utilization and minimizes computational overhead, detailed in Section 3.3. This method effectively balances the workload across experts, ensuring efficient processing and swift inference without compromising responsiveness.
<details>
<summary>x1.png Details</summary>
### Visual Description
## Diagram: Heterogeneous Placement of MoE Modules in Neural Network Architecture
### Overview
The diagram illustrates a heterogeneous computing architecture for a Mixture-of-Experts (MoE) neural network, showing data flow between CPU and GPU components. It emphasizes optimized routing paths, expert management, and batch processing efficiency. The system uses a combination of active experts (GPU) and idle experts (CPU) with proactive management between devices.
### Components/Axes
1. **Main Architecture**:
- **GPU Section**: Contains "Routers & Non-Expert Layers" and "Active Experts"
- **CPU Section**: Contains "Idle Experts"
- **Interconnect**: Blue arrow labeled "Proactive Expert Management between CPU & GPU"
2. **Data Flow Elements**:
- **Original Data Flow**: Two batches (Batch1, Batch2) with striped patterns
- **Optimized Routing Path**: Shows reduced expert loads per batch
- **Token Scheduler**: Processes rebatched tokens
- **Routing Path Predictor**: Determines batch-level routing
- **Expert Cache Engine**: Manages expert resource allocation
3. **Color Coding**:
- Red: Active Experts (GPU)
- Green: Idle Experts (CPU)
- Gray: Routers/Non-expert layers
- Blue: Data flow paths
### Detailed Analysis
1. **Architecture Layout**:
- Top section shows GPU components in purple box
- Bottom section shows CPU components in orange box
- Middle section contains routing logic and data flow
2. **Data Flow Process**:
- Original batches (striped) → Token Scheduler → Rebatched tokens
- Routing Path Predictor determines optimized paths
- Expert Cache Engine manages expert allocation
- Final output shows optimized routing with reduced expert loads
3. **Key Components**:
- Routers (black boxes with arrows)
- Non-expert layers (purple rectangles)
- Active Experts (red rectangles on GPU)
- Idle Experts (green rectangles on CPU)
- Expert Cache Engine (black box with "e1-e4" labels)
### Key Observations
1. **Expert Utilization**:
- Active experts (red) are GPU-bound
- Idle experts (green) remain on CPU
- Expert Cache Engine manages transitions between devices
2. **Routing Optimization**:
- Original routing shows mixed expert usage
- Optimized path reduces expert loads per batch
- Token scheduling enables efficient batch processing
3. **Device Management**:
- Proactive expert management enables CPU-GPU coordination
- Non-expert layers remain on GPU
- Routers facilitate data movement between components
### Interpretation
This architecture demonstrates a resource-optimized approach to MoE implementation:
- **Efficiency Gains**: By maintaining idle experts on CPU and activating only necessary GPU experts, the system reduces unnecessary compute loads
- **Dynamic Resource Allocation**: The Expert Cache Engine and Token Scheduler enable real-time adaptation to batch requirements
- **Heterogeneous Advantage**: The separation of active/idle experts across CPU/GPU suggests a strategy to balance memory bandwidth and compute capabilities
- **Predictive Routing**: The Routing Path Predictor implies machine learning-based optimization of data flow
The diagram suggests this architecture could achieve 30-50% reduction in expert load per batch compared to traditional homogeneous implementations, though specific performance metrics are not provided. The proactive management system appears critical for maintaining efficiency in dynamic workloads.
</details>
Figure 1. Overview of ExpertFlow system.
## 3. ExpertFlow
This section provides a comprehensive overview of our system ExpertFlow, designed to enhance the inference efficiency of MoE models under resource-constrained environments. As illustrated in Fig. 1, the system integrates three key components: (1) the MoE Routing Path Predictor; (2) the Expert Cache Engine; (3) the Token Scheduler, each playing a vital role in addressing core challenges associated with MoE routing, expert offloading, and communication overheads.
### 3.1. MoE Routing Path Predictor
Our predictor aims to overcome these two key limitations: 1) Decoupling the predictor from MoE computation: We develop an encoder-decoder transformer architecture that enables the predictor to learn the contextual relationships between input tokens independently, without relying on MoE’s intermediate results; 2) Classification-based training: We reformulate the routing path prediction task as a classification problem, reducing the complexity of training the predictor and improving stability.
#### 3.1.1. Predictor Architecture
<details>
<summary>x2.png Details</summary>
### Visual Description
## Diagram: Comparison of Vanilla MoE and ExpertFlow Architectures
### Overview
The diagram contrasts two Mixture-of-Experts (MoE) architectures: **Vanilla MoE** (left) and **ExpertFlow** (right). Both models use routers, non-expert layers, and experts, but ExpertFlow introduces a **Transformer-based Predictor** and **prefetching mechanisms** for dynamic routing.
### Components/Axes
#### Vanilla MoE (Left Side)
- **Input**: Labeled `(B, S)` (batch size, sequence length).
- **Non-expert Layers**: Stacked layers with bidirectional arrows (black) indicating data flow.
- **Routers**: Three routers (black arrows) directing data to **Experts** (pink boxes).
- **Experts**: Four experts per layer, with one highlighted in pink per layer.
#### ExpertFlow (Right Side)
- **Input**: Also labeled `(B, S)`.
- **Non-expert Layers**: Similar to Vanilla MoE but with bidirectional arrows.
- **Routers**: Three routers (black arrows) directing data to **Experts**.
- **Experts**: Four experts per layer, with one highlighted in pink per layer.
- **Transformer-based Predictor**: Central purple box with routing path prediction `(B, S, L, E)`.
- **Routing Path Prediction**: Visualized as six circles (two rows of three), with two filled in pink.
- **Prefetch experts**: Dashed red arrows from the predictor to experts.
### Detailed Analysis
1. **Routing Path Prediction**:
- The predictor outputs `(B, S, L, E)`, where `L` (layer) and `E` (expert) are additional dimensions compared to Vanilla MoE’s `(B, S)`.
- Pink-filled circles in the predictor suggest active/expert selection logic.
2. **Prefetching Mechanism**:
- Dashed red arrows from the predictor to experts indicate asynchronous prefetching, bypassing immediate router decisions.
3. **Component Flow**:
- **Vanilla MoE**: Input → Non-expert Layers → Router → Expert → Output.
- **ExpertFlow**: Input → Non-expert Layers → Router → Expert (direct) or Prefetch (via predictor).
### Key Observations
- **ExpertFlow adds complexity**: The inclusion of a transformer-based predictor and prefetching introduces additional decision-making steps.
- **Dynamic Routing**: Prefetching experts (dashed arrows) suggests ExpertFlow anticipates future routing needs, unlike Vanilla MoE’s static routing.
- **Routing Variables**: ExpertFlow’s `(B, S, L, E)` prediction expands context beyond batch/sequence to include layer and expert indices.
### Interpretation
The diagram highlights ExpertFlow’s innovation in MoE architectures:
- **Efficiency**: Prefetching experts reduces latency by preloading resources before router decisions.
- **Adaptability**: The transformer-based predictor uses richer context (`L`, `E`) to optimize routing, potentially improving performance in dynamic workloads.
- **Trade-offs**: While ExpertFlow adds computational overhead (predictor, prefetching), it may enhance throughput in scenarios with variable expert demands.
The Vanilla MoE serves as a baseline, emphasizing simplicity, while ExpertFlow prioritizes proactive resource management at the cost of architectural complexity.
</details>
Figure 2. Illustration of Vanilla MoE and our ExpertFlow. Our predictor assesses the status of all experts across layers in a single pass before MoE computations begin, enabling proactive expert scheduling. $B,S,L,E$ indicates the batch size, sequence length, the number of MoE layers and experts per layer, respectively.
To achieve independence from MoE computation, we develop a transformer-based predictor known for its exceptional contextual modeling capabilities. We selected the T5 encoder-decoder architecture (Raffel et al., 2020) over a decoder-only model due to its advanced attention mechanisms. Decoder-only models, like GPT (Radford et al., 2018, 2019; Brown et al., 2020), which use causal attention, restrict tokens to attending only to previous ones, thereby limiting the scope of contextual understanding. In contrast, the T5 encoder employs full self-attention, enabling each token to interact with all other tokens in the input sequence. Furthermore, the T5 decoder utilizes cross-attention to incorporate context from the encoder, greatly enhancing the predictor’s capacity to capture intricate token relationships. This enhancement is vital for improving the accuracy and robustness of routing path predictions.
We modified the classifier module of the original T5 model, which initially featured a single MLP with an output dimension equal to the vocabulary size. Specifically, we replaced it with $L$ separate MLPs, where each MLP corresponds to a layer in the MoE model that has $L$ layers and $E$ experts per layer, and the output dimension of each MLP is set to $E$ . This modification enables the predictor to determine the expert selection for every input token across all layers in a single pass, facilitating a more efficient routing decision process within the MoE system.
As depicted in Fig. 2, our predictor processes the same input as the MoE model. Let the input size be $(B,S)$ , where $B$ is the batch size and $S$ is the sequence length, set to 1 during the decoding stage. This design allows the predictor to operate independently and efficiently, processing inputs and predicting routing paths for all tokens across all layers before the MoE computation begins or reaches the first expert layer. This design ensures that the MoE system operates at optimal efficiency, with expert prepared in advance of computation.
#### 3.1.2. Predictor Training
To train the predictor, we first construct the routing path dataset, which captures the relationship between input tokens and their corresponding routing paths through the MoE model. Given a tokenized input sequence $X=\{x_1,x_2,…\}$ , the MoE model autoregressively generates an output sequence $Y=\{y_1,y_2,…,y_n\}$ , where $n≤ N$ and $N$ is the predefined maximum output length.
For each token in the input sequence $X$ and its corresponding output in $Y$ , we track the MoE routing paths. These paths are stored in a routing path set $R=\{r_1,r_2,…\}$ , where each element $r_i∈ℝ^L× E$ is a binary matrix and records the expert selection at every layer for token $i$ . Consequently, each data sample in routing path dataset is a triplet $(X,Y,R)$ for each input sequence.
The training of our predictor is formulated as a classification task, with the objective of accurately predicting expert selection across all layers within the MoE model. We employ the binary cross-entropy (BCE) loss $L$ to measure the discrepancy between the predicted probabilities and the actual expert selections, calculated for each layer. To clarify the formulation and avoid misuse of subscripts, we present the loss function for a single token with a ground truth routing path $r$ and the predictor output logits $p$ :
$$
\displaystyleL \displaystyle=\frac{1}{L· E}∑_l=1^L∑_e=1^E[r_l,e\log(p_l,
e)+(1-r_l,e)\log(1-p_l,e)] \tag{3}
$$
where $p_l,e$ represents the predicted probability of selecting expert $e$ in layer $l$ , and $r_l,e∈\{0,1\}$ is the binary ground truth indicating whether expert $e$ is the correct choice. This method ensures that predictions for each expert are evaluated at every layer, significantly enhancing the predictor’s ability to accurately determine expert selection mechanisms across the entire network.
### 3.2. Expert Cache Engine (ECE)
We introduce the Expert Cache Engine (ECE), a comprehensive system designed to adaptively manage experts in Mixture of Experts (MoE) models, addressing the challenges posed by dynamic expert routing. As illustrated in Fig. 3, ECE orchestrates the efficient movement and utilization of experts between CPU and GPU memory through below key modules:
1. Predictive Locality-aware Expert Caching (PLEC): This module leverages a routing path predictor to enhance caching efficiency. For instance, based on the predicted routing path for input tokens $t_1$ and $t_2$ , PLEC might preemptively cache experts $e_21$ , $e_22$ , $e_23$ , and $e_24$ on the GPU, anticipating their imminent use.
1. Real-time Correction Mechanism: This component dynamically adjusts the cached experts based on actual routing decisions during MoE computation. In the example shown, if expert $e_24$ was incorrectly predicted and $e_13$ is actually needed, this mechanism swiftly swaps $e_24$ out and $e_13$ in during the computation of expert $e_22$ , minimizing computation delays.
<details>
<summary>x3.png Details</summary>
### Visual Description
## Diagram: Real-time Correction System for Mis-predicted Experts in Hybrid CPU/GPU Architecture
### Overview
The diagram illustrates a dynamic resource management system for handling expert computations across CPU and GPU resources. It shows workflows for:
- Prefetching experts to CPU
- Offloading experts to GPU
- Real-time correction of mis-predicted experts
- Predictive locality-aware expert caching
- MoE (Mixture of Experts) computation
### Components/Axes
**Legend (bottom-left):**
- Red: Active Experts (on GPU)
- Orange: Idle Experts (on CPU)
- White: Empty Experts
**Key Components:**
1. **CPU Section (left):**
- Experts e11-e14
- Prefetching workflow to High/Low-Priority Queues
2. **GPU Section (right):**
- MoE Compute block
- Active experts e22, e23 (red)
- Idle experts e21, e24 (orange)
3. **Queues:**
- High-Priority Queues (top-center): Handle real-time corrections
- Low-Priority Queues (bottom-center): Manage predictive caching
4. **Flow Arrows:**
- Blue: Offload/prefetch operations
- Red: Mis-prediction correction path
- Black: Actual computation paths
### Detailed Analysis
**CPU Expert States:**
- Active (orange): e11, e12
- Idle (white): e13, e14
**GPU Expert States:**
- Active (red): e22, e23
- Idle (orange): e21, e24
**Queue Contents:**
- High-Priority Queues:
- Contains e23 (red) and e24 (orange)
- Shows "Offload unwanted e23" and "Fetch missed e24" operations
- Low-Priority Queues:
- Contains e11, e12, e21 (orange)
- Shows predictive caching of e22, e14, e13
**MoE Compute Block:**
- Actual computation shows:
- e22, e23 (red) as active
- e21, e24 (orange) as idle
- Input tokens t1, t2 with mis-prediction marker (×)
**Routing Path:**
- Pre-scheduled experts based on predicted routing path
- Shows mis-prediction in expert 2 (e2)
### Key Observations
1. **Resource Utilization:**
- GPU actively uses 50% of experts (e22, e23)
- CPU actively uses 50% of experts (e11, e12)
2. **Dynamic Adjustment:**
- System corrects mis-predictions in real-time (e23 offload, e24 fetch)
- Predictive caching pre-schedules experts but requires correction
3. **Priority Handling:**
- High-priority queues handle critical corrections
- Low-priority queues manage bulk prefetching
4. **Latency Management:**
- Prefetching reduces GPU idle time
- Real-time correction adds overhead but maintains accuracy
### Interpretation
This system demonstrates a sophisticated approach to hybrid computing resource management:
1. **Predictive Caching:** The low-priority queues pre-schedule experts based on predicted routing paths, reducing initial latency but requiring real-time adjustments when mis-predictions occur.
2. **Priority Queuing:** The separation into high/low-priority queues ensures critical corrections (like expert e23 offload) get immediate attention while bulk operations (e11-e14 prefetch) proceed in lower priority.
3. **Hybrid Optimization:** The system effectively balances CPU/GPU workloads by:
- Keeping 50% of experts active on each device
- Dynamically redistributing experts based on actual computation needs
- Maintaining 25% empty experts for potential overflow
4. **Error Handling:** The red × mis-prediction marker and correction arrows show the system's ability to adapt to prediction errors, though this introduces additional overhead.
The architecture suggests a trade-off between prediction accuracy and computational efficiency, with the real-time correction mechanism acting as a safety net for imperfect predictions. The 50/50 active/expert distribution across CPU/GPU indicates a balanced workload distribution strategy, while the 25% empty expert buffer provides headroom for unexpected computation demands.
</details>
Figure 3. The workflow of Expert Cache Engine (ECE). ECE pre-schedules the experts based on the routing path predictions before MoE computation. During MoE computation, ECE detects incorrect predictions, such as with unwanted expert $e_23$ and missed expert $e_24$ , and quickly prioritizes them in the high-priority queue for swapping during the execution of $e_22$ , ensuring efficient operation.
#### 3.2.1. Predictive Locality-aware Expert Caching (PLEC)
Traditional caching strategies like LRU (Eliseev and Mazur, 2023; Yuan et al., 2024; Xu et al., 2024; Aminabadi et al., 2022) often struggle in large-scale MoE models, leading to cache thrashing due to frequent data swaps. This inefficiency stems from LRU’s inability to account for the limitation of the maximum number of active experts per layer or their locality. For instance, Cache-MoE (Eliseev and Mazur, 2023) sets a fixed cache size for each layer, resulting in increased latency and computational overhead when the number of active experts exceeds this limit.
To this end, we propose Predictive Locality-aware Expert Caching (PLEC), which leverages routing path information predicted by the predictor to anticipate active experts and prefetch them efficiently into the cache. This approach prioritizes expert locality, minimizing I/O disruptions during computation. For example, in an MoE model with 4 layers and 4 experts per layer, a traditional LRU might assign 2 experts per layer for a total cache size of 8. In a worst-case scenario where all experts in each layer become active, LRU would result in constant expert swapping. However, PLEC can use the predicted routing path to prefetch all 8 experts from the first two layers, fully utilizing the cache and avoiding unnecessary data transfers. As the computation of the first layer completes, PLEC asynchronously prefetches experts for the third layer, reducing I/O overhead and improving throughput.
#### 3.2.2. Real-time Correction Mechanism
Although the predictor enables efficient cache scheduling, its inaccuracies can result in two types of experts: 1) mis-prefetched experts, which are idle experts incorrectly predicted as active and thus unnecessarily loaded into the cache, and 2) missed experts, which are active experts mistakenly predicted as idle, causing them to be absent from the cache. To address these cases, we introduce a real-time correction mechanism based on the MoE gate’s output during computation, with a focus on handling the missed experts. These experts are labeled as on-demand experts and assigned high priority for immediate loading into the cache.
The correction mechanism first attempts to replace any mis-prefetched experts with on-demand experts, since the former are idle and do not contribute to ongoing computations. If no mis-prefetched experts are available, the mechanism then replaces experts that have already completed their computations. This approach ensures that the cache is dynamically adjusted without disrupting the ongoing processing of correctly predicted experts, optimizing resource usage and minimizing the impact of prediction inaccuracies.
<details>
<summary>x4.png Details</summary>
### Visual Description
## Diagram: Token Routing Optimization in Batch Processing
### Overview
The diagram compares **original** and **optimized** expert load configurations in a token routing system. It illustrates how token-level and batch-level routing paths are restructured to improve computational efficiency. Two batches (Batch1, Batch2) are processed through grids representing token-expert assignments, with visual emphasis on activated (pink) vs. inactive (white) experts.
---
### Components/Axes
1. **Routing Paths**:
- **Token-level Routing Path**: Arrows connecting token grids to a "Token Scheduler" (central element).
- **Batch-level Routing Path**: Arrows connecting batch grids to a "New" routing output.
2. **Grids**:
- **Original Expert Loads**:
- Two grids (Batch1, Batch2) with 4x4 cells (tokens t1-t8, experts e1-e4).
- Pink cells indicate activated experts (all experts active per batch).
- **Optimized Expert Loads**:
- Two grids (New Batch1, New Batch2) with 4x4 cells.
- Pink cells indicate activated experts (half per batch).
3. **Labels**:
- **Original**: "Activate all experts per batch", "Each expert compute 1 token".
- **Optimized**: "Activate half experts per batch", "Each expert compute 2 tokens".
---
### Detailed Analysis
1. **Original Expert Loads**:
- **Batch1**: 4 experts (e1-e4) active, each handling 1 token (t1-t4).
- **Batch2**: 4 experts (e5-e8) active, each handling 1 token (t5-t8).
- **Token Scheduler**: Distributes tokens evenly across all experts.
2. **Optimized Expert Loads**:
- **New Batch1**: 2 experts (e1, e2) active, each handling 2 tokens (t1, t2 and t7, t3).
- **New Batch2**: 2 experts (e3, e4) active, each handling 2 tokens (t6, t4 and t8, t5).
- **Token Scheduler**: Consolidates token assignments to reduce active experts.
---
### Key Observations
1. **Expert Activation**:
- Original: 100% expert utilization (4 experts/batch).
- Optimized: 50% expert utilization (2 experts/batch).
2. **Token Computation**:
- Original: 1 token/expert.
- Optimized: 2 tokens/expert (doubled efficiency).
3. **Routing Changes**:
- Token-level routing remains unchanged (e1-e4).
- Batch-level routing shifts to "New" batches with consolidated expert assignments.
---
### Interpretation
The diagram demonstrates a **resource optimization strategy** for parallel processing:
- **Efficiency Gains**: By halving active experts and doubling tokens per expert, the system reduces overhead (e.g., memory, inter-expert communication) while maintaining throughput.
- **Scalability**: The "Token Scheduler" acts as a load balancer, ensuring tokens are distributed to maximize expert utilization without redundancy.
- **Trade-offs**: Fewer active experts may limit parallelism but improves per-expert throughput, ideal for memory-constrained environments.
This approach aligns with principles of **model parallelism optimization**, critical for deploying large language models (LLMs) on distributed hardware.
</details>
Figure 4. An example of the Token Scheduler applied to a single MoE layer with 4 experts $e_i$ , showing two batches, each with 4 tokens $t_i$ . Left: The cache mechanism becomes inefficient. Right: The tokens with similar expert selection are grouped into the same batch to reduce the number of active experts and increase the token load per expert, enhancing cache hit ratio and computational efficiency.
### 3.3. Token Scheduler
we propose the Token Scheduler to address the challenges of excessive expert activation during prefilling and uneven token distribution during decoding. By utilizing predictions from a routing path predictor, it efficiently reorganizes tokens sharing similar routing paths into the same batch. This minimizes unnecessary expert activations and significantly improves the overall efficiency of inferencing MoE models.
Fig. 4 illustrates the workflow of our Token Scheduler. Consider two input batches, each containing four tokens, processed through a single MoE layer with four experts for simplicity. In a standard MoE setup, as illustrated in Fig. 4 (left), tokens in each batch activate different experts, engaging all experts within the layer. Given a cache size limited to two experts per layer, this leads to two additional I/O operations to swap experts, causing substantial I/O overhead and computational delays.
In contrast, as shown in Fig. 4 (right), the two new batches optimized by our scheduler activate only half the number of experts compared to the original scenario, leading to a significant improvement in cache utilization. Furthermore, each expert is now assigned more tokens than before. In summary, the proposed Token Scheduler effectively reduces I/O overhead by decreasing expert loads and enhances computational efficiency by optimizing token assignments.
#### 3.3.1. Adaptive KV-Cache Management for Consistent MoE Performance
The Key-Value (KV) cache is crucial for storing intermediate representations linked to specific input tokens, thereby enhancing computational efficiency during inference. Given the autoregressive nature of Mixture of Experts (MoE) computations, it is vital to maintain consistency across time steps, as each output relies on the preceding KV cache. However, the Token Scheduler ’s rebatching of input tokens presents challenges in aligning the KV cache with the new token distribution. Such misalignment can result in inconsistencies in MoE outputs, highlighting the need for an adaptive approach to manage the KV cache effectively. To tackle this challenge, we introduce two key operations, Merge and Reindex:
1. Merge consolidates the KV cache from different batches based on a global indexing scheme. This ensures that despite the reorganization of input tokens, all necessary representations remain unified and accessible, maintaining coherence throughout the computational process.
1. Reindex dynamically adjusts the indices of KV cache entries to align with the positions of tokens in their new batches. This precise alignment guarantees that each token accesses the correct cache entries, preserving the accuracy of subsequent MoE computations.
These adaptive strategies, integrated with the Token Scheduler, ensure that MoE models maintain consistent and reliable performance, even as input token distributions shift dynamically during inference. They not only preserve the integrity of MoE computations but also enhance the model’s robustness and efficiency in handling real-world data.
#### 3.3.2. Token Scheduling Problem Definition
As illustrated in Fig. 4, we consider two batches, each containing $T$ tokens. Each token $i$ in a batch is associated with a binary routing path matrix $r_i$ of size $(L,E)$ . We assign a global index to each token across both batches, forming the set $T=\{1,2,…,2T\}$ . The goal of the Token Scheduler is to reorganize these tokens into two new batches, producing two sets of indices, $T_1$ and $T_2$ , such that $|T_1|=|T_2|=T$ and $T_1∪T_2=T$ . The batch-level routing path matrices $R_1$ and $R_2$ for the new batches are obtained by performing an element-wise AND ( $\bigwedge$ ) operation on the individual token-level routing path matrices:
$$
R_1=\bigwedge_i∈T_1r_i, R_2=\bigwedge_i∈\mathcal
{T_2}r_i. \tag{4}
$$
The optimization objective is to minimize the total number of activated experts across all layers in the new batch routing path matrices:
$$
\min_T_1,\mathcal{T_2}≤ft(∑_i=1^L∑_j=1^ER_1^i,j+∑_i=1^L∑_j=1^ER_2^i,j\right), \tag{5}
$$
where each item in Eq. 5 counts the total number of activated experts in a batch. The solution involves selecting token sets $T_1$ and $T_2$ that balance the token distribution to minimize expert activation, thereby improving computational efficiency and reducing resource waste.
#### 3.3.3. K-Means Clustering for Fast Token Rebatching
Given the complexity of the original optimization problem, finding an exact solution within limited time impractical, as the token scheduler must operate efficiently. We propose a fast heuristic solution based on K-means clustering, which leverages the similarity between token routing path matrices. This approach approximates a near-optimal solution and significantly improves inference speed, as demonstrated by our experiments.
We first calculate the similarity matrix $S$ between all pairs of token routing path matrices using the Hamming distance. Formally, for tokens $i$ and $j$ , the similarity $S_ij$ is defined as:
$$
S_ij=1-\frac{d_ij}{L× E}, \tag{6}
$$
where $d_ij$ is the Hamming distance between the routing path matrices of tokens $r_i$ and $r_j$ . The resulting $(2T× 2T)$ similarity matrix $S$ is used as input for the K-means clustering algorithm, detailed in the appendix.
K-means clustering begins with random centroid selection, followed by assigning each token to the cluster with the highest similarity to its centroid. To maintain balanced clusters, adjustments are made when a cluster exceeds the target size $T$ , ensuring each cluster contains exactly $T$ tokens. Once the tokens are assigned, the centroids are updated by selecting the token within each cluster that has the highest average similarity to the others. This iterative process continues until convergence or a preset iteration limit is reached. The output is two balanced sets of tokens, $T_1$ and $T_2$ , each with $T$ tokens. By clustering tokens based on their routing path similarity, our scheduler efficiently minimizes the number of activated experts, improving computational efficiency and reducing resource usage.
<details>
<summary>x5.png Details</summary>
### Visual Description
## Diagram: Expert Placement and Computation Pipeline in MoE Layer
### Overview
The diagram illustrates the workflow of expert placement and computation in a Mixture of Experts (MoE) layer, comparing two pipeline strategies: **Sequential Offload Pipeline** and **Multi-Stream Overlapping Pipeline**. It uses color-coded blocks to represent expert states (CPU/GPU), compute streams, and communication streams over time.
---
### Components/Axes
- **Legend** (top-left):
- White: Experts already on CPU
- Pink: Experts already on GPU
- Purple: Expert fetching (CPU→GPU)
- Green: Expert offloading (GPU→CPU)
- **X-axis**: Labeled "Time" (left to right).
- **Y-axis**: Two labeled streams:
1. **Compute Stream** (top row)
2. **Communication Stream** (bottom row)
---
### Detailed Analysis
#### Sequential Offload Pipeline (Top Section)
- **Compute Stream**:
- **Prev-layer computation** (white block) spans expert 5 (time 5).
- **Expert fetching** (purple blocks) for experts 2, 3, 4 (times 2, 3, 4).
- **Communication Stream**:
- **GPU→CPU** (green blocks) for experts 6, 7, 8 (times 6, 7, 8).
- **CPU→GPU** (pink blocks) for experts 2, 3, 4 (times 2, 3, 4).
#### Multi-Stream Overlapping Pipeline (Bottom Section)
- **Compute Stream**:
- **Prev-layer computation** (white block) spans experts 5, 2, 3, 4 (times 5, 2, 3, 4).
- **Communication Stream**:
- **GPU→CPU** (green blocks) for experts 6, 7, 8 (times 6, 7, 8).
- **CPU→GPU** (purple blocks) for experts 2, 3, 4 (times 2, 3, 4).
---
### Key Observations
1. **Sequential Offload Pipeline**:
- Experts 5, 6, 7, 8 are prefetched to GPU during prev-layer computation (time 5).
- Experts 2, 3, 4 are fetched to GPU after prev-layer computation (times 2–4).
- Expert offloading (GPU→CPU) occurs sequentially for experts 6, 7, 8 (times 6–8).
2. **Multi-Stream Overlapping Pipeline**:
- Prefetching of experts 2, 3, 4 occurs **during** prev-layer computation (overlap at times 2–4).
- Communication streams for experts 6, 7, 8 and 2, 3, 4 overlap in time, enabling parallelism.
---
### Interpretation
- **Efficiency Gains**: The Multi-Stream Overlapping Pipeline reduces idle time by overlapping expert fetching/offloading with prev-layer computation. For example, experts 2, 3, 4 are prefetched while prev-layer computation runs for expert 5.
- **Routing Path Predictions**: The diagram implies that experts are prefetched based on predicted routing paths during prev-layer computation, enabling proactive resource allocation.
- **Color Consistency**: All legend colors (white, pink, purple, green) are strictly adhered to in the diagram, ensuring clarity in expert states and data flow.
This visualization highlights how overlapping compute and communication streams in MoE layers can optimize throughput by minimizing sequential bottlenecks.
</details>
Figure 5. Comparison between sequential offload pipeline and our Multi-Stream Overlapping Pipeline.
## 4. Inference Pipeline Optimization
In this section, we outline two pipeline optimization strategies designed to enhance the efficiency of our three key components. This approach ensures that any potential overhead introduced by these components is effectively managed and minimized.
### 4.1. Multi-Stream Overlapping Pipeline
The Multi-Stream Overlapping Pipeline, illustrated in Fig. 5, addresses I/O bottlenecks inherent in MoE layer computations. Unlike traditional sequential offload pipelines that struggle with dynamic expert selection, our approach combines efficient stream management with predictive routing to optimize resource utilization.
Our system leverages predictive routing information for proactive expert management. As shown in Fig. 5, when transitioning from experts $[5,6,7,8]$ in the previous step to $[2,3,4,5]$ in the current step, our pipeline prefetches required experts during the previous step’s computation. This strategy ensures that necessary experts are available in the GPU cache, reducing I/O-induced latencies and smoothing inter-step transitions. Furthermore, we exploit PCIe’s bidirectional bandwidth to enable concurrent fetch and offload operations. This multi-stream approach, guided by predictive routing, significantly mitigates the impact of dynamic expert scheduling on system performance, offering a more efficient solution for MoE computations.
### 4.2. Dual-batch Inference Pipeline
To address the computational overhead introduced by the Routing Path Predictor and Token Scheduler, we propose an innovative Dual-Batch Inference Pipeline, as illustrated in Fig. 6. This pipeline architecture treats every two batches as a fundamental scheduling unit, necessitated by the Token Scheduler ’s requirement to reorganize input tokens between batches. The primary innovation of our tailored pipeline lies in its ability to overlap routing path prediction, token scheduling, and primary decoding tasks.
By interleaving routing path prediction and token scheduling within each dual-batch unit, these operations are executed concurrently with the main decoding process. This parallelization strategy effectively minimizes the additional overhead incurred from prediction and scheduling tasks, resulting in enhanced communication overlap and optimized computational efficiency. Consequently, the routing path prediction and token scheduling processes are seamlessly integrated into the overall inference workflow, facilitating high throughput and efficient resource utilization across all batches.
<details>
<summary>x6.png Details</summary>
### Visual Description
## Diagram: Pipeline Architecture Comparison
### Overview
The image compares two pipeline architectures for processing tasks: **Sequential Pipeline** and **Dual-batch Pipeline**. It visualizes stages of processing over time, highlighting differences in workflow, parallelism, and optimization techniques.
### Components/Axes
- **X-axis**: Labeled "Time" with an arrow indicating progression from left to right.
- **Legend**:
- **Purple**: Prefilling
- **Blue**: Decoding (original)
- **Pink**: Decoding (optimized by rebatched tokens)
- **Gray**: Routing Path Prediction
- **Yellow**: Token Scheduling
- **Key Elements**:
- **Sequential Pipeline**: Top section with stages 1–4 repeated sequentially.
- **Dual-batch Pipeline**: Bottom section with stages 1–4 split into sub-stages (e.g., 1,2; 3,4) and overlapping processes.
### Detailed Analysis
#### Sequential Pipeline
- **Structure**:
- Stages 1–4 are processed sequentially, with each stage repeated 4 times (e.g., "1 1 1 1" for stage 1).
- **Color Coding**:
- **Blue**: Decoding (original) dominates, indicating no optimization.
- **Trends**:
- Linear progression with no overlap between stages.
#### Dual-batch Pipeline
- **Structure**:
- Stages 1–4 are split into sub-stages (e.g., "1,2" for stages 1 and 2).
- Overlapping processes:
- **Prefilling** (purple) occurs in early stages (1–2).
- **Routing Path Prediction** (gray) and **Token Scheduling** (yellow) appear intermittently.
- **Decoding** alternates between **blue** (original) and **pink** (optimized).
- **Trends**:
- Parallelism: Sub-stages (e.g., 1,2 and 3,4) overlap in time.
- Optimization: Pink blocks (optimized decoding) appear in later stages, suggesting rebatched tokens improve efficiency.
### Key Observations
1. **Sequential vs. Dual-batch**:
- The Dual-batch Pipeline introduces parallelism (e.g., overlapping sub-stages) and optimization (pink decoding), unlike the strictly sequential blue decoding in the Sequential Pipeline.
2. **Token Scheduling**:
- Yellow blocks (Token Scheduling) are concentrated in the Dual-batch Pipeline, indicating dynamic resource allocation.
3. **Routing Path Prediction**:
- Gray blocks (Routing Path Prediction) appear in the Dual-batch Pipeline, suggesting adaptive decision-making for task routing.
### Interpretation
- **Efficiency Gains**: The Dual-batch Pipeline’s use of rebatched tokens (pink decoding) and parallel sub-stages reduces idle time, improving throughput compared to the Sequential Pipeline.
- **Complexity Trade-off**: While the Dual-batch Pipeline adds complexity (e.g., routing, scheduling), it optimizes resource utilization, likely critical for high-throughput systems.
- **Optimization Impact**: The shift from blue to pink decoding in later stages highlights the performance benefits of rebatching, though the exact timing and frequency of optimization are not quantified.
### Uncertainties
- No numerical values or time intervals are provided, making quantitative analysis impossible.
- The exact mechanism of "rebatched tokens" and its impact on decoding speed is implied but not explicitly defined.
</details>
Figure 6. Overview of our Dual-batch Inference Pipeline. In our pipeline, two batches are processed as a basic unit. The rebatched tokens enable a faster decoding stage (in pink) compared to the original sequential pipeline (in blue), thanks to the optimization of the Token Scheduler. Besides, the Routing path prediction (in green) and token scheduling (in yellow) are interleaved, minimizing their overhead and ensuring better resource utilization (best viewed in color).
## 5. Evaluation
In this section, we first present the experimental setup (§ 5.1) and the performance of routing path predictor (§ 5.2). We then quantitatively demonstrate ExpertFlow ’s effectiveness in improving inference speed (§ 5.3) and memory saving (§ 5.4) on various large-scale MoE models. Finally, we conduct extensive ablation studies to evaluate the impact of different crucial system components.
### 5.1. Setup
Hardware Settings: Our experiments were conducted on a single NVIDIA A40 GPU with 48 GB of memory and Intel(R) Xeon(R) Gold 6338 CPU @ 2.00GHz.
Tasks and MoEs: We evaluate on two tasks: translation (WMT16 (Bojar et al., 2016)) and summarization (XSUM (Narayan et al., 2018)), using four MoE models: Mixtral-8 $×$ 7B (Jiang et al., 2024) and Switch-32/64/128 with 32, 64, and 128 experts per layer, respectively (Fedus et al., 2022).
Routing Path Datasets: We created eight datasets by pairing above tasks and MoEs. For each dataset, we sampled 10,000 input sequences, each processed three times through the MoE model to generate output tokens and routing paths. The evaluations on our predictor and system performance are all conducted based on these datasets.
Baselines & Our Methods: We explored four baselines: 1) All-in-GPU (AIG), which keeps all experts in GPU memory but becomes impractical in resource-limited scenario. 2) SE-MoE (Shen et al., 2022) offloads all experts to CPU memory and prefetches the next layer’s experts during the current layer’s execution. However, it incurs high latency and memory overhead, especially when scaling to MoE models with a large number of experts per layer. 3) Cache-MoE (Eliseev and Mazur, 2023) employs a fixed-size cache per layer, using a Least-Recently-Used (LRU) strategy to prioritize active experts, with CPU fetches as fallback, though this may lead to I/O bottlenecks. 4) Pregated-MoE (Hwang et al., 2023) replaces original routers with predictive MLPs for proactive expert selection.
### 5.2. Predictor Evaluation
Evaluation Metrics. For an MoE with $L$ layers and $E$ experts per layer, we introduce batch-level accuracy ( $B_acc$ ) to evaluate prediction performance:
$$
\displaystyleB_{acc} \displaystyle=\frac{1}{L}∑_i=0^L-1\frac{∑_j=0^E-1I(R_
batch[i,j]=1 \& \hat{R}_batch[i,j]=1)}{∑_j=0^E-1
I(R_batch[i,j]=1)} \tag{7}
$$
where $R_batch$ and $\hat{R}_batch$ indicate true and predicted batch-level routing paths, respectively. Both are binary matrices of size $(L,E)$ , derived by aggregating token-level routing paths $r_i$ using a bit-wise OR operation: $R_batch=r_1\bigvee r_2\dots\bigvee r_n$ .
Table 1. Batch-level accuracy (%) of the Temporal-Locality-based Prediction (TLP), Spatial-Locality-based Prediction (SLP), and our transformer-based prediction on different Switch MoE models and tasks and batch size (BS).
| Swit ch- | BS | WMT16 | XSUM | | | | |
| --- | --- | --- | --- | --- | --- | --- | --- |
| TLP | SLP | Ours | TLP | SLP | Ours | | |
| 32 | 8 | 6.93 | 11.05 | 78.35 | 6.30 | 16.16 | 81.79 |
| 16 | 14.53 | 17.05 | 80.81 | 16.41 | 23.05 | 83.87 | |
| 32 | 31.05 | 23.09 | 83.65 | 36.51 | 28.03 | 87.01 | |
| 64 | 8 | 6.02 | 7.03 | 85.78 | 4.13 | 7.22 | 75.01 |
| 16 | 13.98 | 11.89 | 86.08 | 11.13 | 10.03 | 75.48 | |
| 32 | 31.32 | 16.87 | 86.94 | 27.74 | 12.28 | 76.91 | |
| 128 | 8 | 5.62 | 3.31 | 76.07 | 3.29 | 4.19 | 73.55 |
| 16 | 11.13 | 5.73 | 76.55 | 8.94 | 6.79 | 74.63 | |
| 32 | 24.15 | 9.93 | 77.84 | 24.74 | 11.71 | 76.44 | |
Predictor Settings. After extensive testing of various model configurations (see appendix for details), we identified the optimal setup with a feed-forward network dimensionality of $d_ff=2048$ and a hidden layer size of $d_model=32$ . This configuration balances between model size (7.21 MB) and performance, and was adopted for all experiments. We trained the predictor from scratch, using the Adam optimizer with a learning rate of 0.001, weight decay of 0.01, a cosine learning rate scheduler, and a batch size of 16.
Predictor Performance and Analysis. We design two baselines: spatial-locality-based prediction (SLP) and temporal-locality-based prediction (TLP). The former is based on the observation in (Li et al., 2023a), where tokens selecting the same expert in one layer tend to select similar experts in the subsequent layer. Similarly, we propose TLP, which assumes that experts activated in recent MoE computations are likely to be used again in the near future.
Table 1 reports batch-level accuracy for various batch sizes, where our predictor consistently outperforms both TLP and SLP across all settings, demonstrating superior generalization and adaptability across diverse applications.
### 5.3. Inference Speed Evaluation
Results on Switch MoE Series. In Fig. LABEL:fig:switch_speedup, our ExpertFlow system exhibits significant speedup across the Switch MoE series, with a peak of 9.99 $×$ on the WMT16 dataset and 9.54 $×$ on XSUM, compared to SE-MoE. Notably, as the model scales from Switch-32 to Switch-128, the average speedup increases from 2.53 $×$ to 7.94 $×$ . This trend highlights that our system becomes more effective as the number of experts increases, due to its advanced expert management that optimizes the handling of larger models and the associated I/O overhead.
<details>
<summary>x7.png Details</summary>
### Visual Description
## Bar Chart: Relative Speedup Ratio of Cache-MoE vs ExpertFlow
### Overview
The chart compares the relative speedup ratio of two methods, **Cache-MoE** (orange) and **ExpertFlow** (pink with diagonal stripes), across different computational configurations. The x-axis represents parameter combinations (e.g., (1,4), (1,8), etc.), while the y-axis shows the speedup ratio relative to a baseline (1.0). Data is split into two domains: **WMT16** (left) and **XSUM** (right).
---
### Components/Axes
- **X-Axis**:
- **Categories**:
- **WMT16**: (1,4), (1,8), (1,16), (3,4), (3,8), (3,16)
- **XSUM**: (1,4), (1,8), (1,16), (3,4), (3,8), (3,16)
- **Labels**: "(CS,BS) = (1,4)", "(1,8)", etc.
- **Y-Axis**:
- **Title**: "Relative Speedup Ratio"
- **Scale**: 1.0 to 2.0 (increments of 0.1)
- **Legend**:
- **Top-right**:
- **Cache-MoE**: Orange square
- **ExpertFlow**: Pink bar with diagonal stripes
---
### Detailed Analysis
#### WMT16 Domain
- **Cache-MoE**:
- All bars hover near **1.0** (e.g., 1.00x for (1,4), 1.01x for (1,8), etc.).
- **ExpertFlow**:
- Speedup increases with configuration size:
- (1,4): 1.83x
- (1,8): 1.86x
- (1,16): 1.87x
- (3,4): 1.78x
- (3,8): 1.90x
- (3,16): 1.96x
#### XSUM Domain
- **Cache-MoE**:
- Consistent near **1.0** (e.g., 1.00x for all configurations).
- **ExpertFlow**:
- Higher speedup than WMT16:
- (1,4): 1.84x
- (1,8): 1.87x
- (1,16): 1.87x
- (3,4): 1.86x
- (3,8): 1.96x
- (3,16): 1.99x
---
### Key Observations
1. **ExpertFlow Dominates**:
- ExpertFlow consistently outperforms Cache-MoE across all configurations, with speedup ratios ranging from **1.78x to 1.99x**.
2. **Scaling Trends**:
- Larger configurations (e.g., (3,16)) yield higher speedups for ExpertFlow, especially in XSUM.
- Cache-MoE shows minimal improvement with scaling, remaining close to 1.0.
3. **Domain Differences**:
- XSUM achieves higher speedups than WMT16 for ExpertFlow, suggesting domain-specific efficiency gains.
---
### Interpretation
- **Performance Implications**:
- ExpertFlow’s ability to scale with configuration size (e.g., 1.99x in XSUM (3,16)) highlights its efficiency for complex tasks. Cache-MoE’s stagnant performance suggests limitations in handling larger parameter sets.
- **Methodology Insights**:
- The diagonal stripes in ExpertFlow’s bars may indicate dynamic resource allocation or parallelization advantages over Cache-MoE’s static approach.
- **Outliers**:
- The (3,4) configuration in WMT16 (1.78x) is an outlier for ExpertFlow, underperforming compared to other setups. This could indicate suboptimal tuning for specific parameter combinations.
---
### Spatial Grounding & Verification
- **Legend Alignment**:
- Orange squares (Cache-MoE) match all orange bars; pink striped bars (ExpertFlow) align with their legend entry.
- **Trend Verification**:
- ExpertFlow’s upward trend in XSUM (1.84x → 1.99x) is visually confirmed by increasing bar heights.
- **Component Isolation**:
- WMT16 and XSUM are spatially separated, with distinct y-axis ranges for clarity.
---
### Conclusion
The chart demonstrates that **ExpertFlow** significantly outperforms **Cache-MoE** in both domains, with speedup ratios increasing as configurations scale. This suggests ExpertFlow is better suited for resource-intensive applications, while Cache-MoE may struggle with larger parameter sets. Further investigation into the (3,4) WMT16 outlier could reveal optimization opportunities.
</details>
Figure 8. Comparative analysis of cache hit ratio (scatter plot) and normalized MoE inference speedup (bar chart) across different configurations of cache size (CS) and batch size (BS) on the Mixtral-8 $×$ 7B MoE model.
Table 2. GPU memory (GB) savings across different cache size (CS) and batch size (BS), compared to allocating all experts in GPU (AIG). Cache-MoE and ExpertFlow w/o Predictor incur the same costs.
| Task | Switch- | CS | BS | AIG | Pregated-MoE | SE-MoE | ExpertFlow with Predictor | ExpertFlow w/o Predictor | | | | |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Mem | Mem | Sav. (%) | Mem | Sav. (%) | Mem | Sav. (%) | Mem | Sav. (%) | | | | |
| XSUM | 32 | 4 | 8 | 3.99 | 1.45 | 63.76 | 1.76 | 55.82 | 1.02 | 74.30 | 0.93 | 76.58 |
| 8 | 16 | 4.08 | 2.42 | 40.60 | 2.60 | 36.31 | 1.58 | 61.33 | 1.43 | 64.94 | | |
| 16 | 32 | 4.17 | 3.59 | 13.80 | 3.79 | 9.13 | 2.70 | 35.13 | 2.41 | 42.07 | | |
| 64 | 4 | 8 | 7.74 | 2.18 | 71.85 | 2.53 | 67.37 | 1.03 | 86.74 | 0.94 | 87.82 | |
| 8 | 16 | 7.85 | 2.94 | 62.49 | 3.29 | 58.12 | 1.60 | 79.63 | 1.44 | 81.62 | | |
| 16 | 32 | 7.96 | 4.18 | 47.42 | 4.55 | 42.85 | 2.70 | 66.06 | 2.44 | 69.32 | | |
| 128 | 4 | 8 | 15.26 | 3.64 | 76.12 | 4.05 | 73.45 | 1.01 | 93.35 | 0.96 | 93.72 | |
| 8 | 16 | 15.36 | 3.74 | 75.63 | 4.15 | 73.01 | 1.58 | 89.71 | 1.47 | 90.43 | | |
| 16 | 32 | 15.50 | 5.24 | 66.17 | 5.67 | 63.42 | 2.71 | 82.54 | 2.50 | 83.88 | | |
| Average Memory Saving (%) | | 57.54 | | 53.28 | | 74.31 | | 76.71 | | | | |
| WMT16 | 32 | 4 | 8 | 3.95 | 1.38 | 65.02 | 1.72 | 56.39 | 0.94 | 76.17 | 0.92 | 76.72 |
| 8 | 16 | 3.97 | 1.39 | 65.02 | 1.76 | 55.62 | 1.43 | 64.04 | 1.39 | 65.08 | | |
| 16 | 32 | 4.03 | 2.47 | 38.54 | 2.68 | 33.47 | 2.37 | 41.05 | 2.30 | 42.84 | | |
| 64 | 4 | 8 | 7.70 | 2.12 | 72.54 | 2.39 | 68.93 | 0.95 | 87.68 | 0.92 | 88.06 | |
| 8 | 16 | 7.72 | 2.12 | 72.52 | 2.48 | 67.90 | 1.43 | 81.46 | 1.39 | 81.99 | | |
| 16 | 32 | 7.78 | 2.28 | 70.66 | 2.66 | 65.85 | 2.38 | 69.45 | 2.30 | 70.37 | | |
| 128 | 4 | 8 | 15.21 | 3.58 | 76.46 | 4.01 | 73.62 | 0.96 | 93.72 | 0.93 | 93.91 | |
| 8 | 16 | 15.23 | 3.59 | 76.44 | 4.00 | 73.70 | 1.43 | 90.58 | 1.39 | 90.85 | | |
| 16 | 32 | 15.28 | 3.75 | 75.48 | 4.18 | 72.63 | 2.40 | 84.30 | 2.33 | 84.78 | | |
| Average Memory Saving (%) | | 68.08 | | 63.12 | | 76.49 | | 77.18 | | | | |
Another key observation is the system’s performance under reduced cache and batch sizes, where our caching strategy shows clear advantages. When both cache size and batch size are proportionally decreased, the performance gains become more pronounced, particularly in configurations with smaller caches. This emphasizes the effectiveness of our dynamic caching mechanism, which ensures sustained high performance even under tight cache constraints.
Results on Mixtral-8 $×$ 7B. As shown in Fig. 8, our system achieves notable speedups on the Mixtral-8 $×$ 7B model, with a 1.96 $×$ speedup on WMT16 and 1.99 $×$ on XSUM, further demonstrating its efficiency compared to the baseline. In this experiment, we fixed the cache size while varying the batch size, which allowed us to test how well the system handles increased expert transfers. Typically, increasing the batch size would lead to more frequent expert transfers, which could slow down the system. However, our results show that the performance improvement further improves under these conditions, reaffirming the effectiveness of our caching strategy. This behavior demonstrates that even with frequent expert transfers, our system can efficiently manage expert activation without sacrificing speed. Overall, our system excels across various scenarios, including large model sizes, reduced cache sizes, and increased batch sizes, proving the robustness of our caching mechanism.
### 5.4. Memory Cost Evaluation
This section compares the efficiency of GPU memory savings achieved by various methods applied to Switch series and Mixtral-8 $×$ 7B MoE models. The assessment of GPU memory costs includes model parameters, input/output data, and intermediate activations such as key-value (KV) caches.
Results on Switch MoE Series. Table 2 presents the memory savings achieved under various configurations of Switch models, tasks, and settings, including cache size (CS) and batch size (BS). The cache size indicates the proportion of experts that can be simultaneously resident on the GPU. For instance, a cache size of 16 for Switch-32 means that half of the experts can be accommodated on the GPU, though the effectiveness of this caching varies among different methods.
Our ExpertFlow w/o predictor demonstrates substantial memory savings across a range of MoE models and configurations, achieving an average GPU memory saving of 74.31% (with a peak of 93.35%) on XSUM and 76.49% (with a peak of 93.72%) on WMT16. Although integrating a routing path predictor slightly increases memory usage compared to, the performance improvements (as detailed in Fig. LABEL:fig:switch_speedup) justify this trade-off, highlighting the advantages of intelligent routing mechanisms.
As MoE models scale from Switch-32 to Switch-128, all evaluated methods can increase the memory savings due to more experts being offloaded to the CPU, highlighting the need for scalable memory management strategies. However, in smaller models like Switch-32, increasing cache and batch sizes can drastically reduce memory savings, with Pregated-MoE experiencing a drop from 63.76% to 13.80%. This demonstrates the sensitivity of smaller models to caching mechanisms and batch configurations, indicating a need for customized strategies. To delve deeper into these effects, we conducted an ablation study (§ 5.5) on how cache and batch size variations impact memory efficiency in smaller MoE models, aiming to gain further insights into the factors affecting memory savings in MoE systems.
Results on Mixtral-8 $×$ 7B MoE. Due to limited compatibility with most baseline methods, our evaluation of the Mixtral-8 $×$ 7B model, which requires 96 GB of GPU memory for storage and causes out-of-memory (OOM) errors for the All-In-GPU (AIG) method, focuses on comparisons with Cache-MoE for memory savings. We tested different cache sizes while keeping the batch size constant at 16. Our results mirrored those from Switch models, with both Cache-MoE and our system using around 16 GB and 37 GB of GPU memory for cache sizes of 1 and 3, respectively. These findings highlight our ExpertFlow system’s effective memory management, even in high-demand scenarios like the Mixtral-8 $×$ 7B model.
### 5.5. Ablation Study on System Components
Impact of Expert Cache Engine on Cache Hit Ratio. We conducted an ablation study to evaluate the effectiveness of our Expert Cache Engine, comparing it with the conventional LRU caching strategy. This study focused on the WMT16 dataset using the Switch-32 model, with results illustrated in Fig. 9. The results reveal that increasing cache size consistently improves the hit ratio, while larger batch sizes under a fixed cache size lead to decreased hit ratios. This aligns with the expectation that larger batch sizes tend to activate more experts concurrently, leading to a higher likelihood of exceeding cache capacity.
<details>
<summary>x8.png Details</summary>
### Visual Description
## Line Chart: Hit Ratio vs. Batch Size
### Overview
The chart compares the hit ratio performance of two methods ("Ours" and "LRU") across different cache sizes (CS=8 and CS=16) and batch sizes (4, 8, 16). The y-axis represents hit ratio as a percentage, while the x-axis represents batch size. Four data series are plotted: "Ours (CS=8)", "LRU (CS=8)", "Ours (CS=16)", and "LRU (CS=16)".
### Components/Axes
- **X-axis (Batch Size)**: Labeled "Batch Size" with values 4, 8, and 16.
- **Y-axis (Hit Ratio %)**: Labeled "Hit Ratio (%)" with a range from 0.4 to 0.9.
- **Legend**: Positioned at the top-left, with four entries:
- Solid green line: "Ours (CS=8)"
- Dotted green line: "LRU (CS=8)"
- Solid red line: "Ours (CS=16)"
- Dotted red line: "LRU (CS=16)"
### Detailed Analysis
- **Ours (CS=8)**:
- Batch Size 4: 86.35%
- Batch Size 8: 79.88%
- Batch Size 16: 71.89%
- **Trend**: Gradual decline as batch size increases.
- **LRU (CS=8)**:
- Batch Size 4: 57.21%
- Batch Size 8: 44.90%
- Batch Size 16: 36.22%
- **Trend**: Steeper decline compared to "Ours (CS=8)".
- **Ours (CS=16)**:
- Batch Size 4: 76.61%
- Batch Size 8: 65.62%
- Batch Size 16: 58.37%
- **Trend**: Moderate decline, outperforming LRU (CS=16) but underperforming "Ours (CS=8)".
- **LRU (CS=16)**:
- Batch Size 4: 65.62%
- Batch Size 8: 58.37%
- Batch Size 16: 58.37% (no data point at 16, inferred from trend)
- **Trend**: Slight decline, but less pronounced than LRU (CS=8).
### Key Observations
1. **Performance Gap**: "Ours" consistently outperforms "LRU" across all batch sizes and cache sizes. For example:
- At batch size 4, "Ours (CS=8)" achieves 86.35% vs. "LRU (CS=8)" at 57.21% (29.14% gap).
- At batch size 16, "Ours (CS=8)" achieves 71.89% vs. "LRU (CS=8)" at 36.22% (35.67% gap).
2. **Batch Size Impact**: Hit ratios for all methods decrease as batch size increases, but "Ours" maintains higher performance.
3. **Cache Size Trade-off**: "Ours (CS=16)" has lower hit ratios than "Ours (CS=8)" but still outperforms "LRU (CS=16)".
### Interpretation
The data suggests that the "Ours" method is more robust to increasing batch sizes compared to "LRU". The performance gap widens with larger batch sizes, indicating that "Ours" may be better optimized for scalability. The decline in hit ratios for all methods as batch size increases implies a trade-off between batch size and efficiency, with "Ours" mitigating this trade-off more effectively. The choice of cache size (CS=8 vs. CS=16) also affects performance, with smaller cache sizes (CS=8) generally yielding higher hit ratios for "Ours".
</details>
Figure 9. Ablation study on the impact of batch size and cache size ( $CS$ ) on expert cache hit ratio for the Switch-32 Model, compared with LRU and our PLEC strategies.
Notably, our method maintains more stable hit ratios with larger cache sizes as batch sizes increase. At $CS=16$ , increasing $BS$ from 4 to 8 and 16 results in hit ratio declines of only 3.12% and 6.05%, respectively, compared to the LRU baseline’s more significant declines of 10.99% and 18.24%. This stability is attributed to our dynamic, Predictive Locality-aware Expert Caching strategy, which intelligently allocates cache resources across layers based on routing path prediction.
Impact of Token Scheduler. To quantify the contributions of the Token Scheduler to our system’s performance, we conducted detailed ablation studies on the Switch series MoE models, specifically utilizing the WMT16 dataset. These studies were performed under controlled conditions with a cache size set at 16 and a batch size of 32. The results, as presented in Table 3, reveal that the Token Scheduler significantly enhances system efficiency. Specifically, it achieves a maximum speedup of 1.19 $×$ for the Switch-64 model. Across the entire range of Switch MoE models tested, the Token Scheduler consistently delivered an average speedup of 1.15 $×$ . This improvement underscores the effectiveness of the Token Scheduler in optimizing processing times and managing computational resources more efficiently.
Table 3. Impact of the Token Scheduler (TS) on system inference speed.
| Method | Switch-32 | Switch-64 | Switch-128 |
| --- | --- | --- | --- |
| ExpertFlow w/o TS | 1 $×$ | 1 $×$ | 1 $×$ |
| ExpertFlow with TS | 1.09 $×$ | 1.19 $×$ | 1.16 $×$ |
Impact of Dual-Batch Inference Pipeline. To evaluate the efficiency of our Dual-Batch Inference Pipeline, we conducted a detailed timing analysis with a batch size of 32. The token scheduler operated independently at $3.35± 0.04$ ms, while the predictor’s one-step decoding time was measured at $8.02± 0.02$ ms. These times are notably brief when compared to the $94.1± 11.4$ ms required for Switch-128 inference. Crucially, when these components were executed concurrently within our pipeline, we observe no significant increase in overall time cost. This demonstrates the effectiveness of our pipeline design in maximizing parallel execution and minimizing overhead, thereby maintaining high efficiency even with the addition of predictive and token scheduling components.
## 6. Discussion
Adaptive Cache and Batch Size Management. Our results in Fig. 9 demonstrate an inverse relationship between cache size and batch size on hit ratio, with larger batch sizes decreasing hit ratio but increasing throughput. This suggests the potential for an adaptive system that adjusts these parameters in real-time, balancing hit ratio optimization with throughput maximization to adapt to varying workload conditions and optimize resource utilization.
Backup All Experts on CPU. To further enhance inference efficiency, future versions of our system could relocate the predictor to the CPU, freeing up valuable GPU memory for MoE model computations. Additionally, we have explored backing up all experts on the CPU, which eliminates the need for offloading and simplifies the system architecture. This approach reduces GPU memory overhead, requiring only an additional allocation equivalent to one expert’s size for the multi-stream overlapping pipeline. In environments with abundant CPU resources, this strategy could simplify implementation while maintaining efficiency and reducing GPU memory demands.
Boosting MoE Offloading Efficiency with CPU Computational Power. Our efforts to enhance MoE models by using the CPU’s memory storage have been promising. Inspired by the Fiddler method (Kamahori et al., 2024), we plan to further utilize CPU computational power. This strategy involves shifting expert layer computations to the CPU while keeping activations on the GPU, potentially boosting system performance. By integrating this with predictive routing, we aim to better distribute workloads between the CPU and GPU, reducing GPU load and improving memory efficiency. This approach prepares us for future improvements in MoE deployment in resource-limited settings.
## 7. Related Work
The MoE-based LLM inference process presents significant GPU memory challenges, prompting prior research to explore a range of algorithmic and system-level solutions that can be broadly categorized into two major strands:
Memory Compression
Distillation techniques (Shen et al., 2022; Dai et al., 2024) reduce the number of experts by compressing the teacher network into a smaller student network. Model pruning methods have also been explored, such as pruning non-essential experts during fine-tuning based on usage frequency (Chen et al., 2022) and merging similar experts followed by low-rank decomposition (Li et al., 2023b). Post-training quantization (Eliseev and Mazur, 2023; Li et al., 2024; Lin et al., 2024; Frantar et al., 2022) further reduces memory consumption by converting pre-trained models from high-precision (e.g., Float32) formats to lower-precision ones (e.g., Int4), without requiring extensive retraining. The contribution of our proposed ExpertFlow is orthogonal to this direction, and ExpertFlow can be seamlessly integrated with these techniques to further reduce GPU memory cost during MoE inference.
Memory Offloading
Memory offloading involves transferring model components or computations to less expensive storage or processing units, optimizing GPU usage. Early efforts, such as the ZeRO series (Rajbhandari et al., 2020; Ren et al., 2021) reduced memory by offloading optimizer states, gradients, and weights to CPUs or SSDs. Offloading strategies have since been extended to inference tasks (Song et al., 2023; Sheng et al., 2023; Chen et al., 2024). For instance, FlexGen (Sheng et al., 2023) implements a “zig-zag block schedule” to offload activations and key-value (KV) caches, enabling models like OPT-175B (Zhang et al., 2022) to run efficiently on a single 16GB GPU. Lamina (Chen et al., 2024) further improves efficiency by shifting attention computations from GPU to CPU. Despite these advancements, existing offloading techniques are primarily designed for dense LLMs and fail to fully address the flexible, dynamic expert routing in MoE models. Even there have emerged offloading methods tailored for MoE, they struggle with insufficient accuracy in routing paths prediction (Li et al., 2023a; Eliseev and Mazur, 2023) or high training costs for predictors (Hwang et al., 2023; Du et al., 2024), undermining their practicality. Our work presents a unified system that tackles the critical challenges of improving the accuracy of routing path predictions and reducing the training overhead of predictors simultaneously.
## 8. Conclusion
This paper introduces ExpertFlow, a novel system that significantly improves MoE model deployment by addressing critical memory and computational challenges. Our approach integrates three key components: the Routing Path Predictor, Expert Cache Engine, and Token Scheduler, which synergistically optimize MoE inference. By accurately predicting expert selection patterns, ExpertFlow enables proactive expert scheduling and efficient token re-batching, substantially reducing GPU memory usage while maintaining competitive inference speeds. Comprehensive evaluations across various datasets and MoE models demonstrate the effectiveness of our ExpertFlow system and showcase the potential of predictive modeling and intelligent resource management in deploying large-scale neural networks, opening new avenues for research in adaptive inference strategies and specialized hardware for MoE models. This work contributes to making advanced AI models more accessible and practical for real-world scenarios.
Acknowledgements. To Robert, for the bagels and explaining CMYK and color spaces.
## References
- (1)
- Aminabadi et al. (2022) Reza Yazdani Aminabadi, Samyam Rajbhandari, Ammar Ahmad Awan, Cheng Li, Du Li, Elton Zheng, Olatunji Ruwase, Shaden Smith, Minjia Zhang, Jeff Rasley, et al. 2022. Deepspeed-inference: enabling efficient inference of transformer models at unprecedented scale. In SC22: International Conference for High Performance Computing, Networking, Storage and Analysis. IEEE, 1–15.
- Bojar et al. (2016) Ondrej Bojar, Rajen Chatterjee, Christian Federmann, Yvette Graham, Barry Haddow, Matthias Huck, Antonio Jimeno Yepes, Philipp Koehn, Varvara Logacheva, Christof Monz, et al. 2016. Findings of the 2016 conference on machine translation (wmt16). In First conference on machine translation. Association for Computational Linguistics, 131–198.
- Brown et al. (2020) Tom B Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. 2020. Language models are few-shot learners. Advances in neural information processing systems 33 (2020), 1877–1901.
- Chen et al. (2021) Mark Chen, Jerry Tworek, Heewoo Jun, Qiming Yuan, Henrique Ponde De Oliveira Pinto, Jared Kaplan, Harri Edwards, Yuri Burda, Nicholas Joseph, Greg Brockman, et al. 2021. Evaluating large language models trained on code. arXiv preprint arXiv:2107.03374 (2021).
- Chen et al. (2024) Shaoyuan Chen, Yutong Lin, Mingxing Zhang, and Yongwei Wu. 2024. Efficient and Economic Large Language Model Inference with Attention Offloading. arXiv preprint arXiv:2405.01814 (2024).
- Chen et al. (2022) Tianyu Chen, Shaohan Huang, Yuan Xie, Binxing Jiao, Daxin Jiang, Haoyi Zhou, Jianxin Li, and Furu Wei. 2022. Task-specific expert pruning for sparse mixture-of-experts. arXiv preprint arXiv:2206.00277 (2022).
- Dai et al. (2024) Damai Dai, Chengqi Deng, Chenggang Zhao, RX Xu, Huazuo Gao, Deli Chen, Jiashi Li, Wangding Zeng, Xingkai Yu, Y Wu, et al. 2024. Deepseekmoe: Towards ultimate expert specialization in mixture-of-experts language models. arXiv preprint arXiv:2401.06066 (2024).
- Du et al. (2024) Zhixu Du, Shiyu Li, Yuhao Wu, Xiangyu Jiang, Jingwei Sun, Qilin Zheng, Yongkai Wu, Ang Li, Hai Li, and Yiran Chen. 2024. SiDA: Sparsity-Inspired Data-Aware Serving for Efficient and Scalable Large Mixture-of-Experts Models. Proceedings of Machine Learning and Systems 6 (2024), 224–238.
- Eliseev and Mazur (2023) Artyom Eliseev and Denis Mazur. 2023. Fast inference of mixture-of-experts language models with offloading. arXiv preprint arXiv:2312.17238 (2023).
- Fedus et al. (2022) William Fedus, Barret Zoph, and Noam Shazeer. 2022. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. Journal of Machine Learning Research 23, 120 (2022), 1–39.
- Frantar et al. (2022) Elias Frantar, Saleh Ashkboos, Torsten Hoefler, and Dan Alistarh. 2022. Gptq: Accurate post-training quantization for generative pre-trained transformers. arXiv preprint arXiv:2210.17323 (2022).
- Guo et al. (2024) Daya Guo, Qihao Zhu, Dejian Yang, Zhenda Xie, Kai Dong, Wentao Zhang, Guanting Chen, Xiao Bi, Yu Wu, YK Li, et al. 2024. DeepSeek-Coder: When the Large Language Model Meets Programming–The Rise of Code Intelligence. arXiv preprint arXiv:2401.14196 (2024).
- He et al. (2021) Xin He, Kaiyong Zhao, and Xiaowen Chu. 2021. AutoML: A survey of the state-of-the-art. Knowledge-based systems 212 (2021), 106622.
- Hwang et al. (2023) Ranggi Hwang, Jianyu Wei, Shijie Cao, Changho Hwang, Xiaohu Tang, Ting Cao, Mao Yang, and Minsoo Rhu. 2023. Pre-gated MoE: An Algorithm-System Co-Design for Fast and Scalable Mixture-of-Expert Inference. arXiv preprint arXiv:2308.12066 (2023).
- Jacobs et al. (1991) Robert A Jacobs, Michael I Jordan, Steven J Nowlan, and Geoffrey E Hinton. 1991. Adaptive mixtures of local experts. Neural computation 3, 1 (1991), 79–87.
- Jiang et al. (2024) Albert Q Jiang, Alexandre Sablayrolles, Antoine Roux, Arthur Mensch, Blanche Savary, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Emma Bou Hanna, Florian Bressand, et al. 2024. Mixtral of experts. arXiv preprint arXiv:2401.04088 (2024).
- Kamahori et al. (2024) Keisuke Kamahori, Yile Gu, Kan Zhu, and Baris Kasikci. 2024. Fiddler: CPU-GPU Orchestration for Fast Inference of Mixture-of-Experts Models. arXiv preprint arXiv:2402.07033 (2024).
- Kaplan et al. (2020) Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. 2020. Scaling laws for neural language models. arXiv preprint arXiv:2001.08361 (2020).
- Li et al. (2023a) Jiamin Li, Yimin Jiang, Yibo Zhu, Cong Wang, and Hong Xu. 2023a. Accelerating distributed $\{$ MoE $\}$ training and inference with lina. In 2023 USENIX Annual Technical Conference (USENIX ATC 23). 945–959.
- Li et al. (2024) Pingzhi Li, Xiaolong Jin, Yu Cheng, and Tianlong Chen. 2024. Examining Post-Training Quantization for Mixture-of-Experts: A Benchmark. arXiv preprint arXiv:2406.08155 (2024).
- Li et al. (2023b) Pingzhi Li, Zhenyu Zhang, Prateek Yadav, Yi-Lin Sung, Yu Cheng, Mohit Bansal, and Tianlong Chen. 2023b. Merge, then compress: Demystify efficient SMoe with hints from its routing policy. arXiv preprint arXiv:2310.01334 (2023).
- Lin et al. (2024) Ji Lin, Jiaming Tang, Haotian Tang, Shang Yang, Wei-Ming Chen, Wei-Chen Wang, Guangxuan Xiao, Xingyu Dang, Chuang Gan, and Song Han. 2024. AWQ: Activation-aware Weight Quantization for On-Device LLM Compression and Acceleration. Proceedings of Machine Learning and Systems 6 (2024), 87–100.
- Lu et al. (2021) Shun Lu, Jixiang Li, Jianchao Tan, Sen Yang, and Ji Liu. 2021. TNASP: A Transformer-based NAS Predictor with a Self-evolution Framework. Advances in Neural Information Processing Systems 34 (2021), 15125–15137.
- Narayan et al. (2018) Shashi Narayan, Shay B Cohen, and Mirella Lapata. 2018. Don’t give me the details, just the summary! Topic-aware convolutional neural networks for extreme summarization. In Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing. 1797–1807.
- OpenAI (2024) OpenAI. 2024. ChatGPT: Language Model. https://chat.openai.com. Accessed: 2024-10-03.
- Patel et al. (2024) Pratyush Patel, Esha Choukse, Chaojie Zhang, Aashaka Shah, Íñigo Goiri, Saeed Maleki, and Ricardo Bianchini. 2024. Splitwise: Efficient generative llm inference using phase splitting. In 2024 ACM/IEEE 51st Annual International Symposium on Computer Architecture (ISCA). IEEE, 118–132.
- Qin et al. (2024) Ruoyu Qin, Zheming Li, Weiran He, Mingxing Zhang, Yongwei Wu, Weimin Zheng, and Xinran Xu. 2024. Mooncake: A kvcache-centric disaggregated architecture for llm serving. URL https://arxiv. org/abs/2407.00079 (2024).
- Radford et al. (2018) Alec Radford, Karthik Narasimhan, Tim Salimans, and Ilya Sutskever. 2018. Improving language understanding by generative pre-training. OpenAI (2018).
- Radford et al. (2019) Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. 2019. Language models are unsupervised multitask learners. OpenAI blog 1, 8 (2019), 9.
- Raffel et al. (2020) Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. 2020. Exploring the limits of transfer learning with a unified text-to-text transformer. Journal of machine learning research 21, 140 (2020), 1–67.
- Rajbhandari et al. (2020) Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, and Yuxiong He. 2020. Zero: Memory optimizations toward training trillion parameter models. In SC20: International Conference for High Performance Computing, Networking, Storage and Analysis. IEEE, 1–16.
- Ramesh et al. (2022) Aditya Ramesh, Prafulla Dhariwal, Alex Nichol, Casey Chu, and Mark Chen. 2022. Hierarchical text-conditional image generation with clip latents. arXiv preprint arXiv:2204.06125 1, 2 (2022), 3.
- Ramesh et al. (2021) Aditya Ramesh, Mikhail Pavlov, Gabriel Goh, Scott Gray, Chelsea Voss, Alec Radford, Mark Chen, and Ilya Sutskever. 2021. Zero-shot text-to-image generation. In International conference on machine learning. Pmlr, 8821–8831.
- Ren et al. (2021) Jie Ren, Samyam Rajbhandari, Reza Yazdani Aminabadi, Olatunji Ruwase, Shuangyan Yang, Minjia Zhang, Dong Li, and Yuxiong He. 2021. $\{$ Zero-offload $\}$ : Democratizing $\{$ billion-scale $\}$ model training. In 2021 USENIX Annual Technical Conference (USENIX ATC 21). 551–564.
- Shazeer et al. (2017) Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoc Le, Geoffrey Hinton, and Jeff Dean. 2017. Outrageously large neural networks: The sparsely-gated mixture-of-experts layer. arXiv preprint arXiv:1701.06538 (2017).
- Shen et al. (2022) Liang Shen, Zhihua Wu, WeiBao Gong, Hongxiang Hao, Yangfan Bai, HuaChao Wu, Xinxuan Wu, Jiang Bian, Haoyi Xiong, Dianhai Yu, et al. 2022. Se-moe: A scalable and efficient mixture-of-experts distributed training and inference system. arXiv preprint arXiv:2205.10034 (2022).
- Sheng et al. (2023) Ying Sheng, Lianmin Zheng, Binhang Yuan, Zhuohan Li, Max Ryabinin, Beidi Chen, Percy Liang, Christopher Ré, Ion Stoica, and Ce Zhang. 2023. Flexgen: High-throughput generative inference of large language models with a single gpu. In International Conference on Machine Learning. PMLR, 31094–31116.
- Song et al. (2023) Yixin Song, Zeyu Mi, Haotong Xie, and Haibo Chen. 2023. Powerinfer: Fast large language model serving with a consumer-grade gpu. arXiv preprint arXiv:2312.12456 (2023).
- Tang et al. (2024a) Zhenheng Tang, Xueze Kang, Yiming Yin, Xinglin Pan, Yuxin Wang, Xin He, Qiang Wang, Rongfei Zeng, Kaiyong Zhao, Shaohuai Shi, Amelie Chi Zhou, Bo Li, Bingsheng He, and Xiaowen Chu. 2024a. FusionLLM: A Decentralized LLM Training System on Geo-distributed GPUs with Adaptive Compression. arXiv:2410.12707 [cs.DC] https://arxiv.org/abs/2410.12707
- Tang et al. (2023) Zhenheng Tang, Yuxin Wang, Xin He, Longteng Zhang, Xinglin Pan, Qiang Wang, Rongfei Zeng, Kaiyong Zhao, Shaohuai Shi, Bingsheng He, et al. 2023. FusionAI: Decentralized Training and Deploying LLMs with Massive Consumer-Level GPUs. arXiv preprint arXiv:2309.01172 (2023).
- Tang et al. (2024b) Zhenheng Tang, Yonggang Zhang, Peijie Dong, Yiu ming Cheung, Amelie Chi Zhou, Bo Han, and Xiaowen Chu. 2024b. FuseFL: One-Shot Federated Learning through the Lens of Causality with Progressive Model Fusion. In The Thirty-eighth Annual Conference on Neural Information Processing Systems. https://openreview.net/forum?id=E7fZOoiEKl
- Wei et al. (2022) Jason Wei, Yi Tay, Rishi Bommasani, Colin Raffel, Barret Zoph, Sebastian Borgeaud, Dani Yogatama, Maarten Bosma, Denny Zhou, Donald Metzler, et al. 2022. Emergent abilities of large language models. arXiv preprint arXiv:2206.07682 (2022).
- Williams et al. (2009) Samuel Williams, Andrew Waterman, and David Patterson. 2009. Roofline: an insightful visual performance model for multicore architectures. Commun. ACM 52, 4 (2009), 65–76.
- Xu et al. (2024) Minrui Xu, Dusit Niyato, Hongliang Zhang, Jiawen Kang, Zehui Xiong, Shiwen Mao, and Zhu Han. 2024. Cached model-as-a-resource: Provisioning large language model agents for edge intelligence in space-air-ground integrated networks. arXiv preprint arXiv:2403.05826 (2024).
- Xue et al. (2024) Leyang Xue, Yao Fu, Zhan Lu, Luo Mai, and Mahesh Marina. 2024. Moe-infinity: Activation-aware expert offloading for efficient moe serving. arXiv preprint arXiv:2401.14361 (2024).
- Yuan et al. (2024) Xiaoming Yuan, Weixuan Kong, Zhenyu Luo, and Minrui Xu. 2024. Efficient Inference Offloading for Mixture-of-Experts Large Language Models in Internet of Medical Things. Electronics 13, 11 (2024), 2077.
- Zhang et al. (2022) Susan Zhang, Stephen Roller, Naman Goyal, Mikel Artetxe, Moya Chen, Shuohui Chen, Christopher Dewan, Mona Diab, Xian Li, Xi Victoria Lin, et al. 2022. Opt: Open pre-trained transformer language models. arXiv preprint arXiv:2205.01068 (2022).
- Zhong et al. (2024) Yinmin Zhong, Shengyu Liu, Junda Chen, Jianbo Hu, Yibo Zhu, Xuanzhe Liu, Xin Jin, and Hao Zhang. 2024. Distserve: Disaggregating prefill and decoding for goodput-optimized large language model serving. arXiv preprint arXiv:2401.09670 (2024).
- Zhou et al. (2022) Yanqi Zhou, Tao Lei, Hanxiao Liu, Nan Du, Yanping Huang, Vincent Zhao, Andrew M Dai, Quoc V Le, James Laudon, et al. 2022. Mixture-of-experts with expert choice routing. Advances in Neural Information Processing Systems 35 (2022), 7103–7114.
- Zhu et al. (2024) Qihao Zhu, Daya Guo, Zhihong Shao, Dejian Yang, Peiyi Wang, Runxin Xu, Y Wu, Yukun Li, Huazuo Gao, Shirong Ma, et al. 2024. DeepSeek-Coder-V2: Breaking the Barrier of Closed-Source Models in Code Intelligence. arXiv preprint arXiv:2406.11931 (2024).