Summary
Large language models use a special design called Mixture-of-Experts (MoE) to be very powerful without needing too much computing. But this design needs a lot of fast memory and can be tricky to run efficiently on new 3D memory-computer chips. The authors created HDA-MoE, a system that combines different ways of splitting up work and smartly schedules tasks to make the models run faster and use resources better on this new hardware. Their system makes MoE models run 10% to 270% faster compared to older methods. This could help make big language models work better on specialized chips that join memory and computing closely together.
Mixture-of-Experts (MoE)Large Language Models (LLMs)3D Near-Memory Processing (NMP)Hybrid ParallelismDynamic SchedulingTensor Parallelism (TP)Expert Parallelism (EP)Compute UtilizationCommunication OverheadRuntime Scheduling
Authors
Haochen Huang, Shuzhang Zhong, Shengxuan Qiu, Zhe Zhang, Shuangchen Li, Cong Li, Dimin Niu, Hongzhong Zheng, Guangyu Sun, Runsheng Wang, Meng Li
Abstract
Mixture-of-Experts (MoE) architectures have become a key technique for scaling Large Language Models (LLMs), enabling high model capacity with reduced computational cost. However, this efficiency comes at the expense of increased memory capacity and bandwidth demands. Recent 3D Near-Memory Processing (NMP) architectures, which vertically integrate memory and compute through hybrid bonding, provide high internal bandwidth and energy efficiency, making them attractive for accelerating MoE inference. Nevertheless, the distributed memory and compute organization of NMP systems introduces new challenges for mapping MoE workloads. Existing parallelization strategies, such as Tensor Parallelism (TP) and Expert Parallelism (EP), suffer from either high communication costs or unbalanced computation utilization, leading to inferior efficiency. In addition, the dynamic routing behavior of MoE models further complicates efficient deployment. To address these challenges, we present HDA-MoE, a framework that optimizes MoE execution on NMP architectures through hybrid parallel deployment and runtime scheduling. HDA-MoE integrates an offline hybrid parallel mapping algorithm with an online dynamic and adaptive scheduling mechanism to reduce communication overhead while improving computation utilization. Experimental results show that HDA-MoE achieves a speedup of 1.1x--3.4x over TP, 1.1x--1.5x over EP, 1.1x--3.7x over the Hybrid TP-EP compute-balanced baseline, and 1.1x--1.3x over HD-MoE. Source code is available at https://github.com/PKU-SEC-Lab/HDA-MoE-TCAD26.