TriMoE: Augmenting GPU with AMX-Enabled CPU and DIMM-NDP for High-Throughput MoE Inference via Offloading

TL;DR

TriMoE introduces a hybrid GPU-CPU-NDP architecture with intelligent scheduling to efficiently deploy large MoE models, significantly improving inference speed by addressing memory and compute bottlenecks.

Contribution

It proposes a novel architecture and scheduling strategies that optimize expert placement across heterogeneous compute units for MoE inference.

Findings

Achieves up to 2.83x speedup over existing solutions.

Effectively maps experts to compute units based on their memory and compute characteristics.

Demonstrates improved efficiency in large-scale MoE inference.

Abstract

To deploy large Mixture-of-Experts (MoE) models cost-effectively, offloading-based single-GPU heterogeneous inference is crucial. While GPU-CPU architectures that offload cold experts are constrained by host memory bandwidth, emerging GPU-NDP architectures utilize DIMM-NDP to offload non-hot experts. However, non-hot experts are not a homogeneous memory-bound group: a significant subset of warm experts exists is severely penalized by high GPU I/O latency yet can saturate NDP compute throughput, exposing a critical compute gap. We present TriMoE, a novel GPU-CPU-NDP architecture that fills this gap by synergistically leveraging AMX-enabled CPU to precisely map hot, warm, and cold experts onto their optimal compute units. We further introduce a bottleneck-aware expert scheduling policy and a prediction-driven dynamic relayout/rebalancing scheme. Experiments demonstrate that TriMoE achieves up to 2.83x speedup over state-of-the-art solutions.