moe-router-sim

Simulates and compares Mixture-of-Experts routing strategies, measuring load imbalance, token drop rate, expert collapse, and communication overhead across 432 configurations.

For detailed architecture, algorithm descriptions, and design rationale see DESIGN.md.

---

Portfolio Context

Part of a seven-project series modeling the full LLM inference stack:

Project	Focus	Key Finding
kv-cache-compaction-lab	KV-cache compaction	ThresholdCompaction dominates; 11 free-compaction points
prefix-cache-sim	Prefix sharing	LFU dominates under Zipf; multi-turn hit rate 60%+
llm-inference-scheduler	Continuous batching	ChunkedPrefill eliminates starvation
tensor-memory-allocator	GPU tensor allocation	Free-list beats buddy/slab for continuous sizes
llm-serving-sim	End-to-end serving	ChunkedPrefill + LFU: 41% lower TTFT, 94% hit rate
speculative-decoding-sim	Speculative decoding	6.06x max speedup; breakeven at cost_ratio = 0.25
moe-router-sim	MoE routing and load balancing	ExpertChoice best balance; NoisyTopK best practical tradeoff

---

Problem

The MoE router determines which experts process which tokens. A poor routing strategy causes:

• Load imbalance: overloaded experts stall the pipeline
• Token drops: excess tokens are discarded, wasting compute
• Expert collapse: the router always uses the same few experts
• Communication overhead: in multi-GPU setups, all-to-all transfers dominate

---

Routing Strategies

Strategy	Complexity	Balance	Drop Rate	Notes
TopK	O(N log K)	Poor	High	Collapses under skewed distributions
NoisyTopK	O(N log K)	Better	Medium	Noise reduces collapse by 40%
ExpertChoice	O(N * cap)	Best	Low	Each expert selects its top tokens
Hash	O(K)	Perfect	Near zero	No specialization; deterministic

---

Key Findings

1. TopK collapses under realistic distributions

Under skewed and Zipf token distributions:

TopK load_imbalance = 5.52 (mean across all configs)
TopK collapse_events = 5-7 per 20 batches
TopK drop_rate > 1.0 (both K assignments can drop)

2. ExpertChoice achieves the best practical balance

ExpertChoice load_imbalance = 1.40 (vs TopK = 5.52)
ExpertChoice drop_rate = 0.198 (vs TopK = 1.211)
ExpertChoice collapse_events = 0.37 (vs TopK = 3.89)

3. NoisyTopK reduces drops by 40% vs TopK on skewed distributions

TopK drop rate (skewed):      1.244
NoisyTopK drop rate (skewed): 0.741

Adding noise forces exploration and prevents collapse, at no additional algorithmic complexity cost.

4. Capacity factor is the strongest lever

ExpertChoice drop rate by capacity_factor (skewed dist):

capacity_factor = 1: 0.419
capacity_factor = 2: 0.170
capacity_factor = 4: 0.003

Larger capacity directly trades memory for near-zero drops.

5. There is no universally best router

Hash:          perfect balance, zero specialization
ExpertChoice:  best balance with content-aware routing
NoisyTopK:     best practical tradeoff for skewed real-world workloads
TopK:          collapses under realistic distributions; avoid in production

---

Quick Start

# Build
cmake -S . -B build -G Ninja
cmake --build build -j

# Single run
./build/moe_router_sim \
    --router expert_choice \
    --n-experts 8 \
    --top-k 2 \
    --capacity-factor 2 \
    --token-dist skewed \
    --n-tokens 10000

# Full sweep (432 runs)
python3 experiments/sweep_moe.py

# Plots (6 plots)
python3 plots/plot_moe.py

# Analysis
python3 scripts/analyze_moe.py

---

CLI Reference

./build/moe_router_sim [options]

  --router STR          top_k | noisy_top_k | expert_choice | hash
  --n-experts N         total number of experts (default: 8)
  --top-k N             experts activated per token (default: 2)
  --capacity-factor N   capacity = factor * tokens/n_experts (default: 2)
  --n-gpus N            number of GPUs for comm modeling (default: 4)
  --comm-cost-us F      cross-GPU transfer cost per token us (default: 50)
  --noise-std F         Gaussian noise std for noisy_top_k (default: 1.0)
  --token-dist STR      uniform | zipf | skewed (default: uniform)
  --zipf-alpha F        Zipf exponent (default: 1.2)
  --skew-hot-frac F     fraction of hot experts in skewed (default: 0.1)
  --skew-hot-prob F     probability mass on hot experts (default: 0.8)
  --n-tokens N          total tokens to process (default: 10000)
  --batch-size N        tokens per batch (default: 512)
  --seed N              random seed (default: 42)
  --summary-out FILE    append summary CSV

---

Project Structure

moe-router-sim/
|-- include/
|   |-- config.hpp          MoEConfig
|   |-- token.hpp           Token, ExpertAssignment, DispatchResult
|   |-- router.hpp          IRouter + TopK, NoisyTopK, ExpertChoice, Hash
|   |-- expert.hpp          ExpertPool (capacity enforcement)
|   |-- dispatcher.hpp      LocalDispatcher, MultiGPUDispatcher
|   |-- load_balancer.hpp   LoadBalanceStats, LoadBalancer
|   |-- metrics.hpp         BatchRecord, MoEMetrics, Collector
|   +-- simulator.hpp       MoESimulator
|-- src/
|   |-- router.cpp
|   |-- expert.cpp
|   |-- dispatcher.cpp
|   |-- load_balancer.cpp
|   |-- metrics.cpp
|   |-- simulator.cpp
|   +-- main.cpp
|-- experiments/
|   +-- sweep_moe.py        432-run sweep
|-- plots/
|   +-- plot_moe.py         6 plots
|-- scripts/
|   +-- analyze_moe.py      Analysis + key findings
|-- results/
|   |-- sweep_summary.csv   432 rows
|   +-- plots/
|       |-- 01_imbalance_vs_n_experts.png
|       |-- 02_drop_rate_vs_capacity.png
|       |-- 03_effective_throughput_by_dist.png
|       |-- 04_collapse_heatmap.png
|       |-- 05_pareto_imbalance_vs_drop.png
|       +-- 06_cross_gpu_traffic.png
|-- DESIGN.md
|-- LICENSE
+-- README.md