Mixture‑of‑Experts (MoE) has become a mainstream approach to scaling large language models (LLMs) because each token activates only a small subset of experts, reducing inference cost. Yet MoE decoding remains highly memory‑bound and exhibits strongly skewed, time‑varying expert loads.
3D‑stacked DRAM near‑memory processing (NMP) provides the required bandwidth, but 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.
We observe that domain‑specialized token routing creates offsetting expert‑load fluctuations across decode steps, a property we call expert complementarity.
Leveraging this insight, we propose 3D‑MoE, a hardware‑software co‑designed system for MoE decoding on 3D‑stacked NMP. 3D‑MoE uses offline, one‑time complementarity‑guided grouping and locality‑aware mapping to achieve coarse‑grained inter‑group balance. Residual intra‑group hotspots are handled by 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. Cost‑aware online remapping keeps the system robust to workload drift.
Across various 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 only $2.11\%$ additional logic‑die area.
Review