Paper deep dive
FreeBalance: Pre-Routing Online Moe Load Balancing via Residual Workload Prediction
Pengfei Chen, Yize Wu, Shouxu Kuang, Ke Gao, Ling Li
Intelligence
Status: succeeded | Model: Gemma-4-26B-A4B | Prompt: intel-v1 | Confidence: 91%
Last extracted: 8/17/2026, 5:17:29 AM
Summary
The paper introduces FreeBalance, a lossless online load-balancing framework for Mixture-of-Experts (MoE) models in distributed inference. It addresses load imbalance by predicting expert workloads using residual hidden representations from previous layers, allowing expert migration to overlap with pre-routing computations (like attention). This approach reduces the max-to-mean rank load ratio by 32.8% and end-to-end prefill latency by 13.1%.
Entities (6)
Relation Signals (4)
FreeBalance → uses → Residual Workload Prediction
confidence 95% · FreeBalance leverages cross-layer similarities in hidden representations within the residual network to build a lightweight workload predictor.
FreeBalance → reduces → Latency Reduction
confidence 92% · 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%.
FreeBalance → optimizes → Expert Parallelism
confidence 90% · Load imbalance poses a major bottleneck to the efficiency of expert parallelism in distributed inference of Mixture-of-Experts (MoE) models.
Residual Workload Prediction → enables → Pre-Routing Balancing
confidence 88% · This enables proactive expert migration planning before routing decisions are available, creating substantial overlap between weight transfer and computation-heavy pre-routing stages.
Cypher Suggestions (0)
No Cypher suggestions yet.
Abstract
Abstract: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.
Tags
Links
- Source: https://arxiv.org/abs/2608.14205v1
- Canonical: https://arxiv.org/abs/2608.14205v1
Trouble viewing inline? Open PDF directly →
Full Text
39,292 characters extracted from source content.
Expand or collapse full text
FreeBalance: Pre-Routing Online MoE Load Balancing via Residual Workload Prediction Pengfei Chen Yize Wu Shouxu Kuang Ke Gao Ling Li Abstract 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 8.5% of the critical-path latency. Figure 1: Overview of FreeBalance and its four stages: ① residual workload prediction, ② budgeted expert-swap planning, ③ expert-weight migration overlapped with pre-routing attention, and ④ lossless MoE computation. The residual representation from layer ℓ−1 -1 predicts the expert workload of layer ℓ before its router executes. The prediction-guided expert placement is available before the target expert FFN executes. The blue line denotes the FreeBalance execution flow (ours), whereas the black line denotes the original inference flow. 1 Introduction Mixture-of-Experts (MoE) architecture has become a dominant design choice for modern large language models (LLMs). It dramatically increases parameter capacity without proportionally increasing each-token computational cost (3; 4). Specifically, each MoE layer activates only a small subset of experts for each input token through a lightweight routing mechanism, which effectively decouples performance capacity from inference latency (15). In distributed MoE inference, expert parallelism (EP) is commonly adopted for expert-stage acceleration. Experts of each MoE layer are sharded across multiple devices, allowing MoE computation to be executed in parallel. When tokens are routed to experts residing on different devices, the corresponding activation tensors are first dispatched to the corresponding devices via an all-to-all communication, then the results are combined after expert computation with a second all-to-all(9; 14). The effectiveness of EP is largely limited by load imbalance, which arises from uneven routing distribution across ranks. Despite the auxiliary load-balancing loss in the training process(6), inference-time routing distributions can still be highly skewed (10). Since both the all-to-all communication and subsequent expert computation require synchronization across the EP group, the overall execution latency is determined by the most heavily loaded rank. Consequently, lightly loaded devices must remain idle while waiting for stragglers to complete, resulting in reduced device utilization and degraded end-to-end performance. Existing load balancing methods primarily rely on expert-device mapping modification, which changes the assignment of experts to physical devices to redistribute workloads more evenly. Such mappings can be determined offline based on historical routing statistics (5). However, practical LLM serving systems typically handle multi-task requests, where routing distributions can vary significantly across task types, and a static expert-device mapping cannot instantly adapt to dynamic workload changes, as illustrated in Figure 2. In this case, online load balancing is therefore required, which dynamically adjusts expert-device assignments at runtime by migrating expert weights across ranks. However, the routing distribution required to guide expert placement is only available after the routing decision is made. Therefore, the available window for expert migration is restricted to the interval between routing completion and expert execution, which lies directly on the critical inference path and introduces additional latency overhead. We identify the root cause of this problem as the sequential dependency between workload distribution acquisition and expert-map modification. Specifically, the routing stage can only begin after all preceding computations (e.g., attention computation of the current layer) are completed, and expert migration can only be initiated afterward based on the observed routing distribution. In this work, we observe that expert balancing does not necessarily need to wait for the current-layer router execution if the workload distribution can be predicted lightweightly and accurately before routing. Such pre-routing prediction enables expert migration to be overlapped with preceding computation stages, making it possible to effectively hide the balancing overhead. Furthermore, this approach is particularly beneficial for long-sequence inference, where attention computation often constitutes a significant portion of the critical latency. Based on this insight, we propose FreeBalance, an online MoE load balancing framework (Figure 1) that performs pre-routing balancing through residual workload prediction. We observe that the routing distribution of an upcoming MoE layer can be well approximated before routing, via feeding the residual representation from the preceding layer into the corresponding router. This prediction is lightweight compared to the whole inference latency, and most importantly, it provides workload information before the actual routing stage. Based on the predicted workload, FreeBalance determines the online expert mapping adjustments and initiates weight migration immediately, so that the migration process can be overlapped with other pre-routing stages (primarily attention computations) and does not extend the critical path. Moreover, FreeBalance employs a cost model to constrain the amount of migrated experts according to the available overlapping opportunity, preventing overly aggressive adjustments from introducing additional latency. Notably, FreeBalance uses the predicted workload solely for online balancing, while the actual MoE execution still follows the original routing decisions, thereby achieving acceleration without affecting model outputs. Experiments show that FreeBalance reduces the max-to-mean rank load ratio by 32.8%, reduces end-to-end latency by 13.1%. Specifically, our method hides online balancing overhead of averagely 5.1 experts per layer, which would otherwise account for 8.5% of the critical-path latency. 2 Preliminary 2.1 Mixture-of-Expert Architecture An MoE layer contains a router and a set of feed-forward experts. For an input token representation hlh_l of layer l, the router produces a score for each expert and selects the top-k experts. Wl,Il=topk(G(hl)).W_l,I_l=topk(G(h_l)). Only the selected experts process the token, after which their outputs are weighted aggregated. MoE(hl)=∑i∈IlWl,i∗experti(hl),MoE(h_l)= _i∈ I_lW_l,i*expert_i(h_l), where Wl,iW_l,i denotes the routing weight for expert i. The aggregated result is then forwarded as the output of the MoE block (added to hlh_l along with the output of the attention block), such that hl+1=hl+Attention(hl)+MoE(hl+Attention(hl))h_l+1=h_l+Attention(h_l)+MoE(h_l+Attention(h_l)) (1) 2.2 Expert Parallelism Under EP, each expert is placed onto one or more specific devices according to an expert–device mapping (typically an even distribution). Each device hosts a subset of experts and processes tokens routed to them. When tokens are assigned to experts that reside on remote devices, an all-to-all communication is required to dispatch token representations to the corresponding devices. After computation, the expert outputs are communicated back to the original devices via a second all-to-all, where a final weighted aggregation is performed. 2.3 Load Imbalance Load imbalance poses a severe challenge under EP, arising from the skewed token routing distribution across EP ranks. When hot experts are colocated on the same rank, expert-level routing skew translates into inter-rank workload imbalance. As MoE execution requires global synchronization operations of result-gathering communication and expert-output aggregation, the most heavily loaded rank determines the critical path, forcing less-loaded ranks to wait in idle. Furthermore, in practical multi-task serving scenario, routing skew varies significantly across task types rather than remaining fixed. Figure 2 presents the heatmap of workload distributions of two representative tasks on Qwen3-30B. As demonstrated, workload distributions differ substantially across tasks, with only 4.1% of expert-activation overlap across all layers. Figure 2: Per-layer relative expert loads for NarrativeQA and LCC. Experts are sorted independently at each layer by their NarrativeQA load, and the same ordering is used for LCC to expose the workload-dependent shift in expert popularity. 2.4 Load Balancing Offline load balancing typically decides expert placement for more balanced workload distributions from historical routing statistics. Despite its effectiveness, in multi-task scenarios where workload distributions vary significantly across tasks (as discussed in section 2.2), a fixed placement strategy may become suboptimal for the varying workload. In this case, workload-aware online balancing during inference is essential, which dynamically adapts expert-device assignments by migrating expert weights across ranks. The key challenge of online load balancing lies in the additional overhead of expert migration. Since placement decisions can only be made after the routing distribution is computed, the available window for expert migration is limited to a narrow interval between routing and MoE computation, which falls directly on the critical inference path. Consequently, the introduced additional overhead can largely diminish the benefit of improved load balance. 3 Method 3.1 Motivation and Overview As discussed above, the primary bottleneck of online load balancing is the migration overhead incurred on the critical inference path. Since expert placement decisions depend on the routing distribution, expert migration cannot begin until the target layer’s router has completed, i.e., after the preceding attention computation. Consequently, only a narrow window remains for planning and expert-weights migration before MoE computation. We observe that a substantial overlap opportunity naturally exists on the critical path. The pre-routing stages (e.g., attention computation) are computation-heavy, especially for long-sequence workloads, providing a sufficiently long window to hide expert migration latency. However, this opportunity remains unexploited because routing statistics become available only after the router executes, and migration planning can begin only afterward. Consequently, the entire pre-routing computation stage is unavailable for overlapping expert migration. FreeBalance eliminates this limitation by breaking the sequential dependency between routing and migration planning. The key insight is that if the workload distribution can be accurately predicted before routing execution, migration planning no longer needs to wait for the actual routing results. Expert weight migration can therefore be scheduled proactively and overlapped with preceding computation stages. Therefore, we propose residual workload prediction for lightweight and accurate online workload prediction before the actual routing execution. Furthermore, to maximize the latency benefit while limiting migration overhead, FreeBalance employs a cost model to bound the migration budget and control migration planning. Furthermore, migration plans are constructed using pairwise expert swaps, ensuring that every rank preserves the same number of expert slots and memory footprint. Notably, FreeBalance preserves the original routing decisions during MoE execution, ensuring lossless inference. The predicted workload distribution is used solely for expert placement decisions and weight transfer, while the target router determines the final token-to-expert assignments (rather than skipped). The final routing results may alter after the expert migration, yet the overall correctness is totally guaranteed. In conclusion, for each MoE layer ℓ , FreeBalance performs four stages: 1. Residual workload prediction. A lightweight prediction module estimates the token counts of all experts before routing. 2. Budgeted Expert-Swap Planning. A cost model determines the migration budget, and pairwise expert swaps are selected within the available overlap budget. 3. Weight Transfer. Expert weights are exchanged across EP ranks, with the migration latency overlapped with computation-heavy pre-routing stages (mainly the attention stage). 4. Lossless MoE computation. The target router computes the exact routing decisions, and MoE execution proceeds using the migrated expert placement, achieving a more balanced workload without altering the original routing behavior. 3.2 Residual Workload Prediction Residual connections make the hidden representations of adjacent layers highly similar (8; 13; 17). As illustrated in fig. 3 and formalized in eq. 1, hℓh_ is obtained by adding the attention and MoE updates to hℓ−1h_ -1. Because this residual update is small relative to the hidden-state magnitude, the angle between hℓ−1h_ -1 and hℓh_ is small, and their cosine similarity is close to 1. FreeBalance exploits this cross-layer similarity for lightweight and accurate online workload prediction. Since hℓ−1h_ -1 is available before the subsequent attention and routing stages, expert-placement planning and weight migration can begin at pre-routing time. Figure 3: Geometric intuition for cross-layer residual similarity. The attention and MoE updates cause only a small directional change from hℓ−1h_ -1 to hℓh_ . Migration must begin before the target layer routes its tokens, so the final router output is available too late to guide its own placement for that same layer. FreeBalance instead predicts the workload of layerℓlayer_ from hℓ−1h_ -1, the output hidden states of layerℓ−1layer_ -1 and therefore the residual-stream input to layerℓlayer_ . This state is available before the target layer performs attention and normal routing. Because consecutive layers process the same token sequence, hℓ−1h_ -1 retains task- and token-specific information that is useful for anticipating the next router’s aggregate demand. The prediction guides only physical expert placement and never replaces model’s normal routing pass. Let hℓ−1∈ℝN×dh_ -1 ^N× d denote the output of layer ℓ−1 -1 for a prefill batch of N tokens. It is also the residual hidden state entering target layerℓ before attention. Let Hℓ∈ℝN×dH_ ^N× d denote the target router’s normal input after the target layer’s pre-routing transformations. FreeBalance introduces no standalone workload predictor; instead, it invokes the same frozen router gℓg_ twice: G^ℓ G_ =gℓ(hℓ−1), =g_ (h_ -1), early pre-routing logits, pre-routing logits, (2) Gℓ G_ =gℓ(Hℓ), =g_ (H_ ), final routing logits. routing logits. (3) For a linear router, both invocations share the original routing function gℓ(X)=XWℓ+bℓ⊤,g_ (X)=XW_ +1b_ , (4) where Wℓ∈ℝd×EW_ ^d× E and bℓ∈ℝEb_ ^E are the original router parameters. Applying the original top-k rule to G^ℓ G_ yields estimated assignments ^ℓ(i) T_ (i) for token i. Aggregating these assignments gives the predicted expert workload n^ℓ,e=∑i=1N[e∈^ℓ(i)]. n_ ,e= _i=1^N1\! [e∈ T_ (i) ]. (5) Because hℓ−1h_ -1 precedes the target layer’s attention whereas HℓH_ follows its pre-routing transformations, the early assignments are estimates rather than final model decisions. Both invocations nevertheless use the same router parameters. Consequently, FreeBalance adds no 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ℓH_ remains unchanged and exclusively determines the final token-to-expert assignments. The choice of hℓ−1h_ -1 balances timeliness and prediction fidelity. It incorporates all task- and token-dependent transformations through layer ℓ−1 -1, yet is available before attention in layer ℓ . In contrast, HℓH_ is closer to the target routing decision but becomes available only when the normal router is about to execute. Using hℓ−1h_ -1 therefore creates an attention-length migration window while retaining sufficient information to estimate which experts will be hot in the target layer. The early workload estimate is used only to plan physical expert placement, while the normal router invocation remains authoritative and determines which experts process each token. An empty swap plan retains the current placement. Therefore, the rebalancing procedure preserves the model’s routing semantics and baseline execution. Algorithm 1 FreeBalance for Target Layer ℓ 1: hℓ−1h_ -1, gℓg_ , πℓ _ , migration budget BℓB_ , and threshold τ 2: Budget-feasible swap plan ℓS_ and updated placement π′ℓπ^ _ 3: Predict next-layer expert workloads using gℓ(hℓ−1)g_ (h_ -1) 4: Aggregate expert workloads into rank loads under πℓ _ 5: ℓ←∅S_ ← 6: while rank-pair imbalance exceeds τ and budget remains do 7: Select a beneficial swap between overloaded and underloaded ranks 8: if no such swap exists then 9: break 10: end if 11: end while 12: Migrate the selected experts during target-layer attention 13: Commit the placement before the target MoE stage begins 14: Execute the target MoE with the normal router decisions 15: return ℓS_ and π′ℓπ^ _ 3.3 Budgeted Expert-Swap Planning Given the current batch’s expert workload obtained by residual workload pre-routing, the planner seeks a placement that lowers the predicted maximum rank load without scheduling more communication than target-layer attention can hide. It must also produce the same result on every rank. We formulate planning as a budgeted optimization over pairwise swaps and use a deterministic greedy procedure that is lightweight enough to run at layer granularity. Consider expert eae_a on rank a and expert ebe_b on rank b. Exchanging them changes only the loads of these two ranks: La′ L _a =La−n^ea+n^eb, =L_a- n_e_a+ n_e_b, (6) Lb′ L _b =Lb−n^eb+n^ea. =L_b- n_e_b+ n_e_a. (7) We score a candidate by its reduction in a load objective Φ(L)=maxrLr+γ∑r(Lr−L¯)2, (L)= _rL_r+γ _r(L_r- L)^2, (8) where the first term targets the straggler rank and the second discourages moving the bottleneck to a different rank. The predicted benefit is Δ(ea,eb)=Φ(L)−Φ(L′) (e_a,e_b)= (L)- (L ). Candidates with non-positive benefit are discarded. Pairwise exchange has three practical advantages. It preserves the number of experts per rank, requires only bounded staging memory, and makes rollback straightforward because the old and new owners are known. It also avoids replication, which would require the router to divide assignments among multiple copies and would change the dispatch policy. The usable migration budget is derived from a profiled attention time and the topology between the participating ranks. For a swap s, the planner estimates C(s)=αa,b+S(ea)+S(eb)βa,b,C(s)= _a,b+ S(e_a)+S(e_b) _a,b, (9) where S(e)S(e) is the transferred representation of expert e, βa,b _a,b is the measured point-to-point bandwidth, and αa,b _a,b captures launch and protocol overhead. The parameters are calibrated during runtime initialization and may distinguish intra-node and inter-node links. For a fixed input shape, layers that use the same attention mechanism have nearly identical attention latency. FreeBalance therefore profiles attention once per attention type and reuses the measurement for subsequent layers of that type. Let r(ℓ)r( ) denote the first layer that uses the same attention mechanism as target layer ℓ . In a homogeneous model, r(ℓ)=0r( )=0, so the layer-0 attention time is reused throughout the model. For a hybrid-attention model, FreeBalance instead profiles the first occurrence of each attention type, such as linear or sparse attention, and reuses the corresponding measurement for later layers of the same type. The planner admits migrations only within Bℓ=max(0,Tr(ℓ)attn−δ),B_ = (0,T^attn_r( )-δ), (10) where δ is a safety margin for prediction error and stream interference. The budget is enforced per link and per rank rather than only as a global sum, preventing several individually valid swaps from oversubscribing one participant. Every rank receives the same globally aggregated count vector and begins from the same placement map. Each rank enumerates candidate pairs, removes conflicts and non-positive-gain candidates, and sorts the remainder by decreasing benefit-to-cost ratio. Ties are resolved lexicographically by layer identifier, source rank, destination rank, and logical expert identifiers. The planner accepts a candidate if neither expert has already been selected and if adding the transfer respects all communication budgets. Loads and residual budgets are updated after each accepted swap. This deterministic ordering eliminates a plan broadcast: ranks exchange only the compact load statistics required to construct a global view, then reproduce the same swap sequence locally. Determinism also simplifies debugging because a layer, placement map, and predicted count vector uniquely determine the plan. In practice, the planner considers only experts on the most overloaded and most underloaded ranks, which reduces candidate generation from all expert pairs to a small frontier without changing the common heavy-to-light case. Qwen3-30B-A3B-Instruct-2507 Moonlight-16B-A3B-Instruct LongBench Subset Vanilla Vanilla +Ours EPLB EPLB +Ours Vanilla Vanilla +Ours EPLB EPLB +Ours Single-Document QA NarrativeQA 147.6 126.6 135.9 126.6 77.4 66.7 63.6 63.2 Qasper 102.6 94.6 105.6 92.3 57.1 55.4 57.5 55.3 MultiFieldQA (EN) 85.3 74.8 76.9 74.7 43.1 42.5 43.4 42.7 Multi-Document QA HotpotQA 137.3 122.4 121.0 120.9 71.6 69.8 68.2 67.1 2WikiMQA 125.2 110.2 100.1 93.3 66.7 63.1 64.9 62.8 MuSiQue 138.0 123.5 122.6 120.1 73.0 71.4 72.2 69.6 DuReader 142.4 124.5 138.5 123.0 78.1 70.7 78.7 69.5 Summarization GovReport 135.4 118.1 132.7 118.3 70.0 66.7 69.3 66.1 QMSum 137.6 123.2 130.4 122.4 73.0 70.8 72.7 71.4 MultiNews 53.6 48.4 47.6 47.6 27.6 27.3 26.9 26.7 VCSUM 138.1 115.8 125.6 112.9 73.7 62.6 70.5 61.9 Few-Shot Learning TREC 130.0 118.2 118.1 106.8 66.5 63.1 65.9 62.0 TriviaQA 175.9 112.1 122.6 110.7 65.7 64.4 64.8 62.7 SAMSum 131.5 111.8 123.1 111.9 69.3 62.7 66.4 62.8 LSHT 150.0 124.7 151.0 125.3 82.6 71.6 80.0 69.9 Synthetic Tasks PassageCount 132.2 121.8 121.5 121.4 72.1 70.0 69.9 69.6 PassageRetrieval (EN) 137.4 123.8 136.0 122.7 72.1 71.0 69.9 69.7 Code Completion LCC 62.3 55.1 61.2 54.1 37.3 32.9 36.9 33.6 RepoBench-P 129.7 111.8 124.7 111.4 77.3 67.0 77.2 67.9 Mixed-Task Workload Mixed Tasks 83.1 78.0 78.3 75.3 45.9 43.3 44.5 43.1 Table 1: Prefill latency across 19 LongBench subsets and one mixed dataset. Columns compare Vanilla, Vanilla with FreeBalance, EPLB, and EPLB with FreeBalance for Qwen3-30B-A3B and Moonlight-16B-A3B. Each entry reports latency in seconds; lower is better. 4 Evaluation Models We evaluate two MoE models: Qwen3-30B-A3B-Instruct-2507 (18) and Moonlight-16B-A3B-Instruct (12). Qwen3-30B-A3B-Instruct-2507 contains 128 experts, activates the top-8 experts per token, and places 16 experts on each EP rank; each expert occupies 9 MB. Moonlight-16B-A3B-Instruct contains 64 experts, activates the top-6 experts per token, and places 8 experts on each EP rank; each expert occupies 16.5 MB. In both models, one expert constitutes the basic unit of expert migration. Benchmarks We use LongBench (1) as the primary workload suite because it covers six distinct long-context capabilities. Table 1 reports 19 subsets: covering single- and multi-document QA, summarization, few-shot learning, synthetic tasks, and code completion. Each subset is evaluated independently under the same batching and sequence-length configuration. To emulate dynamic multi-task serving, we additionally construct the Mixed Tasks workload as a sequence of prefill steps whose LongBench subset changes from one step to the next. Each individual prefill step draws its requests from one subset. The workload therefore captures step-to-step shifts in routing distributions while retaining the same batch size and input-length configuration as the subset-specific runs. Its row reports latency under this changing-subset execution sequence. For the history-driven EPLB baseline, we construct its profiling state separately by aggregating expert-load statistics from 20 samples from each of all 21 LongBench subsets, for 420 profiling samples in total. Configurations Experiments run on nodes equipped with 8× NVIDIA A800-SXM4 GPUs connected by NVLink. We use expert parallelism across all eight GPUs (EP=8=8). Unless stated otherwise, we use a batch size of 16 with an input length of 8K tokens. Each configuration is warmed up once and then measured for three runs; we report the average of the measured runs. We compare Vanilla with a fixed expert placement, Vanilla with FreeBalance, history-driven EPLB, and EPLB with FreeBalance under the same model weights, routing decisions, batching policies, and parallel configurations. Metrics We report end-to-end prefill latency and the max-to-mean rank-load ratio, where a value of one indicates perfect balance. For residual pre-routing, we measure the cosine similarity between hℓ−1h_ -1 and HℓH_ , the cosine similarity between their router logits, and the top-k hit rate, defined as the fraction of final top-k assignments recovered by the early routing pass. Planner quality is measured by the fraction of executed swaps that reduce the realized imbalance. All GPU timings are collected with synchronized CUDA events. 4.1 End-to-End Effectiveness Vanilla reaches per-layer max/mean ratios up to 2.012.01, whereas FreeBalance reduces them to 1.351.35, an improvement of up to 32.8%32.8\%. The gain spans communication and computation: a more even token distribution reduces both the maximum per-rank activation volume during dispatch and combine and the maximum per-rank expert workload. Since all-to-all collectives and expert execution are gated by the slowest EP rank, mitigating this straggler accelerates the entire MoE stage. Across the 19 evaluated LongBench subsets, FreeBalance reduces Qwen3-30B’s average prefill latency by 13.1%13.1\%. Because its auxiliary work overlaps with existing critical-path computation, as analyzed in Section 4.3, online reconfiguration introduces no separate pause. The load-balance improvement does not translate linearly into wall-clock speedup because other kernels remain on the critical path. 4.2 Prediction and Planning Quality Table 2 compares the early router invocation gℓ(hℓ−1)g_ (h_ -1) with the normal invocation gℓ(Hℓ)g_ (H_ ). Across both models and all three datasets, logit cosine similarity remains between 0.98960.9896 and 0.99520.9952, even though hidden-state cosine similarity ranges from 0.71160.7116 to 0.93160.9316. The resulting top-k hit rate ranges from 0.74190.7419 to 0.84940.8494. These results show that the router logits remain strongly aligned across the two invocation points and provide a useful signal for estimating the target layer’s expert workload. Across layers in which migrations are executed, 95%95\% achieve a lower measured max-to-mean rank-load ratio after applying the complete swap plan. This rate is measured at the layer level rather than for individual swaps. The lack of improvement in the remaining layers arises from runtime variation or interactions among multiple swaps whose isolated benefits are not perfectly additive. 4.3 Overlap Analysis Residual pre-routing and planning consume computation, whereas the measured expert migrations require approximately 33–44 ms per layer. A sequential design would add these costs directly to prefill latency. In the two-stream schedule, the early router invocation and planning execute concurrently with the target layer’s attention. Once the plan is available, expert-weight transfer overlaps with the remaining attention computation. The remaining question is how many swaps can be admitted without exposing a migration tail on the critical path. Model Dataset Hidden cos. Logit cos. Top-k hit Qwen PassageRetrieval 0.7334 0.9952 0.7520 LSHT 0.7242 0.9951 0.7578 TriviaQA 0.7116 0.9947 0.7419 Moonlight PassageRetrieval 0.9316 0.9900 0.8461 LSHT 0.9291 0.9896 0.8494 TriviaQA 0.9234 0.9910 0.8037 Table 2: Pre-routing quality across models and datasets. Router logit and hidden cosine similarities compare the early and normal routing passes; top-k hit measures assignment overlap. Figure 4 shows that speedup is not monotonic in a fixed number of swaps per layer. The best fixed policies achieve 1.0582×1.0582×, 1.0874×1.0874×, and 1.1263×1.1263× speedups at 1K, 2K, and 4K tokens, respectively. By adapting the average swap count from 1.82 to 2.68, FreeBalance reaches 1.1241×1.1241×, 1.1264×1.1264×, and 1.1542×1.1542×, outperforming the fixed policies at all three lengths. These results show that the overlap budget should determine the migration count rather than a fixed policy. Figure 4: Speedup over Vanilla as a function of expert swaps per layer. Lines show zero to four forced swaps per layer; stars show adaptive FreeBalance, whose average swap count ranges from 1.82 to 2.68 across sequence lengths. 4.4 Sensitivity to Sequence Length We evaluate how sequence length affects the balancing benefit and the opportunity to hide expert migration. We fix the batch size to 16 and vary the input length from 1K to 8K tokens. For each configuration, we report time to first token (TTFT) and the max-to-mean rank-load ratio for Vanilla and FreeBalance. We additionally report the measured overlap time available for migration. Length Vanilla Ours Overlap (ms) TTFT Max/Mean TTFT Max/Mean 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 Table 3: Performance under batch size 16. TTFT is reported in seconds, and overlap time is reported in milliseconds. As shown in Table 3, FreeBalance reduces TTFT by 14.7%14.7\% on average across the evaluated sequence lengths. The reduction grows from 11.0%11.0\% at 1K tokens to 23.3%23.3\% at 8K tokens. Meanwhile, the max-to-mean ratio decreases from a Vanilla range of 2.002.00–2.042.04 to a FreeBalance range of 1.331.33–1.371.37. The measured overlap time increases from 12.5512.55 ms to 36.8036.80 ms, indicating that longer sequences provide a larger window for hiding expert migration. 5 Related Work Offline balancing. DeepSeek EPLB periodically replicates and places experts using historical routing statistics (5). While effective for slowly varying workloads by amortizing optimization and weight rearrangement, its historical workload estimation may lag behind abrupt task shifts and cannot directly correct layer- and batch-specific imbalance, especially in multi-task serving where consecutive batches activate different expert subsets. FreeBalance instead predicts the upcoming workload within the current prefill and adjusts placement at layer granularity. Online balancing. Dynamic systems adapt expert placement to observed demand (7; 10), while Harmony coordinates runtime scheduling with model-state movement (11). Recent approaches such as UltraEP and MoonEP rebalance the realized routing distribution using redundant experts, with runtime planning and expert-state movement after routing decisions (16; 2). Such reactive strategies depend on already-observed routing patterns and provide limited opportunity to hide rebalancing overhead before expert computation. FreeBalance instead anticipates upcoming expert demand before the target routing stage, budgets swaps within the available attention window, and overlaps expert migration with main-stream computation. 6 Conclusion FreeBalance addresses dynamic load imbalance in distributed MoE inference by anticipating expert workloads and initiating balancing before target-layer dispatch. By shifting balancing earlier, it adapts to each workload while hiding reconfiguration overhead from the critical path and preserving the original routing decisions and outputs. Experiments show that FreeBalance reduces the max-to-mean rank-load ratio by up to 32.8% and end-to-end prefill latency by 13.1%. References Bai et al. (2024) Y. Bai, X. Lv, J. Zhang, H. Lyu, J. Tang, Z. Huang, Z. Du, X. Liu, A. Zeng, L. Hou, Y. Dong, J. Tang, and J. Li LongBench: a bilingual, multitask benchmark for long context understanding. In Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), p. 3119–3137. External Links: Document Cited by: §4. Chen et al. (2026) Y. Chen, C. Li, Y. Wang, and M. Wei MoonEP: a perfectly balanced expert parallelism library via dynamic redundant experts. Note: GitHub repositoryAccessed: 2026-07-29 External Links: Link Cited by: §5. DeepSeek-AI (2024a) DeepSeek-AI DeepSeek-V2: a strong, economical, and efficient mixture-of-experts language model. arXiv preprint arXiv:2405.04434. External Links: Link Cited by: §1. DeepSeek-AI (2024b) DeepSeek-AI DeepSeek-V3 technical report. arXiv preprint arXiv:2412.19437. External Links: Link Cited by: §1. DeepSeek-AI (2025) DeepSeek-AI Expert parallelism load balancer (EPLB). Note: GitHub repositoryAccessed: 2026-07-25 External Links: Link Cited by: §1, §5. Fedus et al. (2022) W. Fedus, B. Zoph, and N. Shazeer Switch transformers: scaling to trillion parameter models with simple and efficient sparsity. Journal of Machine Learning Research 23 (120), p. 1–39. External Links: Link Cited by: §1. He et al. (2022) J. He, J. Zhai, T. Antunes, H. Wang, F. Luo, S. Shi, and Q. Li FasterMoE: modeling and optimizing training of large-scale dynamic pre-trained models. In Proceedings of the 27th ACM SIGPLAN Symposium on Principles and Practice of Parallel Programming, p. 120–134. External Links: Document Cited by: §5. Lee et al. (2024) W. Lee, J. Lee, J. Seo, and J. Sim InfiniGen: efficient generative inference of large language models with dynamic KV cache management. In 18th USENIX Symposium on Operating Systems Design and Implementation (OSDI 24), p. 155–172. Cited by: §3.2. Lepikhin et al. (2021) D. Lepikhin, H. Lee, Y. Xu, D. Chen, O. Firat, Y. Huang, M. Krikun, N. Shazeer, and Z. Chen GShard: scaling giant models with conditional computation and automatic sharding. In International Conference on Learning Representations, External Links: Link Cited by: §1. Li et al. (2023) J. Li, Y. Jiang, Y. Zhu, C. Wang, and H. Xu Accelerating distributed MoE training and inference with lina. In 2023 USENIX Annual Technical Conference (USENIX ATC 23), Boston, MA, p. 945–959. External Links: Link Cited by: §1, §5. Li et al. (2022) Y. Li, A. Phanishayee, D. Murray, J. Tarnawski, and N. S. Kim Harmony: overcoming the hurdles of GPU memory capacity to train massive DNN models on commodity servers. Proceedings of the VLDB Endowment 15 (11), p. 2747–2760. External Links: Document Cited by: §5. Liu et al. (2025) J. Liu, J. Su, X. Yao, et al. Muon is scalable for LLM training. arXiv preprint arXiv:2502.16982. External Links: Link Cited by: §4. Liu et al. (2023) Z. Liu, J. Wang, T. Dao, T. Zhou, B. Yuan, Z. Song, A. Shrivastava, C. Zhang, Y. Tian, C. Ré, and B. Chen Deja vu: contextual sparsity for efficient LLMs at inference time. In Proceedings of the 40th International Conference on Machine Learning, Proceedings of Machine Learning Research, Vol. 202, p. 22137–22176. Cited by: §3.2. Rajbhandari et al. (2022) S. Rajbhandari, C. Li, Z. Yao, M. Zhang, R. Y. Aminabadi, A. A. Awan, J. Rasley, and Y. He DeepSpeed-MoE: advancing mixture-of-experts inference and training to power next-generation AI scale. In Proceedings of the 39th International Conference on Machine Learning, Proceedings of Machine Learning Research, Vol. 162, p. 18332–18346. External Links: Link Cited by: §1. Shazeer et al. (2017) N. Shazeer, A. Mirhoseini, K. Maziarz, A. Davis, Q. V. Le, G. E. Hinton, and J. Dean Outrageously large neural networks: the sparsely-gated mixture-of-experts layer. In International Conference on Learning Representations, External Links: Link Cited by: §1. Wei et al. (2026) X. Wei, C. Jin, T. Dai, Y. Zhong, S. Yu, C. Yang, B. Wu, Z. Zhang, J. Mai, Q. Zhu, Z. Li, Y. Liu, and G. Luo UltraEP: unleash MoE training and inference on rack-scale nodes with near-optimal load balancing. arXiv preprint arXiv:2606.04101. External Links: Link Cited by: §5. Wu et al. (2025) Y. Wu, K. Gao, L. Li, and Y. Wu EasySpec: layer-parallel speculative decoding for efficient multi-GPU utilization. In Advances in Neural Information Processing Systems, Vol. 38. Cited by: §3.2. Yang et al. (2025) A. Yang, A. Li, B. Yang, et al. Qwen3 technical report. arXiv preprint arXiv:2505.09388. External Links: Link Cited by: §4.