International Conference on Learning Representations (ICLR 2026)

Capacity-Aware Inference:
Mitigating the Straggler Effect in Mixture of Experts

A plug-and-play inference-time dispatching framework that eliminates synchronization bottlenecks in distributed sparse MoE serving through deterministic capacity bounds and local candidate expansion.

Shwai He¹, Weilin Cai², Jiayi Huang², Ang Li¹

¹ CASE Lab, University of Maryland, College Park ² The Hong Kong University of Science and Technology (Guangzhou)

Read Paper (arXiv) OpenReview Code & Models Live Interactive Demo

1.85×

End-to-End Speedup

Mixtral-8x7B-Instruct

30%

MoE Layer Latency Drop

OLMoE-1B-7B

<0.2%

Accuracy Degradation

Lossless Quality Retention

Retraining Needed

Zero-shot Inference Patch

Live Interactive Experiment

Interactive Straggler & Load Balancing Visualizer

Adjust the MoE architecture, token batch size, and capacity factor $\gamma$ in real time. Compare how Vanilla Top-$k$ routing creates severe synchronization stragglers, while Token Drop and Expanded Drop balance multi-GPU workloads.

1. Model Architecture

8 total experts distributed across 4 GPU ranks.

2. Token Batch ($T$) 512 tokens

128 1024 2048

3. Capacity Factor ($\gamma$) 1.00×

0.50 (Aggressive) 1.00 (Balanced) 2.00 (Loose)

4. Routing Strategy

Avg Expert Load ($ar{N}$)

128 tokens

Theoretical ideal balance

Expert Capacity Bound ($C$)

128 tokens

$C = \lceil \gamma \cdot ar{N} ceil$

Straggler Bottleneck ($N_{\max}$)

264 tokens

Straggler causes 2.06× idle wait

Effective Speedup vs Baseline

1.00×

0% dropped tokens

Distributed Expert Load Profile

Normal Load Near Cap Straggler / Overflow Capacity Line

Capacity Limit C

Vanilla Top-k: Unconstrained routing allows hot experts to accumulate huge token backlogs, forcing all GPUs to synchronize on the slowest worker.

Synchronization Idle: 51.5%

The Fundamental Bottleneck

The Distributed MoE Straggler Effect

Learned gating functions route tokens unevenly based on natural linguistic affinities. Under multi-GPU expert parallelism, the whole cluster slows down to match the single most overloaded worker.

Figure 1: Token queue skew creates GPU idle cycles at the synchronization barrier.

Heavy-Tailed Routing Skew

In production MoE models (Mixtral, DeepSeek, OLMoE), 10–20% of the experts receive over 60% of all tokens, while the rest remain substantially underutilized.

Synchronization Barrier Tax

Under Expert Parallelism (EP), all devices must exchange activations via all-to-all communications before proceeding to the next layer, serializing execution on the slowest GPU.

Widespread Across Architectures

The straggler effect persists across language and multimodal MoEs alike, and exacerbates with larger batch sizes and higher expert counts ($E=64, 128$).

Empirical Profiling

Empirical Expert Load Imbalance Across Models

Select a model architecture to inspect its layer-wise token assignment distribution and straggler hotspots.

OLMoE-1B-7B Layer-Wise Distribution

Click to Zoom High-Res

Figure: Layer-wise token routing frequency across all available expert slots.

Root-Cause Straggler Analysis

1. Layer Skew Hotspots 4.2x Peak Skew

In intermediate and deep layers, router logits heavily polarize. Up to 78% of tokens route to only 15% of active experts.

2. Sync Tail Latency 45% Idle Stall

Under Expert Parallelism (EP), All-to-All communication blocks on the most saturated worker, stalling healthy GPUs.

3. Capacity Bounding 1.85x Speedup

Enforcing $\gamma \in [1.0, 1.2]$ caps the long tail and reallocates excess tokens with negligible perturbation to downstream accuracy.

Model Architecture

OLMoE-1B-7B-Instruct

64 Experts, Top-8 routing per token

Serving Speedup: +72% Token/s

Architectural Framework

Deterministic Inference-Time Capacity Bounds

We regulate expert workload using a capacity factor $\gamma$, bounding the maximum computation time per expert without retraining the underlying weights.

$$C = \left\lceil \gamma \cdot ar{N} ight ceil = \left\lceil \gamma \cdot rac{T \cdot k}{E} ight ceil$$

$T$: Token Batch | $k$: Active Experts | $E$: Total Experts | $\gamma$: Capacity Multiplier

1. Capacity-Aware Token Drop

Calculates capacity threshold $C$ per expert. Tokens routed to an expert are ranked by router probability score; overflow tokens beyond $C$ are cleanly dropped, ensuring no GPU exceeds runtime budget $C$.

2. Capacity-Aware Expanded Drop

Instead of discarding overflow tokens, the router dynamically reallocates them to alternative candidate experts situated on the same local GPU rank that have available capacity, maximizing parameter utilization.

Empirical Verification

Performance & Scalability Benchmarks

Extensive multi-GPU benchmarks demonstrating superior throughput, compressed P99 tail latency, and lossless accuracy across major LLM benchmarks.

Tested on 8× NVIDIA A100-80GB

Insight: Expanded Drop maintains steady ~1.85× throughput gain across large batch sizes as token skew increases.

Detailed Experimental Results Table

Language & Multimodal

Model Architecture	Experts ($E$)	Top-$k$	Strategy	Capacity ($\gamma$)	Avg Benchmark Acc	Layer Speedup	End-to-End Speedup
Mixtral-8x7B-Instruct	8	2	Baseline	∞	71.4%	1.00×	1.00×
Mixtral-8x7B-Instruct	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×
OLMoE-1B-7B-Instruct	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×
DeepSeek-V2-Lite-Chat	64	6	Expanded Drop	1.0	68.1% (-0.1%)	1.45×	1.39×

Developer Integration

Quick Start: Zero-Shot Integration

Patch any Hugging Face, vLLM, or DeepSpeed MoE checkpoint in 3 lines of Python without re-training or modifying model weights.

capacity_aware_inference.py

# 1. Import transformers and our lightweight patch engine
import torch
from types import SimpleNamespace
from transformers import AutoModelForCausalLM, AutoTokenizer
from capacity_aware import apply_capacity_aware_moe_patch

# 2. Load standard MoE checkpoint (e.g. Mixtral, DeepSeek, OLMoE)
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")

# 3. Configure Capacity-Aware parameters
config = SimpleNamespace(
    expert_capacity=1.0,       # Capacity factor gamma (C = gamma * N_avg)
    strategy="score",          # Drop heuristic: "score", "first", "last"
    rounds=1,                  # Expansion rounds for candidate reallocation
    capacity_scope="expert",   # "expert" or "device"
)

# 4. Patch the model in-place (Zero Retraining!)
num_patched = apply_capacity_aware_moe_patch(model, config)
print(f"Patched {num_patched} MoE layers with Capacity-Aware routing.")

# 5. Run accelerated inference!
inputs = tokenizer("Capacity-Aware inference solves MoE stragglers by", return_tensors="pt").to("cuda")
outputs = model.generate(**inputs, max_new_tokens=64)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))

# Launch standardized evaluation on Mixtral-8x7B
cd 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

# Launch multimodal evaluation on MMBench
cd VLMEvalKit

python run.py --data MMBench_DEV_EN --model DeepSeek-VL-7B --mode all

Citation

If you find Capacity-Aware Inference helpful in your research, please cite our ICLR 2026 paper:

BibTeX (ICLR 2026)

@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}
}