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

> FreeBalance predicts expert workloads by reusing the frozen router on earlier hidden states, enabling lossless, proactive load balancing that reduces max-to-mean rank load ratios by 32.8% and prefill latency by 13.1%.

- **Source:** [arXiv](https://arxiv.org/abs/2608.14205)
- **Published:** 2026-08-29
- **Permalink:** https://picx.dev/p/Y0iOBN
- **Whiteboard:** https://picx.dev/p/Y0iOBN/image

## Summary

# 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:

| Subset | Qwen Vanilla | Qwen +Ours | Qwen EPLB | Qwen EPLB+Ours | Moon Vanilla | Moon +Ours |
|--------|-------------|------------|-----------|----------------|--------------|------------|
| NarrativeQA | 147.6 | **126.6** | 135.9 | 126.6 | 77.4 | **66.7** |
| TriviaQA | 175.9 | **112.1** | 122.6 | 110.7 | 65.7 | **64.4** |
| VCSUM | 138.1 | **115.8** | 125.6 | 112.9 | 73.7 | **62.6** |
| Mixed Tasks | 83.1 | **78.0** | 78.3 | 75.3 | 45.9 | **43.3** |

- 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

| Model | Dataset | Hidden cos. | Logit cos. | Top-k hit |
|-------|---------|-------------|------------|-----------|
| Qwen | PassageRetrieval | 0.7334 | **0.9952** | 0.7520 |
| Qwen | LSHT | 0.7242 | **0.9951** | 0.7578 |
| Qwen | TriviaQA | 0.7116 | **0.9947** | 0.7419 |
| Moonlight | PassageRetrieval | 0.9316 | **0.9900** | 0.8461 |
| Moonlight | LSHT | 0.9291 | **0.9896** | 0.8494 |
| Moonlight | TriviaQA | 0.9234 | **0.9910** | 0.8037 |

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

| Length | Vanilla TTFT | Vanilla Max/Mean | Ours TTFT | Ours Max/Mean | Overlap (ms) |
|--------|-------------|-----------------|-----------|---------------|--------------|
| 1K | 3.08 | 2.00 | 2.74 | 1.33 | 12.55 |
| 2K | 4.99 | 2.02 | 4.43 | 1.34 | 17.35 |
| 4K | 8.98 | 2.03 | 7.78 | 1.35 | 31.30 |
| 8K | 20.45 | 2.04 | 15.68 | 1.37 | 36.80 |

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.

---

_Markdown view of https://picx.dev/p/Y0iOBN, served by PicX — AI-generated visual whiteboard summaries of research papers._