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

Summary

This paper introduces TEMPO (Time-modeled Expert-Parallel Optimization), a makespan-aware load-balancing dispatcher for Mixture-of-Experts (MoE) model serving. The core argument is that existing production dispatchers balance proxies—token counts or activated-expert counts—which implicitly assume expert computation time is linear in one of those two quantities. The paper demonstrates through measurements on two generations of datacenter GPUs that neither assumption holds, and proposes a calibrated cost model plus a solver that optimizes the actual per-batch makespan (the time of the slowest GPU in the expert-parallel group).

Measured hardware behavior. The paper benchmarks DeepGEMM fp8 masked grouped GEMM kernels and finds a two-regime cost structure per expert: flat in token count up to an inflection n∗ ≈ 156–168 tokens per expert (loading expert weights dominates) and linear beyond it. The per-GPU cost is modeled as a max-affine function: tg = max(a + b·Gg, c + β·Ng), where Gg is the number of activated (expert, replica) pairs on GPU g and Ng is its token count. The activation floor b is physically the cost of streaming an expert's weights from HBM—a roofline check shows fitted b values (1.74 µs for Qwen3-30B shape, 14.8 µs for DeepSeek-V3 shape) sit at 1.13–1.24× the theoretical weight-streaming bound. Above the inflection, compute is also not linear in tokens because grouped GEMM pads every expert's tokens to 128-token M-tiles, so splitting an expert across replicas manufactures padded compute. A refined tile-aware model adds a one-parameter staircase term: t = max(a + bG + b2(T−G), c + βN), with T = Σe ⌈ne/128⌉ and b2 ≈ b/3.

Proxy failures. On recorded batches with the same placement, the dispatches produced by four proxy policies (uniform split, token-LP, activation balancing/METRO, static) differ by 1.4–1.6× in modeled block time at the median (p95 up to 1.7×), and the identity of the expensive proxy flips with the regime. The pivotal empirical fact is that "in real routing traces at decode batch sizes, 92–100% of batches contain both regimes at once—the few hot experts sit deep in the linear region carrying ∼91% of tokens while roughly half of the activated experts remain in the flat region." Token balancing fragments cold experts and pays hidden activation floors; activation balancing piles tokens onto hot replicas.

Problem formulation. The paper formalizes per-batch dispatch as a fixed-charge makespan problem: minimize max g max(a + b·Σ e z e,g, c + β·Σ e x e,g) subject to token shares and activation indicators. Theorem 1 proves this is NP-hard even with 2 GPUs, full replication, and a=c=0 (reduction from Balanced PARTITION). However, the problem is polynomial in each degenerate limit: Lemma 1 shows b→0 reduces to the token LP; Lemma 2 shows β→0 reduces to an optimal semi-matching. Theorem 2 provides an additive approximation guarantee under full replication: a round-robin whole-expert placement A3 achieves M(A3) ≤ OPT + max(b, β·nmax), with the additive term equal to β·nmax in all recorded traces. The paper notes the hardness, like the systems problem, lives in the regime interaction.

Phase diagram. Sweeping synthetic Zipf routing over batch size, skew, replication, and EP size yields a phase diagram where the best fixed policy flips across the map: activation balancing wins at small B (memory-bound), token-LP at large B (compute-bound), and the boundary moves with s, replication, and shape. The flip boundary is analytically predictable via two mechanisms: an average mechanism (B* avg = n*·E eff/(K·n gpus)) and a hot-expert mechanism (B* hot = n*·E eff/(p1·K·n2 gpus)), with B* = min of the two. The analytic boundary lands inside the observed flip band in 12/12 grid columns. TEMPO tracks the per-cell best everywhere (min gain ≥ −0.3%) and wins by up to 8.5% (DSv3) / 10.2% (Qwen3) / 11.8% (DSv2, 1.5× replication) in the mixed zone. Scale extrapolation to EP32–64 shows mixed-zone gains persist (mean 6–6.5% at s=0, max 15.5%).

Algorithm (tempo fast). Four stages: (1) cost-aware greedy seeding (whole-expert placement, descending token count); (2) activation rebalancing via 1- and 2-step augmenting chains (truncated semi-matching); (3) bottleneck local search with partial migrations (ternary-search splits); (4) ensemble with a 1% switching tolerance that scores token-LP and the round-robin certificate A3 under the same model. Against a 10-second MILP, tempo fast achieves makespan ratio mean 1.005 / p95 1.024 / max 1.033, with solve time 1.9 ms at B≤256 (pure Python). Component ablation shows each stage binds in a different phase region: removing partial migrations degrades up to 16.3% (EP32 transition); removing augmenting chains up to 2.7% (mid-B flat); removing the ensemble up to 4.5% (EP64 deep-compute).

Communication and topology extensions. The deployed cost adds a third max-affine piece for all-to-all collectives: tg = max(a + b·Gg, c + β·Ng, c2 + γ·Ng). Term ablation shows dropping the traffic term forfeits more than half the win at B=512–1024; dropping the floor term is the mirror image at B=128. For multi-node EP, the paper adds a two-stage topology-aware split: solve (3) for per-GPU shares, then split token sources across replicas by a same-node-first transportation rule, which is optimal for inter-node traffic among all splits realizing the solved shares. This preserves per-GPU loads provably while minimizing inter-node traffic.

SGLang integration. The design goal is zero marginal in-graph cost. One fused kernel performs probabilistic dispatch plus count collection (atomic increments on a persistent cumulative counter, masked against CUDA-graph padding rows). The solver runs out-of-process in a separate numpy-only worker process, publishing tables race-safely with a fallback to uniform over replicas on torn reads. The paper reports Zero kernels and zero collectives on the critical path, verified by a no-op variant that ties static within noise.

Wall-clock validation. On an 8-GPU Testbed A EP8 microbenchmark with DSv3 expert shape: at B=32 (memory-bound), TEMPO beats EPLB-even by 11–14% and token-LP by 7%; at B=2048 (compute-bound), METRO is 7–11% off best while TEMPO ties token-LP at the top. Across all B, TEMPO is within 5% of the per-B best fixed policy while every fixed policy has a ≥7% failure region. The simulator transfer is validated: pairwise ranking agreement 93% (GEMM pipeline) and 92% (full pipeline), with mean gain-transfer error 5.5 pp and 2.2 pp respectively. The full-pipeline fit has an identifiability issue (a and b nearly collinear over the observed G range), but the solver consumes tightly identified quantities: β (±0.5%) and the flat cost a+bG (±2.8% at G=13).

End-to-end serving results. On Testbed B with Qwen3-235B-FP8 (inside the predicted win region): GovReport throughput +5.0% median with non-overlapping ranges; under moderate Poisson load (8 req/s) p99 TPOT −15.6% (191 vs. 226 ms) with median TTFT −12.5%. ShareGPT is parity (−0.2% throughput, −8.9% p99). On DeepSeek-V3-0324 (outside the win region, communication-dominated), every workload lands at −2 to −3%, indistinguishable from the noop control—at 32 experts per GPU there is no exploitable imbalance left. The pair of models turns the phase diagram from a simulation artifact into a falsifiable, and here twice-confirmed, deployment rule.

Architecture vs. objective separation. Against SGLang's shipped LPLB dispatcher, TEMPO delivers +38–70% request throughput on decode workloads and +65% on ShareGPT. However, a like-for-like port of token-LP into TEMPO's worker shows most of that gap is architectural, the time model's residual edge being stability at the compute-bound point. At 1024 requests, token-LP swings −11.6% to +8.7% across windows while TEMPO stays within noise (spread 4.4 vs. 20 pp).

Staleness and placement interaction. The theory optimizes the current batch; the deployment solves on the previous window's counts. Replay shows one window of staleness returns 5–7 pp of the 22 pp exact-solve headroom and keeps the rest. Under drift-16 (corrupted placement), TEMPO claws back a third of the tail damage (−10% p99, −5% TTFT) but none of the throughput—dispatch repairs the tail; only placement repairs the mean. With tight replica budgets (1.06×), blind uniform splitting backfires (−5.4%) while the cost model recovers parity (+0.3%).

Multi-node EP16. On 2-node Testbed B with Qwen3-235B, the flat table loses −3.5% at the all-to-all-bound point (1024 requests), while the topology-aware split recovers +4.1% (token-LP incidentally +4.0%). The four-way attribution shows the no-solve control captures part of the gain (+5.4% on decode-heavy), the flat solve matches it, and the same-node-first split adds the rest. On DeepSeek-V3 EP16, no balancer helps—when the expert-FFN stage is a minor fraction of step time, balancing it moves nothing—demonstrating that the win region needs three coordinates: moderate experts-per-GPU, sufficient skew, and expert compute a large enough share of the step.

Limitations. The paper is explicit: small-expert shapes (b ≈ 1.7 µs) show no adaptive-dispatch gain; with fresh placement and ample replication all policies tie; the (G, N) model has a floor (DeepGEMM masked kernels prefer uniform per-slot loads, a third cost dimension that decides near-ties); the deployed system is not the theoretical optimum by design (NP-hardness, asynchronous solving, whole-slot granularity, inherited placement). The paper concludes: Balance time, not tokens... dispatching on measured time rather than counted tokens is both the principled and the profitable choice.

Based on this paper, here are the specific improvements I can make to AI systems, and what the improved systems can do:

Improvement 1: Replace proxy-based load balancing with makespan-aware dispatch

• What I improve: Instead of balancing token counts or activated-expert counts (which assume linear compute scaling), I optimize the actual per-batch makespan—the time of the slowest GPU in the expert-parallel group.

• What the improved system can do: It dispatches tokens to experts using a calibrated cost model with two regimes (flat weight-loading cost + linear compute cost), plus a tile-aware staircase term for padded grouped-GEMM compute. This yields 5–12% throughput gains in mixed-regime batches (where hot experts carry 91% of tokens while half of experts sit in the flat region), and eliminates the 1.4–1.6× median gap between proxy policies.

Improvement 2: Use a hybrid solver that combines greedy seeding, augmenting chains, bottleneck local search, and ensemble scoring

Improvement 3: Add communication-aware cost modeling for multi-node deployments

Improvement 4: Integrate dispatch with zero marginal in-graph cost

Improvement 5: Make dispatch robust to staleness and placement drift

Improvement 6: Use the phase diagram as a deployment rule to decide when to adapt

Improvement 7: Separate architecture from objective in dispatcher design

Improvement 8: Add a tile-aware cost term for grouped GEMM padding

Improvement 9: Use roofline-validated cost parameters

Improvement 10: Provide a fallback certificate with additive approximation guarantee

In expert-parallel (EP) MoE serving, every layer synchronizes at the slowest GPU. Dispatchers balance token counts (EPLB, LPLB, UltraEP) or activated-expert counts (METRO), assuming expert time is linear in one. Measurements on two datacenter GPU generations show it is neither: below ! about!156 -- 168 tokens, HBM weight streaming dominates---cost attaches to activated replicas, not tokens; above it, grouped GEMM rounds tokens to 128-tile M-tiles, so splitting an expert adds padded compute. A max-affine profile t= (a+bG,,c+ beta N) captures both regimes. Realistic decode batches hold hot experts in the linear regime and cold in the flat simultaneously; recorded batches show proxy dispatches differ by 1.4 -- 1.6 times in modeled block time (p95 up to 1.7 times), and which proxy wins flips with the regime. We formalize per-batch dispatch as a fixed-charge makespan problem---NP-hard on two fully replicated GPUs, polynomial in degenerate limits---and present, a makespan-aware dispatcher solving it in milliseconds off the critical path; its SGLang integration runs out-of-process and fuses dispatch with count collection into one in-graph kernel. Anchored by an 8-GPU Testbed A microbenchmark, stays within 1% of the best fixed baseline everywhere and wins by up to 15.5% where regimes mix. End-to-end on Testbed B, Qwen3-235B (inside the win region) gains 4 -- 6% throughput and cuts p99 latency by about 15.6%; DeepSeek-V3 (outside, communication-dominated) shows only mechanism cost. A phase diagram, not a universal win, is the claim: it predicts both outcomes before deployment.