Toy Chiplet Simulator for MoE

Multi-chiplet GPU simulator for MoE (Mixture-of-Experts) model execution. Simulates cycle-accurate inference on a 2D mesh of chiplets with configurable expert placement, allocation strategies, and cache/execution policies.

Features

• Topology: 2D mesh (or ring) interconnect; configurable y_chiplets x x_chiplets.
• Models: Predefined configs for DeepSeek, Qwen, Kimi, and Llama-4 (layer count, experts, hidden dims, etc.).
• Allocation: NEAREST, EVEN, EXP_EVEN, RANDOM, OURS (load-balanced).
• Execution: Baseline cache-local / cache-all, or prediction-based (next-token / next-layer).
• Trace-driven: JSON trace files (per-token expert selection) drive the simulation.

Layout

toy_chiplet_sim/
├── main.py              # Entry: run single trace or batch benchmark
├── main_ae.py           # One-click Artifact Evaluation script
├── draw_e2e_ae.py       # AE plotting (handles partial data)
├── data_loader.py       # Trace file loading utilities
├── experiment_config.py # Experiment configuration
├── simulator_runner.py  # Simulation runner helpers
├── test_dataset.py      # Dataset / trace list helpers
├── ep_balancer.py       # Optional: expert load balancing (torch)
├── toy_chiplet_sim/     # Main package
│   ├── config.py        # SystemConfig, TimingConfig
│   ├── models/          # Addresses, Expert, Expert_slice
│   ├── hardware/        # DRAM, Cache, Chiplet, Interconnect
│   └── simulation/      # Events, ResourceManager, Simulator
├── results/             # Simulation outputs (gitignored)
├── figures_ae/          # AE-generated figures
├── moe_profiling_trace/ # Downloaded trace files (auto-created by AE)
├── requirements.txt
└── README.md

---

Artifact Evaluation (AE)

Dependencies

Hardware

• CPU: Any modern x86-64 processor (simulation is single-threaded, no GPU required)
• RAM: >= 64 GB for single model (batch size 16384 caches ~50 GB of trace data in memory); >= 200 GB if running all 4 models with --parallel
• Disk: >= 80 GB free space for 1 model; >= 300 GB for all 4 models (trace files are large)

Software

• OS: Linux (tested on Ubuntu 22.04)
• Python: >= 3.10

Python Package	Min Version	Purpose
numpy	>= 1.20	Core simulation computation
matplotlib	>= 3.5	Figure generation
pandas	>= 1.3	CSV result processing
huggingface_hub	latest	Trace download (auto-installed only if traces not present locally)

All Python dependencies are auto-installed by main_ae.py if missing.

Data

• MoE expert selection traces: Profiled token-level expert routing decisions on MMLU / MMLU_ZH_CN benchmarks.
• Source: HuggingFace Dataset
• Models covered:

  Model	Parameters	HuggingFace Subdirectory
  DeepSeek-R1-AWQ	671B (4-bit)	cognitivecomputations/DeepSeek-R1-AWQ
  Qwen3-235B-A22B-FP8	235B (FP8)	Qwen/Qwen3-235B-A22B-FP8
  Kimi-K2-Thinking	1T	moonshotai/Kimi-K2-Thinking
  Llama-4-Maverick	17B (128E)	meta-llama/Llama-4-Maverick-17B-128E-Instruct

• Format: JSON files, one per query, containing per-layer expert selections.
• The AE script auto-downloads required traces if not found locally.

Quick Start (Recommended)

Run the default AE configuration (Qwen model, 8x3 chiplet, batch 4096/8192/16384):

pip install -r requirements.txt
python main_ae.py

This will:

1. Check for trace files (download from HuggingFace if missing)
2. Backup existing results/ to results_bp/ and run simulations
3. Generate figures in figures_ae/

Estimated time: 8-12 hours for the default configuration (1 model, 1 architecture, 3 batch sizes, 5 strategies).

Full Evaluation

Run all 4 models on both chiplet architectures:

python main_ae.py --models all --archs "8,3;5,5"

Run all 4 models in parallel (one process per model, significantly faster):

python main_ae.py --models all --archs "8,3;5,5" --parallel

Per-model logs are saved to ae_{model}.log. All models finish before plotting.

Estimated time: ~96 hours sequentially; 18-36 hours with --parallel (4 models x 2 architectures x 3 batch sizes x 5 strategies = 120 configurations).

Custom Configurations

# Specific models
python main_ae.py --models qwen,deepseek

# Specific architecture
python main_ae.py --archs "5,5"

# Both architectures
python main_ae.py --archs "8,3;5,5"

# Custom batch sizes
python main_ae.py --batches 4096,8192

# Parallel execution (one process per model)
python main_ae.py --models all --parallel

# Skip steps (useful for re-plotting existing results)
python main_ae.py --skip-download --skip-sim     # only re-generate figures
python main_ae.py --skip-plot                     # only run simulation

AE Output

• results/: Per-model CSV files (qwen_results.csv, etc.) with 20 columns including throughput, hop count, DRAM access stats, and load balance metrics.
• figures_ae/e2e_stacked.png: Stacked bar chart (4 rows x 4 cols) showing throughput speedup and hop reduction across models, batch sizes, and die shapes. Missing models are left blank.
• figures_ae/e2e_summary.csv: Aggregate statistics grouped by model, batch size, and die shape.

Expected Results

The AE reproduces the end-to-end simulation results from the paper. Key claims:

• The proposed allocation strategy (Allo Only / Allo+Pred) achieves higher throughput and lower inter-chiplet hops compared to baseline and standard EP placement.
• The improvement is consistent across different MoE models, batch sizes, and chiplet topologies.

---

Developer Usage

Install / Run

pip install -r requirements.txt
python main.py              # run with current experiment_config.py settings
python main.py qwen         # run single model

API

from toy_chiplet_sim import SystemConfig, TimingConfig, Simulator

config = SystemConfig(
    y_chiplets=3,
    x_chiplets=3,
    allocation_strategy=SystemConfig.AllocationStrategy.NEAREST,
    exe_strategy=SystemConfig.ExeStrategy.PRED_NEXT_TOKEN,
    model_name="deepseek",
)
sim = Simulator(config)
sim.process_trace(["path/to/trace0.json", "path/to/trace1.json"])
sim.process_records()

print("Total time (ms):", sim.current_time / 1e6)
print("Local DRAM hit rate:", sim.stats["local_dram_hit"] / sim.stats["total_access"])
print("Chiplet usage:", sim.avg_chiplet_usage)

Trace Format

JSON: list of iterations; each iteration is a dict keyed by layer id ("0", "1", ...). Value is either null (no MoE) or a 2D list of selected expert ids per token (prefill: multiple tokens; decode: one token per list).

License

Use as needed for research or open-source projects.