3D-MoE: Complementarity-Guided Hardware-Software Co-Design for MoE Decode on 3D-Stacked NMP

Yudong Pan; Lian Liu; Shixin Zhao; Yintao He; Kaiyan Chang; Tianhua Han; Chenyang Ai; Haonan Zhu

cs.AROct 19, 2026

3D-MoE: Complementarity-Guided Hardware-Software Co-Design for MoE Decode on 3D-Stacked NMP

Authors: Yudong Pan, Lian Liu, Shixin Zhao, Yintao He, Kaiyan Chang, Tianhua Han, Chenyang Ai, Haonan Zhu,

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

Abstract

Mixture-of-Experts (MoE) have become a mainstream way to scale Large Language Models (LLMs), because each token activates only a small subset of experts, reducing inference cost. However, MoE decode remains highly memory-bound and exhibits strongly skewed, time-varying expert loads. 3D-stacked DRAM near-memory processing (NMP) provides the bandwidth that decode demands, yet its tiled organization makes load balancing expensive: expert parallelism (EP) preserves locality but suffers from hotspot-induced imbalance, while more distributed strategies weaken locality and increase NoC traffic. To address this locality-balance conflict, we observe that domain-specialized token routing creates offsetting expert-load fluctuations across decode steps, a property we call expert complementarity. Using this insight, we propose 3D-MoE, a hardware--software co-designed system for MoE decode on 3D-stacked NMP. 3D-MoE uses offline, one-time complementarity-guided grouping and locality-aware mapping to establish coarse-grained inter-group balance, handles residual intra-group hotspots through a lightweight Collaborative Memory-Network (CMN) that streams each hot expert's weight tiles once from its host DRAM channel to helper cores over a reconfigured local broadcast path, and uses cost-aware online remapping to remain robust to workload drift. Across MoE models, 3D-MoE achieves an average 4.63 $\times$ MoE-layer speedup over EP and improves end-to-end throughput by 1.58 $\times$ --1.86 $\times$ over the strongest NMP baselines, with 2.11% additional logic-die area.

Figures & tables

Fig. 1: MoE workflow, latency breakdown, and device roofline comparison.

Fig. 2: Expert activation analysis on Qwen3-30B-A3B. (a) Expert hit rates across domains. (b) Inter-layer variation.

Fig. 3: Stacked-DRAM NMP system organization.

Fig. 4: A comparison of prior MoE deployment strategies and our proposed method.

Model	Complementary	Positive	Random
DeepSeek-V2-Lite	0.19	0.67	0.43
Qwen3-30B-A3B	0.25	0.71	0.46
GLM-4.5-Air	0.21	0.69	0.44

TABLE I: Domain-preference similarity across pair types.

Model	LMSys [ 85 ]	ShareGPT [ 49 ]	HumanEval [ 87 ]
DeepSeek (64 E)	7.1 / 82.8% / 75.7%	6.4 / 83.7% / 77.1%	5.8 / 82.2% / 77.3%
Qwen3 (128 E)	10.1 / 82.5% / 74.8%	10.9 / 86.0% / 72.1%	12.2 / 88.3% / 80.5%
GLM-4.5 (128 E)	9.8 / 82.8% / 70.6%	11.3 / 85.6% / 74.5%	12.5 / 87.7% / 79.4%
$E$ : routed-expert count; entries are partners / coverage / cross-dataset retention.

TABLE II: Expert complementarity across models and datasets.

Fig. 5: Complementarity examples from layer 10 of DeepSeek-V2-Lite. (a) High-complementarity experts. (b) Low-complementarity experts.

Fig. 7: Software framework of 3D-MoE system: (a) Offline complementarity-guided expert grouping and mapping. (b) Runtime intra-group expert scheduling. (c) Online adaptive expert remapping.

Notation: for each Transformer layer
$\mathbb{E}$	All experts
$\mathbb{D}$	All decode steps
$\mathbb{G}$	All expert groups
$\mathbb{N}$	All NMP cores
$F_{e,d}$	Routed-token count of expert $e$ at decode step $d$
$C_{e_{1},e_{2},d}$	Pairwise load-offset score of expert $e_{1}$ and $e_{2}$ at decode step $d$

TABLE III: Grouping and mapping notation.

#Collab Cores	2	4	8	16
Ideal latency (ms)	0.84	0.42	0.21	0.105
On-demand remote fetch (ms)	1.08	0.69	0.45	0.36
Single-read + NoC broadcast (ms)	0.94	0.56	0.33	0.20
Intra-group TP (ms)	1.00	0.64	0.42	0.30
CMN (ms)	0.88	0.46	0.25	0.14

TABLE IV: Characterizing hotspot collaboration on NMP.

Collab. cores	Avg. flit lat. (cyc.)	Inter-router link	Inj./ej. channel	Fixed router pipeline	Dynamic wait /congestion
2	10.00	10.0%	20.0%	60.0%	10.0%
4	11.33	11.8%	17.6%	61.8%	8.8%
8	14.00	14.3%	14.3%	64.3%	7.1%
16	18.80	17.0%	10.6%	67.0%	5.3%

TABLE V: Flit-level breakdown of regular NoC broadcast.

Fig. 8: (a) Resource sharing via CMN. (b) CMN microarchitecture.

Level	Configuration
NMP-S/M/L	1/2/4 NMP devices; 64/128/256 GB total capacity
NMP device	1 H100-class xPU die + 4 NMP stacks on one interposer xPU: 0.99 PFLOPS; aggregate xPU–stack HBM BW: 2.68 TB/s
NMP stack	16-GB 8-hi DRAM; 16 channels; 16 NMP cores, $4\times 4$ mesh
NMP core	650 MHz; 1.33 TFLOPS; 166.4 GB/s local DRAM BW; 64-KB buffer
Interconnect	Intra-stack: 512-bit, 2 cycles/hop; Inter-stack: 128 GB/s/direction Inter-device: NVLink-like, 450 GB/s/direction

TABLE VI: NMP system configurations.

Model	Param	$\mathbf{N_{ex}}$	top- $k$	GPU Cfg	NMP Cfg
DeepSeek-V2-Lite [ 8 ]	16B	64	6	1 $\times$ H100	NMP-S
Qwen3-30B-A3B [ 75 ]	31B	128	8	2 $\times$ H100	NMP-M
GLM-4.5-Air [ 67 ]	106B	128	8	4 $\times$ H100	NMP-L

TABLE VII: MoE model configurations.

Fig. 9: Generalization of one-time offline expert grouping. (a) Cross-dataset generalization. (b) Cross-window stability; the dashed curve denotes Ours+Remap. (c) Cross-batch stability. (b) and (c) use Qwen3 on ShareGPT.

Fig. 10: MoE layer decode speedup of all schemes relative to the EP baseline under varying batch sizes.

Fig. 11: End-to-end throughput of all schemes relative to the EP baseline.

Fig. 12: Average compute resource utilization.

Fig. 13: TBT latency distribution on DeepSeek-V2-Lite.

Fig. 14: Ablation of 3D-MoE components.

Fig. 15: Normalized MoE layer decode speedup at different Op/B ratios.

Fig. 16: Normalized throughput for scale-out and low-expert-count workloads.

Components	Area (mm 2 )	Area ratio of logic die	Peak power (W)
DFU	0.56	0.94%	0.172
RCL	0.156	0.28%	0.033
DRU	0.26	0.42%	0.080
Look-Up Table	0.28	0.47%	0.052
Total	1.256	2.11%	0.337

TABLE VIII: Additional area and power of 3D-MoE.

Neural Network Hardware Acceleration

See all 48 papers

ThAME: 3D Memory-Enabled Heterogeneous Accelerator for LLM Mixture of Experts
Jul 19, 2026Pratyush Dhingra, Pramit Kumar Pal, Janardhan Rao Doppa +1AI Accelerator Inference Mixture-of-Experts Inference
Department of Electrical Engineering and Computer Science, Washington State University, Pullman, WA 99163 USA
ELMoE-3D: Leveraging Intrinsic Elasticity of MoE for Hybrid-Bonding-Enabled Self-Speculative Decoding in On-Premises Serving
Apr 16, 2026Yuseon Choi, Jingu Lee, Jungjun Oh +5LLM Serving Mixture-of-Experts Inference
KAIST · Samsung Electronics · GIST
DeaMoE: Efficient MoE Structure for Fast Small-Batch Decoding
Aug 14, 2026Zewen Jin, Shen Fu, Zeping Duan +8Mixture-of-Experts Inference Mixture-of-Experts Models
University of Science and Technology of China · Institute of Artificial Intelligence, Hefei Comprehensive National Science Center · Shanxi University

Jul 19, 2026J/K move · Enter open · S save

ThAME: 3D Memory-Enabled Heterogeneous Accelerator for LLM Mixture of Experts

Pratyush Dhingra, Pramit Kumar Pal, Janardhan Rao Doppa, Partha Pratim Pande

Department of Electrical Engineering and Computer Science, Washington State University, Pullman, WA 99163 USA

AI Accelerator Inference Mixture-of-Experts Inference LLM Inference Acceleration

Open paper AskPDF arXiv

Mixture of Experts (MoE) architectures have emerged as a dominant paradigm for scaling Large Language Models (LLMs). However, MoE inference on conventional hardware is constrained by three fundamental bottlenecks. These encompass the massive memory bandwidth required to fetch non-contiguous expert weights, the non-deterministic scatter-gather traffic generated by input-dependent token routing, and the tail-latency dependency imposed by synchronous expert output aggregation. To address these challenges, we propose ThAME, a three-dimensional (3D) heterogeneous multi-chiplet architecture for MoE inference. ThAME employs Ferroelectric Field-Effect Transistor (FeFET)-based non-volatile and DRAM-based volatile memory chiplets with a co-designed compute mapping strategy that aligns the distinct computational profiles of attention mechanisms and expert routing. Furthermore, we design a specialized Network-on-Chip communication backbone optimized to mitigate the bottlenecks associated with non-deterministic token routing traffic across the combinatorial space of input-dependent MoE traffic patterns. Experimental results demonstrate that ThAME outperforms state-of-the-art counterparts by up to 15.7x in terms of speedup and improves energy efficiency by up to 9.8x.

Abstract

Neural Network Hardware Acceleration

Explore similar work

ThAME: 3D Memory-Enabled Heterogeneous Accelerator for LLM Mixture of Experts

ELMoE-3D: Leveraging Intrinsic Elasticity of MoE for Hybrid-Bonding-Enabled Self-Speculative Decoding in On-Premises Serving

DeaMoE: Efficient MoE Structure for Fast Small-Batch Decoding