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

Yudong Pan; Yintao He; Tianhua Han; Lian Liu; Shixin Zhao; Zhirong Chen; Mengdi Wang; Cangyuan Li

cs.ARMar 1, 2026

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

Authors: Yudong Pan, Yintao He, Tianhua Han, Lian Liu, Shixin Zhao, Zhirong Chen, Mengdi Wang, Cangyuan Li,

Organizations: SKLP, Institute of Computing Technology, Chinese Academy of Sciences, Beijing, China · University of Chinese Academy of Sciences, Beijing, China

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.

Figures & tables

Figure 1. Execution timelines of baseline MoE offloading systems and our proposed architecture.

Figure 2. MoE architecture with shared and routed experts.

Figure 3. Expert Activation across Batch Sizes, Models and Datasets 2 2 2 “Batch size” here refers to offline large-batch inference or zigzag batching that aggregates multiple small batches. . (a) Fine-grained activation for a specific configuration. (b) Summary of activation across all configurations.

Figure 4. Overview of TriMoE. (a) DIMM-NDP architecture. (b) Online expert scheduling & Offline expert relayout-rebalancing.

Figure 5. Compute Characterization. (a) Measured Throughput vs. Token Count. (b) Empirical GPU-CPU-NDP roofline.

GPU	H100 80GB: 819.6 TFLOPS (BF16), 2.04 TB/s Bandwidth, 64 GB/s PCIe bandwidth
CPU	Memory: 8 Channels $\times$ 2 DIMM, 2TB Capacity
	AMX: 90.1 TFLOPS (BF16), 307.2 GB/s Bandwidth
DIMM-NDP	DIMM Parameters
	DDR5-4800, 128GB/DIMM, 4 Ranks $\times$ 8 Bank Group $\times$ 4 Banks
	CCDL=12:CCDS=8:RRDS=8:RRDL=12:FAW=32:RP=34:BL=8:CL=40:WR=72
	NDP

Table 1. System configurations.

Model	params	shared exp	routed exp	top- $k$	expert params
DeepSeek-V2 ( Liu et al., 2024a )	236B	2	160	6	422GB
Qwen3-235B-A22B ( Yang et al., 2025 )	235B	0	128	8	423GB
GLM-4.5-Air ( Zeng et al., 2025a )	106B	1	128	8	190GB

Table 2. MoE models configurations.

Resource	Klotski	En-KTransformers	MoNDE	ours
GPU Compute Utilization	28.6%	57.6%	33.9%	66.0%
CPU Compute Utilization	$\times$	42%	$\times$	74.9%
NDP Compute Utilization	$\times$	$\times$	70.1%	87.8%

Table 3. Resource utilization comparison

Figure 6. MoE decode speedup over Klotski.

Figure 7. End-to-End throughput over Klotski.

Figure 8. Ablation study of TriMoE with batch size 512.

Figure 9. Sensitivity study of NDP count and CPU TFLOPS.

Mixture-of-Experts Inference

Mixture-of-experts inference studies how to run sparsely activated expert models efficiently while preserving their quality.

See all 165 papers

CoX-MoE: Coalesced Expert Execution for High-Throughput MoE Inference with AMX-Enabled CPU-GPU Co-Execution
May 18, 2026Muyoung Son, Yi Chen, Seungjae Yoo +2Expert Offloading GPU Acceleration
KAIST · Daejeon, Republic of Korea
Democratizing MoE inference on commodity GPUs with CoMoE
Oct 7, 2026Ruwen Fan, Yuezhi Zu, Minhui Xie +6Communication-Efficient Mixture-of-Experts Mixture-of-Experts Inference
Jimmy · Alibaba Cloud Computing · Tsinghua University
Fast MoE Inference via Predictive Prefetching and Expert Replication
May 12, 2026Ankit Jyothish, Ali Jannesari, Aishwarya Sarkar +1Expert Load Balancing GPU Acceleration
Iowa State University Ames, USA

May 18, 2026J/K move · Enter open · S save

CoX-MoE: Coalesced Expert Execution for High-Throughput MoE Inference with AMX-Enabled CPU-GPU Co-Execution

Muyoung Son, Yi Chen, Seungjae Yoo, Soongyu Choi+1

KAIST · Daejeon, Republic of Korea

Expert Offloading GPU Acceleration Mixture-of-Experts Inference LLM Inference Acceleration Heterogeneous Computing

Open paper AskPDF arXiv

The Mixture-of-Experts (MoE) architecture improves computational efficiency via sparse expert activation, but throughput-oriented inference faces substantial GPU memory pressure due to a significant parameter size and intermediate data. Prior works attempt to mitigate this using expert offloading with micro-batching or by offloading computation to the CPU. However, the fragmented workload resulting from micro-batching degrades operational intensity, causing expert execution to become memory-bound. Meanwhile, CPU offloading is constrained by slow PCIe transfers and its limited applicability to attention computation in the decode stage. Consequently, these inefficiencies prevent effective system utilization, severely restricting the end-to-end throughput of MoE inference. To address these challenges, this paper proposes CoX-MoE, an Advanced Matrix Extensions (AMX)-enabled CPU-GPU collaborative system that comprehensively optimizes MoE inference by combining coalesced expert execution with strategic workload orchestration for higher throughput. CoX-MoE introduces (i) a coalescing-aware orchestration policy to jointly optimize resource allocation by adopting ordinary batch, instead of micro-batch, for expert computation and selective attention offloading, and (ii) a static expert-aware stratification scheme that pre-assigns frequently activated experts to the GPU, mitigating PCIe transfer overhead and balancing workload for the CPU and GPU during inference. Compared to state-of-the-art frameworks, CoX-MoE delivers significant gains, achieving up to 7.1x and 2.4x higher throughput than FlexGen and MoE-Lightning, respectively.

Abstract

Mixture-of-Experts Inference

Explore similar work

CoX-MoE: Coalesced Expert Execution for High-Throughput MoE Inference with AMX-Enabled CPU-GPU Co-Execution

Democratizing MoE inference on commodity GPUs with CoMoE

Fast MoE Inference via Predictive Prefetching and Expert Replication