Source-linked AI summary
FreeBalance: Pre-Routing Online Moe Load Balancing via Residual Workload Prediction
Pengfei Chen, Yize Wu, Shouxu Kuang, Ke Gao, Ling Li
TL;DR
Dynamic routing can imbalance distributed MoE inference, so FreeBalance predicts upcoming workloads before routing and overlaps expert migration with preceding computation, reducing rank-load imbalance by up to 32.8% and prefill latency by 13.1%.
Problem
Uneven inference-time routing leaves the most heavily loaded expert-parallel rank as a straggler, limiting utilization and increasing distributed MoE latency.
Method
FreeBalance predicts upcoming expert workloads from residual representations, then proactively plans budgeted expert migration before routing while preserving the authoritative routing decisions.
Results
FreeBalance reduces the max-to-mean rank-load ratio by up to 32.8% and average prefill latency by 13.1% across 19 LongBench subsets.
Takeaways & Limitations
Pre-routing balancing can adapt to dynamic workloads while hiding reconfiguration overhead from the critical path and preserving original routing decisions and outputs.
Takeaways & Limitations
The approach assumes early workload estimates are used only for placement, while the normal router remains authoritative for routing semantics and baseline execution.
Abstract
from arXiv · showhide
Load imbalance poses a major bottleneck to the efficiency of expert parallelism in distributed inference of Mixture-of-Experts (MoE) models. The most heavily loaded rank stalls global execution due to skewed routing distributions, directly increasing latency. While offline expert placement can alleviate persistent imbalance, practical multi-task serving workloads exhibit layer- and batch-dependent routing dynamics, making online load balancing indispensable. Existing approaches rely on routing statistics collected after each MoE router, requiring expert weight load or migration to begin only after routing decisions are available, consequently placing migration overhead on the inference critical path. In this work, we observe that online balancing can instead be largely overlapped with computation before target routing (e.g., attention), if routing distributions can be predicted accurately in advance. Therefore, we propose FreeBalance, a lossless online load-balancing framework that overlaps expert migration with preceding computation stages via residual workload prediction. FreeBalance leverages cross-layer similarities in hidden representations within the residual network to build a lightweight workload predictor. This enables proactive expert migration planning before routing decisions are available, creating substantial overlap between weight transfer and computation-heavy pre-routing stages. Furthermore, a cost model constrains the number of swaps to fully hide the synchronization overhead within the available window. Experiments across models and datasets show that FreeBalance reduces the max-to-mean rank load ratio by 32.8% and end-to-end prefill latency by 13.1%. Specifically, our method hides balancing overhead of an average of 5.1 experts per layer, which would otherwise account for about 8.5% of the critical-path latency.
1 Introduction
FreeBalance addresses dynamic load imbalance in distributed MoE inference by predicting upcoming expert workloads before routing and overlapping expert migration with preceding computation. This pre-routing approach aims to hide balancing overhead while preserving lossless MoE computation.
- Background: MoE layers activate only a small subset of experts per token, increasing parameter capacity without proportionally increasing per-token computation.The routing mechanism decouples model capacity from inference latency.
- Background: Expert parallelism shards experts across devices, using all-to-all communication to dispatch activations and combine expert outputs.Distributed execution requires synchronization across the expert-parallel group.
- Problem: Load imbalance from skewed routing distributions makes the most heavily loaded rank determine synchronized execution latency.Inference-time routing can remain highly uneven despite training-time auxiliary load-balancing loss.
- Problem: Static expert-device mappings cannot promptly adapt to routing distributions that vary across multi-task serving requests.Offline mappings rely on historical routing statistics, whereas practical workloads vary by task type.
- FreeBalance: FreeBalance predicts an upcoming layer’s expert workload from the preceding layer’s residual representation before routing, enabling expert migration to overlap with preceding computation.The prediction is lightweight relative to inference latency and supports pre-routing balancing.
2 Preliminary
MoE routing selects top-k experts for each token, while expert parallelism dispatches tokens across devices through all-to-all communication. Skewed, task-dependent routing creates rank imbalance, and conventional online balancing incurs migration overhead on the critical path.
- MoE routing: For each token representation, the MoE router scores experts and selects the top-k experts to process the token.Only selected experts process the token, and their outputs are weighted and aggregated.
- Expert parallelism: Under expert parallelism, tokens routed to remote experts require an all-to-all dispatch, followed by a second all-to-all returning expert outputs.Experts are assigned to one or more devices through an expert–device mapping.
- Load imbalance: Skewed routing makes the most heavily loaded rank determine the critical path, leaving less-loaded ranks idle during global synchronization.The imbalance arises when hot experts are colocated on the same rank and expert-level skew becomes inter-rank workload imbalance.
- Workload variation: Only 4.1% of expert activations overlap across all layers for two representative Qwen3-30B tasks, demonstrating substantial task-dependent routing variation.This variation can make fixed offline expert placement suboptimal for multi-task serving.
- Online balancing challenge: Online migration decisions made after routing leave only a narrow interval before MoE computation, placing migration overhead directly on the critical inference path.The added overhead can largely diminish the benefit of improved load balance.
3 Method
FreeBalance predicts target-layer expert workloads before routing, using cross-layer residual similarity and the frozen router to plan proactive migration during preceding attention computation. A budgeted, deterministic swap planner enables balanced placement while preserving the target router’s original decisions and lossless execution.
- Lossless MoE computation: FreeBalance preserves lossless inference because the normal target router still determines every token-to-expert assignment after migration.Predicted workloads affect only physical expert placement and weight transfer; an empty swap plan retains the current placement.
- Budgeted expert-swap planning: A deterministic greedy planner selects beneficial pairwise expert swaps under a migration budget that limits communication to what target-layer attention can hide.The planner lowers predicted maximum rank load, discards non-positive-benefit candidates, and produces the same result on every rank.
- Weight transfer: Expert weights migrate during computation-heavy pre-routing stages, mainly attention, and placement is committed before the target MoE stage begins.Attention latency is profiled once per attention type and reused for subsequent layers with the same mechanism.
- Residual workload prediction: Residual hidden-state similarity lets FreeBalance estimate target-layer workloads from h_ℓ−1 before target-layer attention and routing.Because residual updates are small and adjacent layers process the same tokens, h_ℓ−1 retains task- and token-specific information while creating an attention-length migration window.
- Residual workload prediction: The same frozen router is invoked early on h_ℓ−1 for workload estimation and normally on H_ℓ for authoritative routing, adding no predictor weights or training state.The early invocation produces only an E-element workload vector for placement planning; the normal invocation remains unchanged.
4 Evaluation
The evaluation measures FreeBalance on two MoE models using LongBench and a dynamic Mixed Tasks workload. Experiments compare fixed placement, FreeBalance, history-driven EPLB, and their combination using latency, load balance, prediction, and planner-quality metrics.
- Models and workloads: Experiments evaluate Qwen3-30B-A3B-Instruct-2507 and Moonlight-16B-A3B-Instruct with different expert counts, top-k routing, experts per rank, and expert sizes.Qwen3 uses 128 experts with top-8 routing and 16 experts per rank; Moonlight uses 64 experts with top-6 routing and 8 experts per rank.
- Models and workloads: LongBench provides 19 independently evaluated subsets spanning six long-context capabilities under shared batching and sequence-length settings.The subsets cover single- and multi-document QA, summarization, few-shot learning, synthetic tasks, and code completion.
- Models and workloads: The Mixed Tasks workload changes its LongBench subset between consecutive prefill steps, capturing routing-distribution shifts while preserving batch size and input length.Each prefill step draws requests from one subset, and the reported row measures latency for the changing-subset sequence.
- Experimental setup: Experiments use eight NVLink-connected NVIDIA A800-SXM4 GPUs with EP=8, batch size 16, and 8K-token inputs, comparing Vanilla, FreeBalance, history-driven EPLB, and their combination.Each configuration is warmed up once, measured for three runs, and reported as the average.
- Metrics: Evaluation reports prefill latency, max-to-mean rank-load ratio, residual pre-routing similarity and top-k recovery, and the fraction of swaps that reduce realized imbalance.The load ratio equals one under perfect balance, and GPU timings use synchronized CUDA events.
4.1 End-to-End Effectiveness
FreeBalance improves MoE load balance and reduces end-to-end prefill latency by overlapping online reconfiguration with existing critical-path computation. It lowers per-layer max/mean load ratios while accelerating execution gated by the slowest EP rank.
- Load balance: FreeBalance reduces per-layer max/mean ratios from up to 2.01 with Vanilla to 1.35, improving the ratio by up to 32.8%.More even token distribution lowers maximum per-rank activation volume and expert workload during dispatch, combine, and execution.
- End-to-end latency: 13.1%: FreeBalance reduces Qwen3-30B’s average prefill latency across 19 LongBench subsets.Auxiliary work overlaps existing critical-path computation, so online reconfiguration introduces no separate pause.
4.2 Prediction and Planning Quality
FreeBalance’s early router invocation provides a strong signal for estimating target-layer expert workload, despite lower hidden-state similarity. Applying complete swap plans improves the measured max-to-mean rank-load ratio in 95% of migration-executed layers.
- Prediction quality: Logit cosine similarity remains 0.9896–0.9952 across both models and all three datasets, indicating strong alignment between early and normal router invocations.The early invocation is gℓ(hℓ−1), while the normal invocation is gℓ(Hℓ).
- Prediction quality: Hidden-state cosine similarity ranges from 0.7116 to 0.9316, while the resulting top-k hit rate ranges from 0.7419 to 0.8494.These results support using early router logits to estimate target-layer expert workload.
- Planning quality: 95% of layers with executed migrations achieve a lower measured max-to-mean rank-load ratio after applying the complete swap plan.The rate is measured at the layer level; remaining layers show no improvement because of runtime variation or non-additive interactions among multiple swaps.
4.3 Overlap Analysis
FreeBalance overlaps residual prediction, planning, and expert-weight transfer with attention computation, avoiding direct addition of migration costs to prefill latency. Adaptive migration budgets outperform fixed swap policies across sequence lengths by selecting swap counts according to the available overlap window.
- In the two-stream schedule, router invocation and planning overlap target-layer attention, while expert-weight transfer overlaps the remaining attention computation.
- The best fixed policies achieve 1.0582×, 1.0874×, and 1.1263× speedups at 1K, 2K, and 4K tokens, respectively.
- FreeBalance reaches 1.1241×, 1.1264×, and 1.1542× speedups at 1K, 2K, and 4K tokens by adapting average swaps from 1.82 to 2.68.It outperforms fixed policies at all three sequence lengths, indicating that the overlap budget should determine migration count.
4.4 Sensitivity to Sequence Length
At batch size 16, FreeBalance’s balancing benefit and migration-overlap opportunity increase with sequence length from 1K to 8K tokens. Across these lengths, it lowers TTFT and max-to-mean rank-load ratios while measured overlap time expands.
- Sensitivity to Sequence Length: 14.7% average TTFT reduction across evaluated sequence lengths, increasing from 11.0% at 1K tokens to 23.3% at 8K tokens.The evaluation fixes batch size at 16 and varies input length from 1K to 8K tokens.
- Sensitivity to Sequence Length: The max-to-mean rank-load ratio decreases from a Vanilla range of 2.00–2.04 to a FreeBalance range of 1.33–1.37.These results are reported for each evaluated sequence length.
- Sensitivity to Sequence Length: Measured overlap time increases from 12.55 ms to 36.80 ms, providing a larger window to hide expert migration for longer sequences.The study reports measured overlap time alongside TTFT and max-to-mean rank-load ratio.
5 Related Work
Offline balancing methods such as DeepSeek EPLB periodically replicate and place experts using historical routing statistics. This can amortize optimization and weight rearrangement for slowly varying workloads, but may lag abrupt task shifts and miss layer- and batch-specific imbalance in multi-task serving.
- Offline balancing: DeepSeek EPLB periodically replicates and places experts using historical routing statistics.Its approach is designed around offline balancing.
- Offline balancing: Historical workload estimation can lag behind abrupt task shifts.This limitation arises when workload patterns change faster than the statistics reflect.
- Offline balancing: Offline balancing cannot directly correct layer- and batch-specific imbalance in multi-task serving.Consecutive batches may activate different expert subsets, making historical estimates less responsive.
- Offline balancing: FreeBalance instead predicts the upcoming workload rather than relying solely on historical routing statistics.The supplied passage introduces this as its contrast with offline balancing.
6 Conclusion
FreeBalance anticipates expert workloads to initiate balancing before target-layer dispatch, adapting to dynamic workloads while hiding reconfiguration overhead and preserving routing decisions and outputs. It reduces rank-load imbalance and end-to-end prefill latency.
- Approach: FreeBalance anticipates expert workloads and initiates balancing before target-layer dispatch for dynamic distributed MoE inference.This shifts balancing earlier to adapt to each workload.
- Benefits: FreeBalance hides reconfiguration overhead from the critical path while preserving the original routing decisions and outputs.
- Results: 32.8%: FreeBalance reduces the max-to-mean rank-load ratio by up to 32.8%.
- Results: 13.1%: FreeBalance reduces end-to-end prefill latency by 13.1%.