Venkata Pavan Kumar Miriyala, German Sviridov, Bingxu Chen, Haris Javaid | 2025-12 | https://doi.org/10.1145/3769695.3771675 Category: framework | Tags: moe, inference, load-balancing, expert-parallelism, optimization Read: 2026-04-16
Proposes a latency-optimal algorithm for expert replication and reallocation during distributed MoE inference that jointly minimizes load imbalance and data movement overhead, formulated first as ILP then solved via a polynomial-time heuristic, achieving up to 12.5% latency reduction and 2× more frequent load balancing than prior approaches.
Expert parallelism (EP) is the dominant strategy for distributing Mixture-of-Experts (MoE) model inference across multiple devices, where each device hosts a subset of experts. However, the dynamic nature of token routing means some experts receive far more tokens than others, creating workload imbalance — all devices must wait for the slowest one, wasting hardware utilization and increasing latency. Prior approaches either (a) add auxiliary loss functions during training to encourage uniform token distribution, which constrains model quality, or (b) dynamically replicate/reallocate experts across devices at inference time, which introduces high data movement overhead from shuffling expert weights.
This paper addresses the data movement overhead problem head-on. The authors formulate the expert replication and reallocation problem as an Integer Linear Programming (ILP) optimization that jointly minimizes two objectives: the load imbalance across devices AND the total data movement cost incurred during the rebalancing phase. This joint optimization is the key insight — prior work treated these objectives separately, leading to suboptimal tradeoffs where rebalancing itself became a bottleneck. The ILP formulation establishes the theoretical optimum but is too expensive to solve at runtime. The authors then design a lightweight heuristic algorithm that solves the problem in polynomial time, making it practical for online use during inference serving.
Experimental results demonstrate up to 12.5% reduction in MoE execution latency over naive expert assignment. Critically, because the heuristic is fast enough, it enables load balancing to be performed 2× more frequently than existing approaches, further improving responsiveness to shifting workload patterns. The method is designed for scale-up network environments (e.g., within a single node or tightly-coupled accelerator cluster). During the workshop Q&A, the authors noted it could be adapted to heterogeneous environments via a normalization step.
What it shows: Illustration of how MoE expert parallelism distributes experts across devices, and how uneven token routing creates load imbalance where the slowest device becomes the bottleneck.
Why it matters: Establishes the core problem — in EP, each device processes tokens for its assigned experts, but dynamic routing means some devices get overwhelmed while others are idle.
Detailed description: Diagram showing N devices, each hosting a subset of E experts. Input tokens flow through a router that assigns each token to top-k experts. The resulting distribution is skewed — some devices receive many more tokens than others. The execution time is determined by the max-loaded device (the "straggler"), leaving other devices idle. This is the fundamental waste that motivates the paper.
What it shows: The optimization framework that jointly minimizes load imbalance and data movement cost during expert reallocation.
Why it matters: This is the core technical contribution — prior work optimized these objectives independently, leading to suboptimal solutions.
Detailed description: The ILP has decision variables for expert-to-device assignments. The objective function combines two terms: (1) a load balance term minimizing the maximum workload across devices (min-max formulation), and (2) a data movement cost term penalizing expert weight transfers between devices. Constraints ensure each expert is assigned to at least one device and capacity limits per device are respected. The joint formulation means a rebalancing action is only taken if its benefit (reduced imbalance) exceeds its cost (data movement).
What it shows: Comparison of the polynomial-time heuristic against the optimal ILP solution across different workload distributions.
Why it matters: Demonstrates that the heuristic achieves near-optimal results while being orders of magnitude faster — the key enabler for practical deployment.
Detailed description: Performance comparison showing the heuristic's latency closely tracks the ILP optimum across varying load distributions, while the solve time drops from ILP's exponential worst-case to the heuristic's polynomial time. This speed difference is what enables the 2× more frequent rebalancing.
| Method | MoE Exec Latency | Rebalancing Frequency | Data Movement Overhead |
|---|---|---|---|
| Naive assignment | baseline | N/A | N/A |
| Prior dynamic rebalancing | improved | 1× | high |
| This work (heuristic) | up to -12.5% | 2× | minimized |
Takeaway: The joint optimization achieves the best latency reduction while enabling more frequent rebalancing because the algorithm itself is fast and its decisions minimize data movement.
The system inserts a load balancing layer between the MoE router and device-level execution:
[Token Batch] → [MoE Router (top-k)] → [Token-to-Expert Assignment]
→ [Workload Monitor: observe per-expert load]
→ [Load Balancer: ILP/Heuristic decides expert placement]
→ [Expert Migration Engine: move/replicate expert weights]
→ [Device Execution: each device processes its assigned tokens]
→ [All-to-All gather results] → [Output]
| Stage | Input → Output | Location | Latency | Data |
|---|---|---|---|---|
| Token routing | tokens → per-expert token lists | CPU/GPU | ~μs | routing decisions |
| Workload observation | token counts → load profile | CPU | ~μs | E integers |
| Load balancing decision | load profile + current assignment → new assignment | CPU | ~ms (heuristic) | assignment matrix E×D |
| Expert migration | expert weights transferred | GPU↔GPU via interconnect | ~ms | expert param size (tens of MB each) |
| MoE computation | tokens × expert weights → outputs | GPU HBM | main cost | depends on model |
The paper's core insight is that hotspot #1 (expert migration) was being ignored in the optimization objective of prior work. By including migration cost in the objective, the algorithm avoids excessive weight transfers.
| Innovation | Mechanism | Benefit | Cost/Tradeoff |
|---|---|---|---|
| Joint ILP formulation | Objective = $\alpha \times \text{max\_load} + \beta \times \text{migration\_cost}$ | Avoids rebalancing actions where migration cost exceeds balance benefit | Requires tuning $\alpha/\beta$; ILP is NP-hard |
| Polynomial-time heuristic | Greedy descent: find straggler → evaluate best expert move → execute if net positive | Runtime-practical (ms-scale), enables 2× more frequent rebalancing | May miss global optimum; greedy can get stuck in local optima |
| Expert replication | Replicate hot experts to multiple devices instead of just moving | Reduces straggler load without removing expert from source | Increases total memory usage; requires coherent handling |
| Min-max load formulation | Minimize maximum device load (not average or variance) | Directly targets the straggler — the actual bottleneck in synchronous EP | May leave non-straggler devices suboptimally loaded |
Core scheduling algorithm (heuristic):
Input: current_assignment, workload_profile, device_set
while improvement_possible:
d_max = argmax_d load(d)
e_hot = argmax_e tokens(e, d_max)
best_target = None, best_gain = 0
for d_cand in device_set \ {d_max}:
if can_fit(e_hot, d_cand):
gain = balance_improvement(move e_hot → d_cand)
cost = migration_cost(e_hot, d_max → d_cand)
net = α * gain - β * cost
if net > best_gain:
best_target = d_cand; best_gain = net
if best_target:
migrate(e_hot, d_max → best_target)
else:
break
Complexity: $O(E \times D \times K)$ per rebalancing — polynomial, practical for online use.
| Scenario | Workload Pattern | SLO / Goal | Why existing systems fail |
|---|---|---|---|
| MoE inference on scale-up cluster | Dynamic token routing creates load imbalance across EP devices | Minimize per-request latency | Naive assignment creates stragglers; dynamic rebalancing has excessive migration overhead |
| High-frequency load shifting | Expert popularity changes across batches | Responsive rebalancing | Prior optimization approaches too slow to run frequently |
Primary bottleneck: Communication-bound — the data movement cost of expert weight migration is the key bottleneck that prior approaches failed to optimize. After optimization, the remaining bottleneck shifts to compute (the actual expert computation, which is now better balanced).
| Metric | Definition | Unit | Direction |
|---|---|---|---|
| MoE execution latency | Time for the MoE layer forward pass under EP | ms | Lower is better |
| Load imbalance | Max device load / average device load | ratio | Lower is better |
| Rebalancing frequency | How often load balancing can be triggered | per N batches | Higher is better |
| Data movement overhead | Total bytes transferred during expert migration | GB | Lower is better |
| Optimization | Metric | Baseline | After | Improvement | Conditions |
|---|---|---|---|---|---|
| Joint load balance + migration cost | MoE exec latency | naive assignment baseline | up to -12.5% | 12.5% reduction | Scale-up, EP setting |
| Fast heuristic vs prior approach | Rebalancing frequency | 1× (prior) | 2× | 2× more frequent | Same overhead budget |
| ILP optimality bound | Heuristic gap | ILP optimal | near-optimal | small gap | Experimental verification |
Before: scheduling-bound (straggler from load imbalance)
→ After joint optimization: communication-bound → compute-bound
(migration cost is minimized, remaining bottleneck is actual expert computation)
N/A — This is an algorithm/technique paper, not a full system. No API, deployment, or configuration details are provided. Integration into existing frameworks (vLLM, SGLang, TRT-LLM) would require adaptation.
| Layer | Impact |
|---|---|
| Algorithm | The ILP formulation provides a principled optimization framework that could inform training-time auxiliary loss design for MoE routing |
| Kernel | No new kernels needed — operates above kernel level. Expert computation kernels unchanged |
| LLM | Model-agnostic; benefits any MoE architecture (DeepSeek-V3, Mixtral, etc.) without modification |
| Agent | Indirectly reduces latency for agent systems using MoE models |
| Ops | The rebalancing frequency and α/β tradeoff are operationally relevant knobs for SLO tuning |
| Feature | This Work | ExFlow | Libra | CRAFT | FineMoE/LPLB |
|---|---|---|---|---|---|
| Scope | Inference | Inference | Inference | Inference | Training |
| Optimization | ILP + heuristic | Dynamic reallocation | Prediction-based | Cost-aware replication | LP-based scheduling |
| Considers migration cost | Yes (joint) | No | No | Yes (different formulation) | N/A |
| Scale | Scale-up | Multi-node | Multi-node | Multi-node | Multi-node |
| Rebalancing speed | Fast (polynomial) | Medium | Fast (prediction) | Medium | Medium |
| Latency reduction | 12.5% | [varies] | up to 19.2% throughput | 1.14× throughput | up to 47.6% throughput |