Libra: Effective yet Efficient Load Balancing for Large-Scale MoE Inference

Official implementation of Libra (ICLR 2026), built on SGLang v0.4.10.

Libra is a dynamic load-balancing system for distributed Mixture-of-Experts (MoE) inference under expert parallelism (EP). It predicts the next layer's expert activations, plans locality-aware expert replication from that prediction, and shards the tokens of hot experts across their replicas. It restructures the execution of each MoE layer so that these load-balancing costs are hidden under MoE computation. On 8× H200 GPUs with Qwen3-235B-A22B and GLM-4.5, Libra improves prefill throughput by up to 19.2%.

Table of Contents

• Key Ideas
• Repository Layout
• Requirements
• Installation
• Configuration
• Running a Single Experiment
• Reproducing the Paper
• License
• Citation

Key Ideas

• Two-Stage Locality-Aware Execution. Each MoE layer is split into two phases:
  • MoE_local processes tokens routed to experts that reside on (or were replicated to) the token's own GPU.
  • MoE_remote processes all other tokens after they are dispatched to other GPUs.
  • MoE_local starts right after gating, which hides the token dispatch (AllGather) and the CPU-side token sharding.
• Lookahead predictor. The gate of layer i+1 is applied to the hidden states of layer i, which evolve slowly across layers, to predict the experts that layer i+1 will activate.
• Locality-aware expert replication planning. It runs in two phases from that prediction, giving each GPU up to N extra experts:
  • Phase 1: each GPU imports the L = N·α non-resident experts most requested by its own tokens.
  • Phase 2: the hottest expert of the most-loaded GPU is replicated to the least-loaded eligible GPU, repeatedly.
  • Replicas are copied GPU-to-GPU with PyTorch SymmetricMemory into even/odd double buffers while the previous layer computes.
• Adaptive token sharding. Remote tokens are iteratively moved from overloaded GPUs to the least-loaded GPUs holding a replica of their expert. This runs on the CPU (Cython) in parallel with MoE_local.

Repository Layout

This repository holds the documentation and two git submodules:

Path	Branch	Contents
Libra-Core/	iclr2026_libra	SGLang v0.4.10 + Libra. Prefill-throughput runs of SGLang, EPLB and Libra, and the EPLB expert-placement generation.
Libra-Core/	iclr2026_lina	SGLang v0.4.10 + our re-implementation of Lina (Li et al., USENIX ATC 2023), used as a baseline.
Libra-Internal/	iclr2026_libra	SGLang v0.4.10 patched to record per-token routing (Token-to-Expert Mapping Tables, TEMT). It also builds Lina's expert-selection-path tables and simulates imbalance ratios offline.

Where the Libra implementation lives in Libra-Core (iclr2026_libra):

Path	What it implements
python/sglang/srt/models/qwen3_moe.py, glm4_moe.py	Two-stage execution, lookahead prediction and expert replication in the decoder layers.
python/sglang/srt/models/qwen2_moe.py, deepseek_v2.py	The model-level layer loop (Qwen3MoeModel inherits Qwen2MoeModel) and the MoE block used by GLM-4.5.
python/sglang/load_balance/	Expert replication planning (prefetch) and adaptive token sharding (rebalance), in Cython.
python/sglang/srt/layers/moe/	FusedMoE support for masked tokens and double-buffered expert weights, and the combine kernel.
python/sglang/srt/utils.py	SymmetricMemory buffers, P2P replication plans, routing exchange, and the scheme switch get_libra_scheme().
python/sglang/bench_one_batch_test.py	The benchmark driver used by every experiment.

Requirements

• A single node with 8 NVIDIA Hopper GPUs connected by NVLink/NVSwitch. All results were obtained on 8× H200 (141 GB); GLM-4.5 in BF16 does not fit on 8× 80 GB GPUs.
• Python 3.10 and CUDA 12.8 (tested).
• Disk space for the model weights (Qwen3-235B-A22B and GLM-4.5) and for the TEMTs used by the baselines and the imbalance-ratio simulation (up to about 0.6 TB).
• Building Lina's expert-selection-path tables needs a host with several hundred GB of RAM.

The exact hardware and software configuration is listed in REPRODUCE.md.

Installation

Clone the repository with its submodules:

git clone --recursive https://github.com/SNU-ARC/Libra.git
cd Libra

Create two environments: one for Libra-Core (shared by both of its branches) and one for Libra-Internal.

# Libra-Core: SGLang, EPLB, Libra, and Lina (same dependencies on both branches)
conda create -n libra python=3.10 -y
conda activate libra
cd Libra-Core
git checkout iclr2026_libra
pip install -e "python[all]"   # also builds the Cython load balancer in python/sglang/load_balance
cd ..

# Libra-Internal: TEMT generation and imbalance-ratio simulation
conda create -n libra_internal python=3.10 -y
conda activate libra_internal
cd Libra-Internal
pip install -e "python[all]"
pip install cython
cd prefetch_rebalance && python setup.py build_ext --inplace && cd ..   # Cython planner used by the simulator
cd ..

Libra-Core is installed in editable mode, so git checkout iclr2026_lina inside Libra-Core switches the installed code to the Lina baseline without reinstalling.

Configuration

All paths are configured with environment variables; every variable is optional.

Variable	Default	Meaning
LIBRA_MODEL_DIR	$HOME	Directory containing the model weights Qwen3-235B-A22B/ and GLM-4.5/.
LIBRA_DATA_DIR	$HOME/hf_datasets	Local copies of the eight evaluation datasets (layout in REPRODUCE.md).
LIBRA_ARTIFACT_DIR	$HOME/libra_artifacts	Artifacts shared between Libra-Internal and Libra-Core: TEMTs, Lina's path tables, and the simulated imbalance ratios.
CUDA_VISIBLE_DEVICES	0,1,2,3,4,5,6,7	GPUs used by the launcher scripts.

The load-balancing scheme is selected with LIBRA_SCHEME (sglang, eplb, libra, or lina). The launcher scripts set it for you.

Running a Single Experiment

Each experiment is one prefill-throughput run on 8 GPUs (DP attention + EP MoE, TP = DP = EP = 8). To run Libra once, from Libra-Core on branch iclr2026_libra:

bash single_node_scripts/ep_test.sh Qwen3-235B-A22B 1024 4 EP bookcorpus train 0 1 400 496 8 4 4096 libra

The arguments are, in order:

Argument	Example	Meaning
MODEL	Qwen3-235B-A22B	Qwen3-235B-A22B or GLM-4.5, loaded from $LIBRA_MODEL_DIR/MODEL.
INPUT_LEN	1024	Tokens per request.
BS	4	Requests per GPU (DP rank) per iteration.
EP	EP	Label used only for the profiler output directory.
DATASET	bookcorpus	One of bookcorpus, codeforces, deepseek_prover, fineweb, gsm8k, hellaswag, humanevalplus, lmsys_chat_1m.
MODE	train	Hugging Face split to read.
START_PORTION, END_PORTION	0 1	Percent slice of the split (train[0%:1%]).
START_IDX, END_IDX	400 496	Range of token batches to evaluate; each iteration consumes 8 batches, one per GPU.
N	8	Maximum number of extra (replicated) experts per GPU.
L	4	Experts imported per GPU by locality in Phase 1 (L = N·α).
SEQ_LENS_SUM	4096	Tokens per GPU per iteration; must equal INPUT_LEN × BS.
SCHEME	libra	sglang or libra.

The run appends one JSON line per iteration to latency_throughput/${MODEL}_${SCHEME}_${DATASET}_${SEQ_LENS_SUM}.txt. Each line includes prefill_throughput, the prefill throughput in tokens/s measured on GPU 0.

The other launchers follow the same argument order:

• EPLB: single_node_scripts/eplb_test.sh, branch iclr2026_libra. It takes the number of redundant experts per layer in place of N L, and has no SCHEME argument.
• Lina: single_node_scripts/ep_test_{qwen3,glm45}.sh, branch iclr2026_lina. It takes N only, and has no L or SCHEME argument.

Reproducing the Paper

REPRODUCE.md gives step-by-step instructions for reproducing Figure 8 (prefill throughput of SGLang, EPLB, Lina and Libra) and Figure 9 (throughput and imbalance-ratio fluctuation on a dynamic workload).

License

Libra is released under the Apache License 2.0 (see Libra-Core/LICENSE and Libra-Internal/LICENSE). It is built on SGLang v0.4.10; files modified from SGLang carry a notice stating so.

Citation

@inproceedings{yang2026libra,
  title     = {Libra: Effective yet Efficient Load Balancing for Large-Scale {MoE} Inference},
  author    = {Yang, Jaehoon and Kim, Yushin and Moon, Seokwon and Park, Yeonhong and Lee, Jae W.},
  booktitle = {The Fourteenth International Conference on Learning Representations},
  year      = {2026}
}