FreeBalance: Pre-Routing Online MoE Load Balancing via Residual Workload Prediction

Summary (Overview)

• Core contribution: FreeBalance is a lossless online load-balancing framework for Mixture-of-Experts (MoE) models that predicts expert workloads before routing decisions are made, enabling expert migration to overlap with computation-heavy pre-routing stages (e.g., attention) rather than sitting on the critical inference path.
• Key innovation: The framework exploits cross-layer residual similarity in hidden representations—since $h_{\ell-1}$ and $h_\ell$ are highly similar due to small residual updates—to invoke the same frozen router on the earlier hidden state, producing a lightweight workload prediction without any additional trainable parameters.
• Results: FreeBalance reduces the max-to-mean rank load ratio by 32.8% and end-to-end prefill latency by 13.1% across models and datasets, hiding balancing overhead for an average of 5.1 experts per layer (which would otherwise account for 8.5% of critical-path latency).
• Lossless guarantee: The predicted workload is used only for expert placement decisions; the actual router still determines final token-to-expert assignments, preserving model outputs exactly.
• Practical design: A cost model constrains the number of expert swaps to fit within the available attention-time window, and a deterministic greedy planner ensures all ranks agree on the swap plan without broadcast communication.

---

Introduction and Theoretical Foundation

Background: MoE and Expert Parallelism

Mixture-of-Experts (MoE) architectures have become dominant in modern LLMs by dramatically increasing parameter capacity without proportionally increasing per-token computational cost. Each MoE layer activates only a small subset of experts per token through a lightweight routing mechanism:

$$W_l, I_l = \operatorname{topk}(G(h_l))$$ $$MoE(h_l) = \sum_{i \in I_l} W_{l,i} \cdot \operatorname{expert}_i(h_l)$$

In distributed inference, expert parallelism (EP) shards experts across devices. Tokens routed to remote experts require all-to-all communication for dispatch and result combination, with global synchronization across the EP group.

The Load Imbalance Problem

Load imbalance arises from skewed routing distributions across ranks. Since MoE execution requires global synchronization, the most heavily loaded rank determines the critical path, forcing less-loaded ranks to idle. This problem is exacerbated in multi-task serving scenarios where routing distributions vary significantly across task types—the paper shows only 4.1% expert-activation overlap between two representative tasks (NarrativeQA and LCC) on Qwen3-30B.

The Sequential Dependency Bottleneck

The root cause of online balancing overhead is the sequential dependency between workload acquisition and expert-map modification:

1. Routing can only begin after preceding computations (e.g., attention) complete
2. Expert migration can only begin after routing decisions are available
3. The migration window is therefore restricted to the narrow interval between routing and MoE computation—directly on the critical path

Key Insight

The paper observes that expert balancing does not need to wait for the current-layer router if the workload distribution can be predicted lightweightly and accurately before routing. This enables migration to overlap with preceding computation stages, particularly beneficial for long-sequence inference where attention dominates critical latency.

The residual network structure provides the theoretical foundation:

$$h_{\ell+1} = h_\ell + \text{Attention}(h_\ell) + \text{MoE}(h_\ell + \text{Attention}(h_\ell)) \tag{1}$$

Because the residual update is small relative to hidden-state magnitude, the angle between $h_{\ell-1}$ and $h_\ell$ is small, and their cosine similarity is close to 1.

---

Methodology

Overview: Four Stages per MoE Layer

For each MoE layer $\ell$, FreeBalance performs:

1. Residual workload prediction: Estimate token counts for all experts before routing
2. Budgeted expert-swap planning: Determine migration budget and select pairwise swaps
3. Weight transfer: Exchange expert weights across EP ranks, overlapped with attention
4. Lossless MoE computation: Execute with original routing decisions on migrated placement

Residual Workload Prediction

FreeBalance introduces no standalone predictor—it invokes the same frozen router $g_\ell$ twice:

$$\widehat{G}_\ell = g_\ell(h_{\ell-1}), \quad \text{early pre-routing logits} \tag{2}$$ $$G_\ell = g_\ell(H_\ell), \quad \text{final routing logits} \tag{3}$$

For a linear router:

$$g_\ell(X) = XW_\ell + \mathbf{1}b_\ell^\top \tag{4}$$

where $W_\ell \in \mathbb{R}^{d \times E}$ and $b_\ell \in \mathbb{R}^E$ are the original router parameters. Applying the top-k rule to $\widehat{G}_\ell$ yields estimated assignments, aggregated into predicted expert workloads:

$$\widehat{n}_{\ell,e} = \sum_{i=1}^{N} \mathbf{1}\left[e \in \widehat{\mathcal{T}}_\ell(i)\right] \tag{5}$$

Key properties:

• No additional predictor weights, training procedure, or checkpoint state
• The early invocation produces only an E-element workload vector for placement planning
• The normal invocation on $H_\ell$ remains authoritative for final assignments
• Using $h_{\ell-1}$ creates an attention-length migration window while retaining sufficient task/token information

Budgeted Expert-Swap Planning

Pairwise swap formulation: Exchanging expert $e_a$ on rank $a$ with expert $e_b$ on rank $b$ changes only those two ranks' loads:

$$L_a' = L_a - \widehat{n}_{e_a} + \widehat{n}_{e_b} \tag{6}$$ $$L_b' = L_b - \widehat{n}_{e_b} + \widehat{n}_{e_a} \tag{7}$$

Load objective to minimize:

$$\Phi(L) = \max_r L_r + \gamma \sum_r (L_r - \overline{L})^2 \tag{8}$$

The first term targets the straggler rank; the second discourages moving the bottleneck elsewhere. The predicted benefit is $\dot{\Delta}(e_a, e_b) = \Phi(L) - \Phi(L')$.

Migration cost model:

$$C(s) = \alpha_{a,b} + \frac{S(e_a) + S(e_b)}{\beta_{a,b}} \tag{9}$$

where $S(e)$ is the transferred representation size, $\beta_{a,b}$ is measured point-to-point bandwidth, and $\alpha_{a,b}$ captures launch/protocol overhead.

Budget constraint:

$$B_\ell = \max(0, T_{r(\ell)}^{\text{attn}} - \delta) \tag{10}$$

where $\delta$ is a safety margin and $T_{r(\ell)}^{\text{attn}}$ is the profiled attention time (reused across layers with the same attention mechanism). The budget is enforced per link and per rank, not just as a global sum.

Deterministic greedy planning: All ranks receive the same globally aggregated count vector, enumerate candidate pairs, remove conflicts and non-positive-gain candidates, and sort by decreasing benefit-to-cost ratio. Ties are resolved lexicographically. This eliminates the need for a plan broadcast.

---

Empirical Validation / Results

Experimental Setup

• Models: Qwen3-30B-A3B-Instruct-2507 (128 experts, top-8, 16 experts/rank, 9 MB/expert) and Moonlight-16B-A3B-Instruct (64 experts, top-6, 8 experts/rank, 16.5 MB/expert)
• Hardware: 8× NVIDIA A800-SXM4 GPUs with NVLink, EP=8
• Benchmarks: 19 LongBench subsets plus a Mixed Tasks workload with step-to-step task shifts
• Default config: batch size 16, input length 8K tokens

End-to-End Effectiveness

Table 1 (excerpt) shows prefill latency improvements across LongBench subsets:

• Vanilla max/mean ratio up to 2.01 reduced to 1.35 (32.8% improvement)
• Average prefill latency reduction of 13.1% on Qwen3-30B

Prediction Quality

Table 2: Pre-routing quality metrics

Despite hidden-state cosine similarity ranging from 0.7116 to 0.9316, logit cosine similarity remains between 0.9896 and 0.9952, with top-k hit rates of 0.7419–0.8494. This demonstrates that router logits remain strongly aligned across invocation points.

Overlap Analysis

• Expert migrations require approximately 3–4 ms per layer
• 95% of executed swap plans achieve a lower measured max-to-mean ratio
• Figure 4 shows speedup is not monotonic in fixed swaps per layer; adaptive FreeBalance (1.82–2.68 avg swaps) outperforms all fixed policies at 1K, 2K, and 4K token lengths

Sensitivity to Sequence Length

Table 3: Performance under batch size 16

TTFT reduction grows from 11.0% at 1K to 23.3% at 8K tokens (14.7% average), confirming that longer sequences provide larger windows for hiding migration overhead.

---

Theoretical and Practical Implications

Theoretical Significance

1. Breaking the sequential dependency: The paper demonstrates that online load balancing need not be reactive—workload prediction enables proactive placement, fundamentally changing the timing constraints of expert migration.

2. Residual similarity as a free prediction signal: The observation that frozen routers can be invoked on earlier hidden states (with logit cosine similarity > 0.99) provides a zero-cost prediction mechanism requiring no additional training or parameters.

3. Budget-constrained balancing: The cost model formalizes the trade-off between balancing benefit and migration overhead, showing that the overlap window (not the balancing potential) should determine migration aggressiveness.

Practical Implications

1. Lossless acceleration: FreeBalance preserves model outputs exactly, making it safe for production deployment without accuracy concerns.

2. Multi-task adaptability: Unlike offline methods (e.g., DeepSeek EPLB) that lag behind task shifts, FreeBalance adapts at layer granularity to each batch's routing distribution.

3. Scalability: The deterministic planner eliminates plan broadcast overhead, and the frontier-based candidate generation (considering only overloaded/underloaded ranks) keeps planning lightweight enough for layer-granularity execution.

4. Complementarity: FreeBalance works alongside existing methods (EPLB + FreeBalance consistently outperforms either alone), suggesting it can be integrated into existing serving stacks.

---

Conclusion

FreeBalance addresses dynamic load imbalance in distributed MoE inference by anticipating expert workloads before target-layer dispatch and initiating balancing proactively. Key takeaways:

1. Residual workload prediction using the frozen router on earlier hidden states provides accurate, zero-cost workload estimates before routing.
2. Budgeted expert-swap planning with a cost model ensures migration fits within the available attention-time window, hiding overhead from the critical path.
3. Lossless execution preserves original routing decisions and model outputs exactly.
4. Empirical results: 32.8% reduction in max-to-mean rank load ratio, 13.1% end-to-end prefill latency reduction, with balancing overhead for 5.1 experts per layer fully hidden.

Future directions suggested by this work include extending the approach to decode-phase inference, exploring prediction-based balancing for other parallelization strategies, and investigating whether residual similarity can enable other proactive optimizations in distributed inference systems.