3D-MoE:面向 3D 堆叠近存处理器的 MoE 解码互补性引导软硬件协同设计
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 · et al.
中文摘要
混合专家(Mixture-of-Experts, MoE)通过仅激活少量专家降低大语言模型推理成本,但其解码阶段受限于内存带宽且专家负载呈时变偏斜分布。3D 堆叠 DRAM 近存处理(NMP)虽提供高带宽,但片上网络(NoC)的瓦片化组织导致负载均衡代价高昂:专家并行(EP)保持局部性却易产生热点,分布式策略则增加通信开销。针对此局部性与均衡性的冲突,研究观察到领域专用令牌路由在解码步骤间产生互补的专家负载波动。据此提出 3D-MoE 系统,采用离线互补性引导分组与局部性感知映射实现粗粒度均衡,通过轻量级协作内存网络(CMN)将热点专家权重瓦片从主机 DRAM 通道流式传输至辅助核心以处理组内热点,并利用成本感知在线重映射应对负载漂移。在多个 MoE 模型上,3D-MoE 相比 EP 实现平均 4.63 倍的 MoE 层加速,端到端吞吐量较最强 NMP 基线提升 1.58 至 1.86 倍,逻辑芯片面积仅增加 2.11%。
关键要点
- 01指出 MoE 解码在 3D 堆叠 NMP 上面临局部性与负载均衡的冲突,EP 策略易导致热点不均。
- 02发现领域专用令牌路由产生的专家负载波动具有互补性,可作为优化负载均衡的关键洞察。
- 03提出 3D-MoE 系统,结合离线互补性分组、CMN 辅助热点权重传输及在线重映射机制。
- 04实验显示 3D-MoE 相比 EP 实现 4.63 倍 MoE 层加速,端到端吞吐量提升 1.58-1.86 倍。
- 05该方案在保持性能提升的同时,仅增加 2.11% 的逻辑芯片面积,具备硬件实现可行性。
解读
尚无解读。
原始英文摘要
arXiv:2610.26991v1 Announce Type: cross 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.