Shwai He1 • Weilin Cai2 • Jiayi Huang2 • Ang Li1
1CASE Lab, University of Maryland, College Park | 2The Hong Kong University of Science and Technology (Guangzhou)
🌐 Interactive Demo & Website • 📄 arXiv Paper • 💬 OpenReview • 🚀 Quick Start • 📝 Citation
- 🌟 Key Highlights
- 📰 News & Updates
- 🔍 Problem: The Distributed MoE Straggler Effect
- 💡 Core Methodology
- 📊 Benchmark Results
- 📁 Repository Structure
- ⚙️ Installation & Setup
- 🚀 Quick Start & Usage
- 🔬 Reproducing Paper Experiments
- 📝 Citation
- 🤝 Acknowledgements & Contact
- ⚡ Up to 1.85× End-to-End Speedup: Substantially accelerates sparse MoE generation during multi-GPU distributed serving.
- 🛡️ Zero Retraining Required: A plug-and-play inference-time dispatching mechanism compatible with existing pretrained weights.
-
🎯 Lossless Quality Preservation: Retains
$>99.8%$ benchmark accuracy on MMLU, GSM8K, ARC-Challenge, and HellaSwag. - 📉 Tail Latency Mitigation: Eliminates synchronization bottlenecks across Expert Parallel (EP) and Tensor Parallel (TP) ranks.
- 🌐 Broad Architecture Coverage: Out-of-the-box support for DeepSeek-V2/V3/Lite, Mixtral-8x7B/8x22B, OLMoE-1B-7B, and Qwen-MoE models.
- [Jan 2026] 🌟 Paper Accepted to ICLR 2026! Full paper and open-source implementation released.
- [Mar 2025] 📄 Preprint released on arXiv:2503.05066.
In sparse Mixture-of-Experts (MoE) architectures, each token dynamically selects a small subset of
Because learned gating functions naturally route hot topics to popular "specialist" experts, token distributions across experts are heavily skewed. Under standard unconstrained top-$k$ routing:
- The Overloaded Expert Bottleneck: A few hot experts receive disproportionately massive token queues.
- The Distributed Synchronization Barrier: All non-overloaded GPU ranks finish their execution early and sit idle, waiting at the barrier for the slowest expert to complete computation.
-
Severe Tail Latency Inflation: The step execution time is governed entirely by the maximum load:
$$\text{Step Latency} \propto \max_{e \in {1 \dots E}} \text{Load}(e)$$
Figure 1: Illustration of the MoE straggler effect in distributed expert-parallel inference. Overloaded Expert 1 forces GPU 2, 3, and 4 to wait at the global barrier.
To eradicate synchronization stragglers without requiring costly model retraining, we introduce Capacity-Aware Inference, which enforces deterministic per-expert load ceilings at inference time.
Incoming Tokens [T, D]
│
┌───────┴───────┐
│ Router Gate │ Compute router logits & affinity scores
└───────┬───────┘
│
┌───────▼────────────────────────────────────────┐
│ Routing Strategy Decision │
├───────────────────────┬────────────────────────┤
│ Token Drop (TD) │ Expanded Drop (ED) │
├───────────────────────┼────────────────────────┤
│ Enforce C = γ · N̄ │ Expand candidate pool │
│ Drop overflow tokens │ Redirect to low-load │
│ beyond capacity C │ local experts first │
└───────┬───────────────┴────────┬───────────────┘
│ │
└───────────┬────────────┘
│
[ Balanced Expert Dispatch ]
┌───────────┬───────────┐
▼ ▼ ▼
[GPU 0] [GPU 1] [GPU 2]
Expert 1 Expert 2 Expert 3
│ │ │
└───────────┼───────────┘
▼
Zero Barrier Synchronization Idle
- Computes expert capacity bound
$C$ . - Ranks candidate tokens per expert according to router probability scores (or sequential position).
- Retains the top-$C$ tokens for the expert and safely drops/masks lower-affinity overflow tokens.
- Bounds the maximum single-expert compute time to
$C$ .
- Instead of dropping overflow tokens immediately, the router expands candidate selection into local under-utilized experts on the same GPU rank.
- Tokens exceeding the capacity of their primary expert are gracefully redirected to candidate experts with spare capacity without triggering cross-node communication.
- Achieves superior load balancing, higher expert parameter utilization, and better quality preservation at aggressive speedup ratios.
We calculate the maximum token capacity
Where:
-
$T$ : Total batch token count in the current forward pass. -
$k$ : Number of activated experts per token (e.g.,$k=2$ for Mixtral,$k=8$ for OLMoE). -
$E$ : Total number of experts in the MoE layer. -
$\bar{N}$ : Average expected token load per expert ($\bar{N} = \frac{Tk}{E}$ ). -
$\gamma$ : Capacity Factor (expert_capacity, default range:$0.6 \le \gamma \le 1.2$ ).
Figure 2: Conceptual mechanics of Capacity-Aware Token Drop (left) vs. Capacity-Aware Expanded Drop (right).
| Model Architecture | Total Experts ( |
Active Experts ( |
Strategy | Capacity Factor ( |
Avg Benchmark Acc | MoE Layer Speedup | End-to-End Speedup |
|---|---|---|---|---|---|---|---|
| Mixtral-8x7B-Instruct | 8 | 2 | Baseline | 71.4% | 1.00× | 1.00× | |
| 8 | 2 | Token Drop | 1.0 | 71.1% | 1.48× | 1.34× | |
| 8 | 2 | Expanded Drop | 1.0 | 71.6% (+0.2%) | 1.92× | 1.85× | |
| OLMoE-1B-7B-Instruct | 64 | 8 | Baseline | 63.8% | 1.00× | 1.00× | |
| 64 | 8 | Token Drop | 0.8 | 63.2% | 1.24× | 1.18× | |
| 64 | 8 | Expanded Drop | 0.8 | 63.5% (-0.3%) | 1.32× | 1.28× | |
| DeepSeek-V2-Lite-Chat | 64 | 6 | Baseline | 68.2% | 1.00× | 1.00× | |
| 64 | 6 | Expanded Drop | 1.0 | 68.1% (-0.1%) | 1.45× | 1.39× | |
| Qwen-MoE-Chat | 32 | 4 | Baseline | 70.5% | 1.00× | 1.00× | |
| 32 | 4 | Expanded Drop | 0.9 | 70.4% (-0.1%) | 1.38× | 1.31× |
Distributed serving tests on 8× NVIDIA A100/H100 GPUs demonstrate that capacity bounds drastically compress P99 tail latencies caused by expert load variance:
Figure 3: Detailed latency decomposition on OLMoE. Capacity-aware dispatch significantly reduces expert computation, permutation, and inter-GPU communication waiting times.
Capacity-Aware Inference generalizes directly to multimodal sparse MoEs without architectural modifications:
Figure 4: Multimodal evaluation on MMBench comparing Baseline, Token Drop, and Expanded Drop under varying capacity constraints.
Capacity-Aware-MoE/
├── capacity_aware/ # Core runtime capacity-aware routing engine
│ ├── __init__.py # Package exports
│ └── capacity_patch.py # Generic PyTorch monkey-patch for HF/vLLM routers
├── scripts/ # Automated evaluation and cluster launcher scripts
│ └── run_qwen35_capacity_lmeval.sh # Comprehensive Qwen3.5-MoE SLURM/batch eval runner
├── lm-evaluation-harness/ # Standardized LLM evaluation harness integration
│ ├── lm_eval/capacity_aware/ # Harness re-exports
│ └── runs_prune/
│ ├── eval_baseline.sh # Baseline unconstrained evaluation script
│ └── eval_capacity.sh # Capacity-aware evaluation script
├── VLMEvalKit/ # Multimodal MoE evaluation suite (MMBench, etc.)
├── docs/ # Academic project page assets & visualizer
│ ├── index.html # Modern interactive project website
│ └── Figures/ # Architectural SVGs, benchmark charts, and plots
├── requirements.txt # Environment dependencies
└── README.md # Repository documentation
# Create a fresh conda environment
conda create -n capacity-moe python=3.10 -y
conda activate capacity-moe
# Clone the repository
git clone https://github.com/CASE-Lab-UMD/Capacity-Aware-MoE.git
cd Capacity-Aware-MoE
# Install root dependencies
pip install -r requirements.txt# Install language evaluation harness
cd lm-evaluation-harness
pip install -e .
cd ..
# (Optional) Install multimodal evaluation kit
cd VLMEvalKit
pip install -e .
cd ..You can apply Capacity-Aware Inference to any Hugging Face MoE model with a single function call:
import torch
from types import SimpleNamespace
from transformers import AutoModelForCausalLM, AutoTokenizer
from capacity_aware import apply_capacity_aware_moe_patch
# 1. Load your favorite MoE model
model_id = "mistralai/Mixtral-8x7B-Instruct-v0.1"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
model_id,
torch_dtype=torch.bfloat16,
device_map="auto"
)
# 2. Configure capacity-aware parameters
config = SimpleNamespace(
expert_capacity=0.85, # Capacity factor γ (0.5 - 2.0)
strategy="score", # Strategy: "score", "first", "last", "random", "overselect"
rounds=1, # Expansion rounds for candidate reallocation
capacity_scope="expert", # "expert" (per expert) or "device" (per GPU device)
capacity_devices=8, # Number of EP ranks if capacity_scope="device"
)
# 3. Patch the MoE routers in-place (Zero Retraining!)
num_patched = apply_capacity_aware_moe_patch(model, config)
print(f"Successfully patched {num_patched} MoE layers with Capacity-Aware routing.")
# 4. Run accelerated inference as usual
inputs = tokenizer("Capacity-Aware Inference in MoE models solves", return_tensors="pt").to("cuda")
with torch.no_grad():
outputs = model.generate(**inputs, max_new_tokens=64)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))cd lm-evaluation-harness
CUDA_VISIBLE_DEVICES=0,1,2,3 \
PRETRAINED="mistralai/Mixtral-8x7B-Instruct-v0.1" \
TASKS="mmlu,gsm8k,arc_challenge,hellaswag" \
BATCH_SIZE=8 \
bash runs_prune/eval_baseline.shcd lm-evaluation-harness
CUDA_VISIBLE_DEVICES=0,1,2,3 \
PRETRAINED="mistralai/Mixtral-8x7B-Instruct-v0.1" \
TASKS="mmlu,gsm8k,arc_challenge,hellaswag" \
EXPERT_CAPACITY=1.0 \
STRATEGY=score \
ROUNDS=1 \
BATCH_SIZE=8 \
bash runs_prune/eval_capacity.sh| Parameter | Default | Options / Range | Description |
|---|---|---|---|
PRETRAINED |
None |
HF repo / path | Checkpoint identifier |
EXPERT_CAPACITY |
1.0 |
0.4 - 2.0
|
Capacity multiplier |
STRATEGY |
score |
score, random, first, last, overselect
|
Token prioritization heuristic |
ROUNDS |
1 |
1, 2, 3
|
Candidate expansion rounds for Expanded Drop |
CAPACITY_SCOPE |
expert |
expert, device
|
Scope for capacity calculation |
BATCH_SIZE |
1 |
Integer | Micro-batch size per GPU rank |
TASKS |
hellaswag |
Benchmark names | Comma-separated tasks for lm-evaluation-harness |
cd VLMEvalKit
python run.py \
--data MMBench_DEV_EN \
--model DeepSeek-VL-7B \
--mode allTo reproduce the full suite of experiments reported in our ICLR 2026 paper:
# 1. Evaluate Qwen-3.5 / Qwen-MoE series
bash scripts/run_qwen35_capacity_lmeval.sh
# 2. Run capacity sweep for Mixtral-8x7B across gamma values [0.6, 0.8, 1.0, 1.2, 1.5]
for cap in 0.6 0.8 1.0 1.2 1.5; do
EXPERT_CAPACITY=$cap STRATEGY=score bash lm-evaluation-harness/runs_prune/eval_capacity.sh
doneAll evaluation logs, accuracies, and token distribution statistics will be automatically saved under results/ and summarized in your terminal.
If you find Capacity-Aware Inference, our codebase, or our paper helpful in your research, please cite our ICLR 2026 work:
@inproceedings{he2026capacityaware,
title={Capacity-Aware Inference: Mitigating the Straggler Effect in Mixture of Experts},
author={He, Shwai and Cai, Weilin and Huang, Jiayi and Li, Ang},
booktitle={International Conference on Learning Representations (ICLR)},
year={2026},
url={https://arxiv.org/abs/2503.05066}
}Or the preprint reference:
@article{he2025capacityawareinferencemitigatingstraggler,
title={Capacity-Aware Inference: Mitigating the Straggler Effect in Mixture of Experts},
author={Shwai He and Weilin Cai and Jiayi Huang and Ang Li},
journal={arXiv preprint arXiv:2503.05066},
year={2025},
url={https://arxiv.org/abs/2503.05066}
}This work was conducted at the CASE Lab at the University of Maryland, College Park, in collaboration with The Hong Kong University of Science and Technology (Guangzhou).
For questions, issues, or contributions, please feel free to open a GitHub Issue or reach out via email:
- Shwai He:
shwai@umd.edu - Ang Li:
angli@umd.edu