💻 computer science

TEMPO: Makespan-Aware Expert-Parallel Load Balancing Across Memory- and Compute-Bound Regimes

TEMPO introduces a makespan-aware expert-parallel load balancing dispatcher that models non-linear expert execution times across memory- and compute-bound regimes to dynamically optimize token distribution, achieving up to 15.5% throughput gains and significant latency reductions in mixed-regime scenarios where traditional linear-count-based methods fail.

Original authors: Jie Li, Chenxin Jia, Jinliang Shen, Cunzhuang Liu, Ruiyi Ding, Jianwen Xian, Kang He, Chengru Song

Published 2026-08-14

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Jie Li, Chenxin Jia, Jinliang Shen, Cunzhuang Liu, Ruiyi Ding, Jianwen Xian, Kang He, Chengru Song

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 massive, high-speed pizza delivery service for a city that never sleeps. You have a fleet of identical delivery drivers (the GPUs) and a central kitchen with hundreds of different specialty chefs (the "experts" in an AI model). Every time a customer orders a pizza, the system has to decide which chefs work on it and which driver takes the finished pie. In the world of Artificial Intelligence, specifically with a type of model called a "Mixture-of-Experts" (MoE), this is exactly what happens. These models are like giant brains made of thousands of smaller, specialized sub-brains. When the AI thinks, it doesn't use its whole brain at once; it picks a few specific experts to handle the job.

The big challenge is keeping the whole team moving at the same speed. If one driver gets stuck with a huge, complicated order while everyone else is idle, the whole delivery is delayed. The time it takes to finish a batch of orders is determined by the slowest person in the group. For years, the standard rule for balancing this workload was simple: "Just split the number of orders evenly." If you have 100 orders, give 10 to each of your 10 drivers. This seemed logical, like sharing a pile of apples equally. But what if some apples are heavy rocks and others are light feathers? Or what if the kitchen has a rule where picking up a new chef takes a fixed amount of time, no matter how many pizzas they make? The old rules assumed that time was always directly tied to the number of orders. This paper asks: what if that assumption is wrong?

The researchers behind this paper, working at KlingAI, discovered that the old "count the orders" rule is actually a trap. They found that in modern AI hardware, the time it takes to process an expert isn't just about how many tokens (words or data chunks) it sees. It's a two-faced beast. Sometimes, the time is dominated by the sheer effort of loading the expert's "recipe" (weights) from the memory bank, which takes a fixed amount of time regardless of how small the order is. Other times, once the recipe is loaded, the time grows linearly with the number of orders. The old methods, which only looked at the number of orders, were missing the hidden cost of loading the recipe. They were trying to balance a pile of feathers and rocks by counting them, not weighing them.

To fix this, the team built a new dispatcher called TEMPO (Time-modeled Expert-Parallel Optimization). Instead of just counting tokens, TEMPO acts like a smart traffic controller that understands the physics of the kitchen. It uses a special "cost model" that measures exactly how long it takes to load a chef's recipe and how long it takes to cook the pizza. It realizes that if you have a "cold" expert (one that hasn't been used in a while), splitting its tiny order across two drivers is a disaster because you have to pay the "loading fee" twice. But if you have a "hot" expert with a mountain of orders, splitting it is fine.

The paper shows that TEMPO doesn't just guess; it calculates the perfect balance for every single batch of requests in milliseconds. They tested this on real AI models and found that the old methods were often 15% slower or caused significant delays for the last few customers in a line. TEMPO, however, keeps the line moving smoothly. It's like switching from a rule that says "everyone gets the same number of apples" to a rule that says "everyone gets the same amount of work," taking into account that some apples are heavy and some chefs are slow to wake up.

The researchers were very careful to show exactly where this new method works and where it doesn't. They proved that if the "hot" experts are so numerous that the system is just overwhelmed by the sheer volume of data (the "compute-bound" regime), the old token-counting method is actually fine. But in the real world, where some experts are busy and others are resting, and where the "loading fee" is high, TEMPO shines. They even mapped out a "phase diagram," which is like a weather map for AI traffic, predicting exactly when the new method will save time and when the old method is good enough.

In the end, this paper isn't just about a faster algorithm; it's about changing how we think about balancing work in AI. It teaches us that in the complex, high-speed world of modern AI, you can't just count things. You have to understand the hidden costs of moving data and the specific shape of the work. By measuring the actual time it takes to do the job rather than just counting the items, TEMPO makes AI models faster, more efficient, and ready to handle the massive demands of the future. It turns a chaotic kitchen into a well-oiled machine, ensuring that no driver is left waiting while the pizza sits under the heat lamp.

Technical Summary: TEMPO – Makespan-Aware Expert-Parallel Load Balancing

1. Problem Definition

In Mixture-of-Experts (MoE) serving, the latency of a layer is determined by the makespan (the maximum time spent on any single GPU) due to synchronization points at the end of every layer. Current production dispatchers (e.g., EPLB, LPLB, UltraEP, METRO) attempt to balance load by optimizing proxies: either the number of tokens per GPU or the number of activated experts per GPU. These proxies rely on implicit linearity assumptions—that expert execution time is linearly proportional to token count or expert count, respectively.

However, measurements on datacenter GPUs (specifically using DeepGEMM fp8 grouped GEMM kernels) reveal that expert execution time is non-linear and follows a two-regime structure:

Memory-Bound (Flat) Regime: Below a threshold of $n^* \approx 156\text{--}168$ tokens per expert, the cost is dominated by streaming expert weights from High Bandwidth Memory (HBM). In this regime, time is constant per activated expert replica, regardless of token count.
Compute-Bound (Linear) Regime: Above $n^*$ , the cost becomes linear with token count. However, due to grouped GEMM implementations that round tokens up to 128-token M-tiles, splitting an expert across multiple replicas introduces "padded compute," making the cost non-linear in a way that simple token counting misses.

Consequently, real-world decode batches often contain a mixed regime: "hot" experts with many tokens (linear regime) and "cold" experts with few tokens (flat regime) coexist. Existing proxies fail because token-balancing strategies fragment cold experts (paying hidden activation floors), while activation-balancing strategies overload hot replicas. No single fixed proxy is optimal across all operating points.

2. Methodology

Cost Modeling

The authors formalize the per-expert execution time $t$ as a max-affine function:
$t = \max(a + bG, c + \beta N)$
Where $G$ is the number of activated expert replicas and $N$ is the token count.

$b$ (Activation Floor): The cost of streaming weights for an additional replica (memory-bound).
$\beta$ (Token Slope): The cost per token in the compute-bound regime.
Tile Awareness: For prefill or large batches, a secondary parameter captures the "staircase" effect of 128-token tiling, refining the model to $t = \max(a+bG+b_2(T-G), c+\beta N)$ .

This model is calibrated via a black-box, ten-minute microbenchmark on the deployed kernel, measuring wall-clock times across $(G, N)$ grids.

Problem Formulation

The dispatch problem is formalized as a fixed-charge makespan minimization problem. The authors prove this problem is NP-hard even with two GPUs and full replication, as it requires balancing two competing constraints (activation count vs. token count) simultaneously. However, the problem is polynomial in degenerate limits (pure token balancing or pure activation balancing).

The TEMPO Algorithm

TEMPO (Time-modeled Expert-Parallel Optimization) is a dispatcher that solves this makespan problem in milliseconds, operating off the critical path. Its solver, tempo fast, employs four stages:

Cost-Aware Greedy Seeding: Places whole experts on replicas with the lowest marginal cost, avoiding the fragmentation of cold experts.
Activation Rebalancing: Uses augmenting chains to rebalance expert counts in the flat regime (where token moves are ineffective).
Bottleneck Local Search: Performs partial token migrations to reduce the makespan in the linear regime.
Ensemble with Switching Tolerance: Scores the current solution against a Token-LP (Linear Programming) baseline and a Round-Robin heuristic. It only switches if the new solution is >1% better, ensuring stability and inheriting the theoretical guarantees of the simpler algorithms.

System Integration

TEMPO is integrated into SGLang with a focus on zero overhead on the critical path:

In-Graph: A single fused CUDA kernel handles probabilistic dispatch and count collection, eliminating per-layer collectives required by previous LP solvers.
Out-of-Process: The solver runs in a separate process, avoiding GIL contention and ensuring the dispatch decision does not block the inference graph.
Topology Awareness: For multi-node setups, a two-stage split ensures tokens are routed to replicas on the same node first, minimizing inter-node traffic without altering the core solver.

3. Key Results

Theoretical and Simulation

Phase Diagram: A calibrated phase diagram demonstrates that the optimal fixed policy flips between activation balancing (memory-bound) and token balancing (compute-bound) based on batch size and expert skew. TEMPO tracks the optimal policy across this entire map.
Approximation Guarantee: Under full replication, TEMPO provides an additive approximation guarantee of $OPT + \max(b, \beta n_{max})$ .
Simulation Gains: In simulated environments, TEMPO achieves up to 15.5% improvement over the best fixed baseline in mixed-regime scenarios (specifically for EP32–64 extrapolations).

Empirical Evaluation (Testbed A & B)

Microbenchmarks: On an 8-GPU Testbed A, TEMPO stays within 1% of the best fixed baseline at every evaluated grid point and wins by up to 15.5% where regimes mix.
End-to-End Serving (Qwen3-235B):
- On long-context and decode-heavy traffic, TEMPO improves throughput by 4–6% and reduces p99 inter-token latency by ~15.6%.
- On a 2-node EP16 setup with topology-aware splitting, gains reach 4–7%.
- Crucially, the paper notes that much of the gap between TEMPO and SGLang's shipped LPLB is architectural (removing in-graph collectives); when LPLB is ported to the same architecture, the residual gain of TEMPO is primarily stability at compute-bound points.
End-to-End Serving (DeepSeek-V3): On a model where communication dominates and experts are smaller per GPU (outside the "win region"), TEMPO returns only mechanism costs, performing on par with static placement. This confirms the paper's claim that the benefits are conditional.
Staleness and Placement: The system is robust to placement staleness (up to 200ms drift). However, if the placement layer is significantly stale (hot experts moved to cold replicas), dispatch alone cannot recover throughput; only the placement layer can fix the mean, while dispatch repairs the latency tail.

4. Significance and Claims

The paper's primary claim is not a universal performance win, but a phase diagram that predicts when time-based dispatch is beneficial.

The "Win Region": Adaptive dispatch yields significant gains only when:
1. The model has large experts (high activation floor $b$ ).
2. Batches are in a mixed regime (containing both hot and cold experts).
3. Placement is stale enough to cause imbalance but not so stale that only weight migration can fix it.
4. Expert compute constitutes a significant share of the step time (not dominated by communication).
Methodological Shift: The paper argues that "balancing time, not tokens" is the principled approach for MoE serving. It demonstrates that token-count proxies are systematically wrong in memory-bound regimes and that a calibrated cost model is necessary to navigate the interaction between memory and compute bottlenecks.
Practicality: The work provides a deployable solution (TEMPO) that integrates seamlessly into existing serving stacks (SGLang) with negligible overhead, proving that complex makespan optimization can be performed in milliseconds without blocking inference.

The authors explicitly state that for small-expert shapes or communication-dominated models, adaptive dispatch offers no advantage over static placement, and the system correctly degrades to parity rather than incurring a penalty. The contribution is the mapping of these conditions and the algorithm to navigate them.