Hotdry.

Article

Stratum:3D堆叠DRAM与MoE推理的系统-硬件协同设计

Stratum通过Mono3D DRAM的近内存计算与分层数据放置,将MoE模型解码吞吐量提升8.29倍,为内存密集型AI推理提供系统-硬件协同优化范式。

2026-05-18ai-systems

Mixture of Experts(MoE)架构已成为大语言模型扩展的主流选择。通过稀疏激活机制,MoE 模型能在参数量级达到千亿甚至万亿的同时,保持与小型密集模型相近的推理计算成本。然而,这种参数稀疏性带来的并非免费的午餐 —— 每次前向传播仍需从显存加载大量专家权重,内存带宽迅速成为推理延迟的主导因素。传统 GPU-HBM 架构在面对 MoE 层的海量数据搬运时,显存墙问题愈发突出。

Stratum 是近期在 MICRO 2025 发表的一项系统 - 硬件协同设计研究,针对 MoE 推理的内存瓶颈提出了创新性解决方案。其核心思路是将新型单片 3D 堆叠 DRAM(Mono3D DRAM)与近内存处理(Near-Memory Processing, NMP)相结合,通过精细化的数据分层放置策略,最大化利用 3D 内存的高内部带宽特性。

Mono3D DRAM 的技术优势与分层机制

Mono3D DRAM 采用单片集成工艺,通过在同一晶圆上逐层制造 DRAM 单元实现垂直扩展。与依赖硅通孔(TSV)互联的传统 HBM 不同,Mono3D DRAM 使用 Cu-Cu 混合键合技术,垂直互联间距仅为 1µm,相较 HBM 的 10µm TSV 间距大幅提升。这一工艺差异带来两个关键收益:内部带宽可达 19 至 30 TB/s,远超 HBM 的规格;同时消除了 TSV 制造和复杂封装的高昂成本。

然而,激进的垂直缩放引入了一个被传统内存系统忽视的特性 —— 层间访问延迟差异。由于字线(Wordline)在 3D 堆叠的阶梯结构中需要更长的布线距离,深层单元的激活延迟(tRCD)显著高于浅层。Stratum 团队测量发现,在 1024 层堆叠中,最快层与最慢层的访问延迟差异可达 1.6 倍。

Stratum 将这一 "缺陷" 转化为设计机会:将 Mono3D DRAM 划分为多个逻辑层级(评估中采用 8 层),构建内部内存分层架构。通过将高频访问数据映射到低延迟层级、低频数据映射到高延迟层级,系统在不增加硬件复杂度的情况下实现了数据局部性优化。

异构计算架构与计算映射

Stratum 系统由三大组件构成:Mono3D DRAM 堆叠、集成近内存处理器的逻辑层芯片,以及通过 2.5D 硅中介层互联的 GPU(xPU)。逻辑层芯片上的 NMP 采用分层架构设计,包含多个处理单元(PU),每个 PU 配备近存储器处理单元集群(PE)、专用功能引擎(SFE)和环网路由器。

在计算任务划分上,Stratum 遵循 MoE 推理的计算特性:预填充阶段(Prefill)计算密集,由 GPU 负责;解码阶段(Decode)内存密集,卸载至 NMP 执行。这种异构分工充分发挥了两种处理单元的架构优势 ——GPU 擅长大规模并行计算,NMP 则利用 Mono3D DRAM 的高带宽优势处理内存受限操作。

NMP 内部采用张量并行策略处理专家计算。输入矩阵在各通道间复制,权重矩阵按维度分区,通过环网实现 all-gather 和 reduce-scatter 通信。注意力计算则利用头级并行,多个 PU 协同处理不同注意力头,通过交错执行线性运算和 Softmax 操作隐藏通信延迟。

主题感知的专家放置策略

MoE 模型的稀疏激活特性意味着不同输入样本激活的专家子集存在显著差异。Stratum 利用这一观察,在系统层面实现了智能化的专家数据放置。

具体流程如下:首先,使用轻量级主题分类器(基于 67M 参数的 DistilBERT)对输入查询进行主题标记;然后,根据预收集的各主题专家使用统计,预测当前批次的高概率激活专家(热专家);最后,将热专家权重放置于 Mono3D DRAM 的快速层级,冷专家放置于慢速层级,非 NMP 数据(由 GPU 处理)则放置于最慢层级 —— 因其访问受限于中介层带宽而非内存内部带宽。

当批次主题发生变化时,系统通过近内存的行交换缓冲区在不同层级间迁移专家数据。得益于 Mono3D DRAM 的高内部带宽,专家交换的时间和能耗开销均低于 1%,几乎可以忽略不计。

性能评估与工程启示

在 OLMoE、Mixtral、Qwen2.5 和 Llama-4 等模型的评估中,Stratum 相比 GPU 基线(vLLM on NVIDIA A6000/H100)实现了最高 8.29 倍的解码吞吐量提升和 7.66 倍的能效提升。与非分层 Mono3D DRAM 设计相比,分层机制带来了 1.32 至 1.45 倍的额外性能增益。

主题感知的专家放置策略被证明至关重要:准确的热专家预测可将吞吐量提升 1.32 至 1.51 倍。这表明系统级优化(数据放置策略)与硬件创新(3D 内存架构)具有显著的协同效应。

从工程实践角度,Stratum 的设计提供了以下可落地的参考:

内存分层粒度:当垂直层数达到 512 层以上时,分层机制即可产生 17-18% 的性能提升,建议以 8 层为初始分层粒度进行探索。

计算 - 内存协同:对于内存密集型解码阶段,近内存计算可将数据搬运开销降至最低,适合作为 GPU 的专用卸载引擎。

预测驱动的数据放置:轻量级主题分类(如 67M 参数模型)足以指导专家放置决策,分类准确率与最终推理性能正相关。

局限与展望

Mono3D DRAM 作为新兴内存技术,其商业化成熟度和供应链稳定性仍需时间验证。此外,虽然专家交换开销在评估中低于 1%,但在极高频次的主题切换场景下,累积开销的影响有待进一步研究。

Stratum 的价值在于展示了系统 - 硬件协同设计的潜力:当硬件特性(层间延迟差异)被软件策略(分层数据放置)巧妙利用时,可以突破传统架构的性能瓶颈。随着 3D 堆叠内存技术的演进,类似的协同优化思路有望在更广泛的 AI 推理场景中得到应用。


资料来源

  • Pan Y, Xia Z, Hsu PK, et al. Stratum: System-Hardware Co-Design with Tiered Monolithic 3D-Stackable DRAM for Efficient MoE Serving. In: MICRO 2025 - 58th IEEE/ACM International Symposium on Microarchitecture. IEEE Computer Society; 2025:1-17. DOI: 10.1145/3725843.3756043
  • 论文预印本: arXiv:2510.05245

ai-systems

内容声明:本文无广告投放、无付费植入。

如有事实性问题,欢迎发送勘误至 i@hotdrydog.com