HDA-MoE - IEEE TCAD accepted

IEEE TCAD Accepted HD-MoE at ICCAD 2025 Python 3.10 MIT License

Hybrid Parallelism and Dynamic, Adaptive Scheduling for Mixture-of-Experts with 3D Near-Memory Processing

Haochen Huang, Shuzhang Zhong, Shengxuan Qiu, Zhe Zhang, Shuangchen Li, Cong Li,
Dimin Niu, Hongzhong Zheng, Guangyu Sun, and Meng Li

🎉 Accepted by IEEE Transactions on Computer-Aided Design of Integrated Circuits and Systems (TCAD).

HDA-MoE is a deployment and runtime framework for efficient Mixture-of-Experts (MoE) inference on distributed 3D near-memory processing (NMP) systems. It co-optimizes expert placement, on-chip communication, runtime scheduling, and hardware-aware routing so that sparse MoE execution matches the compute, bandwidth, topology, and memory constraints of the target system.

This repository contains the implementation and evaluation workflow for the accepted TCAD paper and builds on the conference work HD-MoE (ICCAD 2025 paper).

• End-to-end speedup: 1.1x-3.4x over tensor parallelism, 1.1x-1.5x over expert parallelism, 1.1x-3.7x over compute-balanced hybrid TP-EP, and 1.1x-1.3x over HD-MoE.
• Offline system-aware mapping: jointly balances node computation, communication volume, link contention, and per-node expert-weight capacity.
• Online adaptation: combines hotspot-aware pre-broadcast scheduling with hardware-aware gating for time-varying expert activation.
• Broad evaluation: covers Mixtral, DeepSeek, Qwen2, and Qwen3.5 across Mesh, Torus, and Fat-tree interconnects.
• Routing fidelity: hardware-aware gating retains comparable task accuracy; non-Top-1 boundary substitutions reach 0.949 average output similarity across the four evaluated models.

1. Node Balance formulates hybrid TP-EP expert placement as a capacity-aware linear program to reduce computation imbalance and communication volume.
2. Link Balance maps logical placements to physical nodes with topology-aware Bayesian optimization, reducing link congestion on Mesh, Torus, and Fat-tree networks.
3. Dynamic Scheduling predicts short-term expert hotspots and pre-broadcasts high-priority experts without adding token communication.
4. Hardware-aware Gating regularizes low-impact routing choices with marginal computation and communication costs while preserving the original routing weights.

Overview of the HDA-MoE framework

HDA-MoE first generates an offline hybrid deployment that balances logical work and physical network traffic. At runtime, dynamic scheduling increases the available hardware supply for emerging hotspots, while hardware-aware gating reduces routing demand that would otherwise create new compute or communication bottlenecks.

HDA-MoE end-to-end speedup ranges

The figure summarizes the end-to-end time-between-token speedup ranges reported in the accepted paper across evaluated models, hardware configurations, mesh sizes, and interconnect topologies. Hardware-aware routing maintains comparable accuracy to the original gating policy.

git clone --recursive https://github.com/angerybob/HDA-MoE.git
cd HDA-MoE
conda create -n hda-moe python=3.10 -y
conda activate hda-moe
# Select the PyTorch build that matches your CUDA environment.
pip install torch==2.6.0 torchvision==0.21.0 torchaudio==2.6.0
pip install -r requirements.txt
pip install -e "./fastchat[model_worker,llm_judge]"

The placement optimizer requires a valid Gurobi license. Gating replay and evaluation from the bundled artifacts can run without solving a new placement.

The following command replays one Qwen2 layer at 5 TFLOPS and 50 GB/s. Small --hd-comp and --hd-bw values are interpreted as TFLOPS and GB/s, respectively.

python evaluation/scripts/simulate_hd_gating_from_scores.py \
 --scores-npz expert_trace/qwen/score/gating_score_reasoning.npz \
 --output-json /tmp/hda_qwen_replay.json \
 --model-name qwen --top-k 8 --layers 0 \
 --reward-comp -18000 --reward-comm -0.0001 \
 --hd-mesh-rows 4 --hd-mesh-cols 8 \
 --hd-comp 5 --hd-bw 50 --chunk-size 32 --device cpu

python evaluation/scripts/e2e_hda.py \
 --cwd . --model ds --batch 32 \
 --mesh 4 8 --comp 5 --bw 50 --topology mesh \
 --results-json /tmp/hda_e2e.json

Change --topology to torus or fat_tree to exercise the additional interconnect models. For a fresh offline placement, run simulator.py for the desired model, layer, hardware setting, topology, and optional --memory-factor; optimizer.sh is the multi-layer driver.

The repository includes the routing traces and placement artifacts required by the bundled evaluations. Checkpoint-dependent fidelity evaluation requires the corresponding Hugging Face model checkpoint and is intentionally kept separate from the code release.

HDA-MoE: Hybrid Parallelism and Dynamic, Adaptive Scheduling for Mixture-of-Experts with 3D Near-Memory Processing
Haochen Huang, Shuzhang Zhong, Shengxuan Qiu, Zhe Zhang, Shuangchen Li, Cong Li, Dimin Niu, Hongzhong Zheng, Guangyu Sun, and Meng Li.
Accepted by IEEE Transactions on Computer-Aided Design of Integrated Circuits and Systems.

The DOI and final IEEE Xplore link will be added after online publication. Until then, please use the accepted-manuscript citation:

@article{huang2026hdamoe,
 author = {Haochen Huang and Shuzhang Zhong and Shengxuan Qiu and Zhe Zhang and
 Shuangchen Li and Cong Li and Dimin Niu and Hongzhong Zheng and
 Guangyu Sun and Meng Li},
 title = {{HDA-MoE}: Hybrid Parallelism and Dynamic, Adaptive Scheduling for
 Mixture-of-Experts with 3D Near-Memory Processing},
 journal = {IEEE Transactions on Computer-Aided Design of Integrated Circuits and Systems},
 year = {2026},
 note = {Accepted}
}

The conference predecessor is:

@inproceedings{huang2025hdmoe,
 author = {Haochen Huang and Shuzhang Zhong and Zhe Zhang and Shuangchen Li and
 Dimin Niu and Hongzhong Zheng and Runsheng Wang and Meng Li},
 title = {{HD-MoE}: Hybrid and Dynamic Parallelism for Mixture-of-Expert LLMs
 with 3D Near-Memory Processing},
 booktitle = {2025 IEEE/ACM International Conference on Computer-Aided Design (ICCAD)},
 year = {2025},
 pages = {1--9},
 doi = {10.1109/ICCAD66269.2025.11240984}
}

The model-evaluation path builds on FastChat, and functional-correctness evaluation builds on HumanEval. Their original licenses are preserved in the corresponding submodules.

The HDA-MoE code is released under the MIT License. The FastChat and HumanEval submodules retain their respective Apache-2.0 and MIT licenses.