💬 NLP

GEM: GPU-Variability-Aware Expert to GPU Mapping for MoE Systems

GEM is a framework that optimizes Mixture-of-Experts (MoE) model inference latency by mapping experts to GPUs based on hardware variability and token load patterns, thereby mitigating straggler effects and improving end-to-end performance by up to 16.5%.

Original authors: Sourish Wawdhane, Avinash Kumar, Poulami Das

Published 2026-05-20

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Sourish Wawdhane, Avinash Kumar, Poulami Das

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are running a high-end restaurant kitchen (the GPU cluster) where you have a team of specialized chefs (the Experts) preparing a massive order of dishes for many customers at once.

In a modern "Mixture-of-Experts" (MoE) AI model, the kitchen doesn't have one giant chef doing everything. Instead, it has many smaller, specialized chefs. For every single word the AI generates (a "token"), a manager (the Router) decides which two chefs are needed to cook that specific word.

The Problem: The "Slowest Chef" Rule

In this kitchen, there is a strict rule: The whole team must wait for the slowest person to finish before they can move to the next step.

If you have 8 chefs, and 7 of them finish their tasks in 10 seconds, but one chef takes 12 seconds because they are slower or have too much work, the entire kitchen is stuck waiting for those extra 2 seconds. In the world of AI, this waiting time is called a "straggler," and it kills the speed of the whole system.

The paper identifies two main reasons why a chef becomes a straggler:

Bad Assignment: The manager accidentally gave the busiest, most popular recipes to just one chef, while others sat idle.
Hardware Variability: Even if the work is perfectly split, some chefs are just naturally slower than others due to their specific hardware (like an older oven or a tired arm). The paper found that in a group of identical-looking GPUs, the fastest one can be nearly 28% faster than the slowest one.

The Old Way: "Equal Work, Equal Time"

Previous solutions tried to fix this by giving every chef the exact same number of dishes to cook. They thought, "If everyone has the same amount of work, everyone will finish at the same time."

But this fails because:

Different Speeds: A fast chef can cook 14% more dishes in the same amount of time as a slow chef. If you give them the exact same pile of work, the fast chef finishes early and waits, while the slow chef is still struggling.
Hidden Patterns: Some chefs are busy almost all the time (Consistent Experts), while others are only busy together in short, intense bursts (Temporal Experts). Old methods missed these "bursty" patterns. If two chefs who always get busy at the exact same time are assigned to the same slow machine, the whole kitchen grinds to a halt.

The New Solution: GEM (GPU-variability-aware Expert Mapping)

The authors propose GEM, a smart system that acts like a genius kitchen manager who knows exactly how fast each chef is and how the orders come in.

GEM uses two clever tricks:

1. The "Proportional Load" Strategy
Instead of giving everyone the same number of dishes, GEM gives the fast chefs more dishes and the slow chefs fewer dishes.

Analogy: Imagine a race. If one runner is 14% faster, you don't give them the same distance as the slower runner. You give them a longer track so they both cross the finish line at the exact same moment.
GEM calculates exactly how much extra work the fast GPUs can handle so that everyone finishes the layer at the same time.

2. The "Pattern Detective" Strategy
GEM watches the kitchen for a short while (just 16 steps) to learn two things:

Who is always busy? (The Consistent Experts).
Who gets busy together? (The Temporal Experts).
Analogy: If Chef A and Chef B always get a huge rush of orders at the same time, GEM makes sure they are not assigned to the same slow oven. It spreads them out to different stations so they don't bottleneck each other.

How GEM Works (The 4-Step Process)

Watch and Learn: GEM watches the AI for a tiny moment to see which experts get used and when.
Test the Hardware: It runs a quick test to see exactly how fast each specific GPU is under different loads. It does this smartly by testing only at specific "milestones" (like checking a car's speed every 32 miles instead of every mile) to save time.
The Search: It runs a simulation to find the perfect arrangement of chefs to ovens. It tries swapping chefs around until it finds a setup where the "slowest" chef finishes as fast as possible.
Deploy: It locks in this new arrangement. The AI starts running, and because the work is balanced for the actual speed of the machines, the whole system runs smoother.

The Results

When the authors tested this on five different powerful AI models:

Speed Boost: The AI finished tasks 7.9% faster on average.
Best Case: In some situations, it was 16.5% faster.
Smoother Experience: The "tail latency" (the worst-case delays that make users feel like the AI is stuttering) improved even more, by up to 16.9%.

In short, GEM stops the AI from waiting for the slowest part of the system by giving the fast parts more work and the slow parts less, while making sure no two "busy" experts are stuck on the same slow machine. It turns hardware differences from a weakness into a tool for better performance.

Technical Summary: GEM – GPU-Variability-Aware Expert-to-GPU Mapping for Mixture-of-Experts Models

1. Problem Statement

Mixture-of-Experts (MoE) models improve inference efficiency by activating only a subset of "expert" sub-networks per token. However, their performance in real-world serving environments is severely bottlenecked by straggler GPUs. In lock-step execution, a batch must wait for the slowest GPU to finish processing a layer before proceeding.

The paper identifies two primary causes for stragglers that existing solutions fail to address simultaneously:

Expert Imbalance: Certain experts are used disproportionately more often than others. If these "heavily used" experts are co-located on a single GPU, that GPU becomes a bottleneck.
Hardware Variability: Modern GPUs exhibit inherent performance variability due to process variations, environmental factors, and power management. Even with perfectly balanced token loads, a physically slower GPU will act as a straggler.

Prior works attempt to solve this by either introducing artificial uniformity in expert utilization (degrading model accuracy) or dynamically redistributing experts (incurring memory overheads). Crucially, these approaches are agnostic of hardware variability, often placing heavily used experts on the slowest available GPUs, thereby perpetuating the straggler problem.

2. Methodology: GEM Framework

The authors propose GEM (GPU-variability-aware Expert Mapping), a framework designed to map experts to GPUs by explicitly accounting for both task-specific token loads and hardware performance characteristics. GEM operates via a four-step pipeline:

Step 1: Expert Utilization Trace

GEM captures the token routing patterns of a specific task during online inference. It counts the number of tokens routed to each expert to build a utilization trace.

Insight: The paper distinguishes between consistent experts (used frequently across most time steps) and temporal experts (used infrequently but heavily in specific, correlated bursts).
Optimization: Experiments show that a trace length of 16 time steps is sufficient to capture both consistent and temporal patterns without delaying deployment.

Step 2: Performance Variability Profiling

GEM profiles the latency of MoE layers across different GPUs as a function of token load.

Challenge: Fine-grained profiling (every token count) is too slow (hours), while coarse sampling misses the "staircase" nature of MoE latency (where latency jumps at specific batch sizes).
Solution: GEM samples token counts only at tile boundaries (multiples of 32 or 64 tokens, where latency jumps occur). This reduces profiling time from hours to minutes (e.g., 3 minutes for Llama-4-Scout) while maintaining accuracy.

Step 3: Variability-Aware Expert Mapping

GEM uses an iterative heuristic search to find the optimal expert-to-GPU mapping.

Scoring Function: Instead of balancing token counts, GEM minimizes the sum of straggler latencies across all time steps. The score $S(M)$ is calculated as:
$S(M) = \sum_{t \in T} \max_{g} (C_g(n_g(M, t)))$
Where $C_g$ is the latency cost for GPU $g$ given token count $n_g$ , derived from the variability profiles.
Key Insights Applied:
1. Non-Uniform Load Balancing: Faster GPUs are assigned proportionally more tokens than slower GPUs so that all GPUs finish processing a layer at approximately the same time.
2. Correlation Handling: GEM detects correlated temporal experts (experts used simultaneously in bursts) and ensures they are placed on different GPUs. Furthermore, it avoids placing these high-load experts on the slowest GPUs.
Search Algorithm: The system starts with multiple initial mappings (adding noise to utilization metrics to ensure diversity) and iteratively swaps expert pairs to reduce the total straggler latency score until convergence.

Step 4: Deployment

Once the optimal mapping is determined, GEM loads the expert weights onto their assigned GPUs at model load time. The MoE router continues to dispatch tokens as usual, but the physical location of the experts is now optimized for the specific hardware cluster.

3. Key Contributions

Identification of Dual Bottlenecks: The paper highlights that stragglers arise from the combination of non-uniform expert utilization and inherent GPU hardware variability, a problem overlooked by prior static or dynamic balancing methods.
GEM Framework: The first framework to perform hardware-variability-aware expert mapping for MoE models.
Novel Mapping Strategy:
- Latency-Based Balancing: Moving away from equal token distribution to proportional token distribution based on GPU speed.
- Temporal Expert Detection: Identifying and separating correlated temporal experts to prevent simultaneous load spikes on single GPUs.
Efficient Profiling: A tile-boundary sampling strategy that captures hardware variability profiles in minutes rather than hours.

4. Experimental Results

The authors evaluated GEM on five state-of-the-art MoE models (Mixtral-8x7B, Mixtral-8x22B, Llama-4-Scout, Hunyuan-A13B, Qwen3-30B-A3B) using ShareGPT and CodeContests datasets. They emulated three levels of GPU variability (Low, Moderate, High) on a 4-GPU node.

End-to-End Latency: GEM improved latency by an average of 7.9% and up to 16.5% compared to the baseline (vLLM's default linear mapping).
- High variability setups saw the most significant gains (up to 16.5%).
- Even in low variability setups, GEM provided improvements (up to 4.5%) by optimizing for correlated temporal experts.
Tail Latency (p90 TPOT): GEM reduced the 90th-percentile time-per-output-token by an average of 9.1% (up to 16.9%), significantly improving perceived responsiveness.
Profiling Overhead: The variability profiling and mapping process takes less than 4 minutes total, which is negligible compared to the hours or days of inference served in production.

5. Significance and Claims

The paper claims that GEM is the first hardware variability-aware expert mapping framework for MoE inference. Its significance lies in transforming hardware variability from a liability into a mapping constraint that can be exploited to improve performance.

The authors emphasize that:

Simply balancing token counts is insufficient because it ignores hardware heterogeneity.
Ignoring temporal experts (which are often correlated) leads to suboptimal mappings even if average loads are balanced.
As system scales increase (e.g., from 4 to 64 GPUs), the gap between the fastest and slowest GPUs widens monotonically, making variability-aware mapping increasingly critical for large-scale deployments.

The work demonstrates that by aligning expert placement with both the statistical nature of token routing and the physical reality of GPU performance, significant latency reductions can be achieved without retraining models or incurring substantial memory overheads.