Stratum: System-Hardware Co-Design with Tiered Monolithic 3D-Stackable DRAM for Efficient MoE Serving

文章标题：Stratum：一种采用分层单片3D可堆叠DRAM进行高效MoE服务的系统-硬件协同设计
作者/机构：
- Yue Pan (加州大学圣地亚哥分校)
- Zihan Xia (加州大学圣地亚哥分校)
- Po-Kai Hsu (佐治亚理工学院)
- Lanxiang Hu (加州大学圣地亚哥分校)
- Hyungyo Kim (伊利诺伊大学厄巴纳-香槟分校)
- Janak Sharda (佐治亚理工学院)
- Minxuan Zhou (伊利诺伊理工学院)
- Nam Sung Kim (伊利诺伊大学厄巴纳-香槟分校)
- Shimeng Yu (佐治亚理工学院)
- Tajana Rosing (加州大学圣地亚哥分校)
- Mingu Kang (加州大学圣地亚哥分校)

A1 主要贡献

大型语言模型（LLMs）的持续发展催生了混合专家（Mixture of Experts, MoE）架构，该架构通过稀疏门控机制，每个输入只激活少数专家子网络，从而在拥有数十亿参数容量的同时，保持与小得多的模型相当的推理成本。然而，MoE层带来的海量数据对硬件部署构成了巨大挑战。

核心问题：
当前服务于MoE模型的硬件基础设施面临内存容量和计算能力的双重压力。尽管高带宽内存（HBM）技术提供了一定的解决方案，但其通过硅中介层提供的带宽仍然不足，导致GPU计算资源利用率低下，尤其是在内存密集型的LLM解码阶段。现有的近内存处理（NMP）或内存中处理（PIM）方案，无论是基于HBM基础裸片还是直接在内存裸片上集成计算单元，都分别受到TSV I/O连接数量有限、DRAM工艺实现逻辑单元的PPA（性能-面积-功耗）开销大等问题的制约。

研究目标：
为解决服务大型MoE模型的挑战，本文提出一个名为Stratum的系统-硬件协同设计方案，旨在集成一种新兴的内存技术——单片3D可堆叠DRAM（Mono3D DRAM）、近内存处理（NMP）和GPU加速，以构建一个高吞吐量、高能效且成本效益高的MoE服务系统。

创新点：
本文的主要贡献如下：
* 首次提出基于Mono3D DRAM的MoE服务系统-硬件协同设计方案Stratum。该方案通过3D混合键合技术将高密度Mono3D DRAM裸片与高性能逻辑裸片进行异构集成，并利用2.5D硅中介层将此Mono3D DRAM堆栈与GPU集成。这为传统的基于GPU-HBM的MoE服务系统提供了一种高吞吐量且具有成本效益的替代方案。

  • 硬件层面的创新设计。引入了一种内存内分层（in-memory tiering）机制，利用了Mono3D DRAM因垂直扩展而产生的固有访问延迟差异。此外，还提出了一种专为基于混合键合的Mono3D DRAM设计的NMP处理器，该处理器为专家计算和注意力计算集成了优化的数据映射与通信策略。

  • 系统层面的算法与系统协同优化。研究发现专家激活频率与用户请求的主题存在非均匀性。基于此，将专家分为“热”和“冷”两类，并分别分配到MonoE3D DRAM的快速和慢速层中。提出的主题感知服务系统，利用一个轻量级主题分类器预测请求主题，并根据主题对请求进行排队和分发，同时遵循定义的服务等级目标（SLOs）。

  • 跨层次评估验证了方案的有效性。通过设备、电路、算法和系统层面的综合评估表明，与最先进的GPU基线相比，Stratum在各种实际的MoE服务场景中，解码吞吐量提升高达8.29倍，能效提升高达7.66倍。

A3 背景知识

2.1 单片3D可堆叠DRAM (Monolithic 3D-Stackable DRAM)

Mono3D DRAM技术特性。Mono3D DRAM是一种有前途的DRAM持续扩展技术，受到了学术界和工业界的广泛关注【索引57, 59, 67, 69, 92】。与传统的2D DRAM技术相比，它通过垂直扩展提供了显著更高的内存密度。这得益于诸如纳米片场效应晶体管（FETs）等先进技术，这些技术提供了更紧密的栅极控制并支持堆叠通道架构；同时，也借鉴了3D NAND闪存的制造技术，包括逐层沉积、用于超薄介电隔离的高深宽比蚀刻以及密集的垂直集成【索引16, 36, 46, 83】。

结构与成本优势。Mono3D DRAM采用单片3D可堆叠的水平1T1C DRAM单元，并整合了字线（WL）阶梯结构和垂直连接的位线（BL），以连接多层内存单元，如图2所示。HBM由于TSV制造的低良率和复杂的裸片堆叠封装工艺而成本高昂，而Mono3D DRAM通过避免使用TSV并利用单片3D集成（在同一晶圆上顺序构建额外的DRAM层）提高了可扩展性，从而具有成本优势。此外，Mono3D DRAM通过更薄的裸片和单片集成带来的改进的垂直导热性，也获得了散热上的优势。

内部带宽优势。除了成本和散热优势外，Mono3D DRAM还为逻辑层提供了增强的内存带宽。它利用异构集成技术【索引16, 46】，并采用铜-铜（Cu-Cu）混合键合技术实现内存单元和逻辑外围电路之间的高速数据传输。图3比较了在同一2.5D集成平台上的Mono3D DRAM与HBM。HBM的内部带宽受限于TSV，其间距较粗（10 µm）【索引79】，导致带宽有限且面积开销大，从而降低了内存密度。相比之下，Mono3D DRAM利用DRAM和逻辑基础裸片之间的Cu-Cu混合键合，其间距要精细得多（1 µm）【索引9】，通过后端制程（BEOL）金属布线连接，实现了极高的内部带宽。


图2：采用垂直堆叠水平1T1C DRAM单元的单片3D可堆叠DRAM。位线垂直布线以避免感知裕度变化，字线通过阶梯结构布线。由于字线阶梯结构，不同层的激活延迟有所不同。

外部带宽限制与NMP的必要性。尽管内部带宽更高，但如图3所示，Mono3D DRAM的外部带宽限制与HBM类似，都受限于中介层I/O接口的带宽。此外，先前的工作【索引63】指出，数据传输到外部处理器（包括在逻辑基础裸片上的路由和通过中介层I/O接口）会产生巨大的能耗。这些低效性凸显了在逻辑裸片上与Mono3D DRAM一同集成近内存处理（NMP）的必要性，以利用内部带宽并提高能效。

垂直扩展带来的挑战与机遇。尽管Mono3D DRAM具有实现超大内存容量的潜力，但其垂直可扩展性受到各层之间访问延迟显著差异的限制。如图2所示，位于阶梯结构底部的WL由于WL布线线性延长而经历了增加的寄生电容和电阻。当Mono3D DRAM扩展到数百层时，这种延迟不平衡变得非常显著。与其基于最坏情况的访问延迟进行设计，不如拥抱这种延迟异构性来提升系统级性能。这一挑战自然地催生了一种称为“内存内分层”（in-memory tiering）的架构方法，将在第3节详细讨论。值得注意的是，Mono3D DRAM的扩展趋势与3D NAND闪存一致，因为它利用了已扩展至400层以上的类似制造工艺【索引70】。此外，最近的白皮书表明，将这种扩展延伸到500甚至1000层是可行的【索引22, 45】。鉴于这些进展和垂直扩展的预期轨迹，我们假设最多可堆叠1024个字线（WL），以反映近期的可行性。


图3：在带有xPU裸片的2.5D集成平台上的HBM与Mono3D DRAM对比。HBM和Mono3D DRAM分别通过TSV和Cu-Cu混合键合连接到逻辑基础裸片。

2.2 混合专家LLMs (Mixture of Expert LLMs)

MoE模型的优势。正如LLM扩展定律【索引52】所示，密集型Transformer模型的准确性随规模增大而提高，但其训练和服务的成本也随之增加。近期的MoE模型，如OLMoE【索引60】、Mixtral【索引51】、Deepseek V3【索引27】、Time MoE【索引74】、DBRX【索引25】、LLaMA-4【索引4】和Kimi-K2【索引78】，通过每个token只激活一小部分专家，提供了一种引人注ทธิ的替代方案。这种稀疏激活提高了训练的可扩展性，并使得在不相应增加预训练成本的情况下实现巨大的参数量成为可能【索引54】，同时保持了与较小密集模型相当的推理成本【索引32】。另一方面，MoE模型需要一个路由机制，其中门控网络使用学习到的路由器参数从token表示（FFN输入或中间激活）中计算专家分配分数，以确定稀疏的专家选择模式【索引32】。每个token随后被分派到其选定的一个或多个专家进行独立处理，当每个token使用多个专家时，它们的输出通常通过使用路由分数的加权聚合来组合，以产生该层的最终输出【索引27, 32, 51】。


图1：密集型基于Transformer的LLM（左）与混合专家（MoE）LLM（右）的架构。

MoE模型的硬件部署挑战。MoE模型中MLP模块的切换特性引入了独特的硬件部署挑战。首先，MoE模型规模庞大，专家权重占总大小的主导地位，例如在Mixtral 8x7B中超过95%【索引51】，给GPU内存带来巨大压力。其次，每个token的专家使用情况动态变化且事先未知，当专家分布在不同计算单元时，会导致负载不平衡【索引27】。近期的工作旨在通过预先预测专家使用情况来减少通信开销。ExpertFlow【索引39】采用轻量级代理模型来预测路由路径，而MoE Infinity【索引85】则利用跨层激活分析来统计性地预测专家选择。在混合GPU和近内存处理系统中，Duplex【索引89】根据延迟模型和批量大小，动态地将专家计算分派到GPU或NMP单元。

专家的领域特化现象。在训练期间，MoE模型通常包含一个专家不平衡损失项，以防止“饥饿”现象（即一个或多个专家被选择的频率远低于其他专家），从而鼓励更均匀的专家利用率【索引32, 51】。然而，随着训练的进行，专家之间自然会产生领域特化【索引13, 58, 87】。随着专家数量的增加和共享专家的引入，这种特化变得越来越明显，共享专家巩固了通用知识，增强了被路由专家的领域特异性【索引4, 24, 27, 60】。基于这一观察，最近的工作探索了利用专家对特定领域的亲和性来加速纯GPU环境下的推理【索引33, 81, 87】。

基于主题的专家使用亲和性。我们通过性能分析观察到，专家的使用情况与查询的主题有明显的关系：特定主题会显著更频繁地激活某些专家。图4展示了一个例子，其中LLaMA-4 Scout在MMLU子集中的数学和逻辑相关主题上表现出超过90%的领域特定专家亲和性。在我们的服务系统中，我们利用了主题特定的专家亲和性，首先进行离线分析，收集不同主题下专家命中率（即使用概率）的统计数据。在在线服务期间，调度器中的一个轻量级主题分类器为批次中所有传入的查询分配主题标签。基于此分类，系统将常用专家映射到更快的Mono3D DRAM层以优化访问延迟，具体讨论见第5节。


图4：LLaMA-4 Scout（16个专家）的专家命中率分析。

A2 方法细节

3 Stratum 概述

3.1 系统概述

Stratum系统配置。Stratum处理系统由一个xPU裸片和可配置数量的单片3D可堆叠DRAM芯片组成，通过硅中介层连接，并具备近内存计算能力。我们展示了三种不同的示例配置（图5），以适应不同规模的模型，使用不同数量的Mono3D DRAM芯片。Stratum-L使用一个NVIDIA H100计算裸片作为xPU，并配有六个通过中介层互连的Mono3D DRAM芯片。Stratum-S使用一个NVIDIA RTX A6000裸片作为xPU，配有一个提供32GB内存的Mono3D DRAM芯片。Stratum-XL由两个Stratum-L模块组成，总共提供384GB内存，用于服务更大的模型。这些配置适应了不同的计算和内存需求，并可以使用如NVLink【索引50】之类的跨芯片互连进行扩展。


图5：Stratum的示例配置。

Mono3D DRAM芯片架构。每个Mono3D DRAM芯片由顶部的内存裸片和底部的逻辑裸片组成，它们通过Cu-Cu混合键合互连以提供高内部带宽。此外，为了利用Mono3D DRAM垂直各层之间的访问延迟差异，我们在内存裸片内部引入了内存分层机制。底部的逻辑裸片实现了一个强大的近内存处理器（NMP），以支持LLM推理，而无需总是将数据取到主处理器，具体细节在§3.2中阐述。

Stratum服务系统流程。图6描述了基于Stratum的服务系统的流程。在现实的服务场景中，用户提交的查询主题各不相同。当用户发送推理请求时，主处理器使用一个轻量级的主题分类器来确定查询的主题。这些请求随后被带有主题标签地加入到服务队列中。调度器会定期将服务队列中的推理请求分组，然后分派到Stratum处理系统。为了提升用户体验，一个关键的服务等级目标（SLO）是首个令牌生成时间（TTFT），它确保请求在开始处理前不会等待太久。当SLO允许时，调度器会优先将相同主题的请求批处理，以最大化专家放置带来的好处。内存映射器通过查询预先分析的专家使用表，为批次构建聚合的专家命中预测，并生成一个专家到Mono3D DRAM层的目标映射。在处理每个带有不同主题标签的新批次之前，会执行专家交换以满足目标布局。考虑到每个阶段的算术强度，计算映射器将预填充（prefill）阶段分配给xPU，将解码（decode）阶段分配给Stratum NMP，这与【索引67】中的策略类似。此外，轻量级的主题分类由主处理器执行。


图6：基于Stratum的服务系统。

3.2 Stratum近内存处理

NMP整体架构。图7展示了Stratum NMP的架构，它将处理组件组织在内存层次结构的多个级别——包括芯片级、通道级和bank级——以利用3D集成的优势。这一架构决策旨在加速注意力和专家计算，这两者是MoE模型中的基本瓶颈。

芯片级处理器。图7(a)展示了逻辑裸片处理器与Mono3D DRAM裸片的集成。逻辑裸片由多个处理单元（PU）组成，每个PU都与一个专用的Mono3D DRAM通道耦合。这些PU通过一个双向环形片上网络互连，该网络旨在优化LLM工作负载中的数据通信模式，如reduce-scatter和all-gather。请注意，环形网络仅在NMP模式下使用。在常规内存操作模式下，逻辑裸片NMP保持非活动状态，确保对传统内存访问模式的干扰最小。在NMP模式下，xPU通过标准的DRAM接口将输入（如查询、隐藏令牌向量等）流式传输到Mono3D DRAM bank中的保留行。计算完成后，xPU通过访问专用地址空间来检索处理后的结果。

通道级处理单元（PU）。每个PU旨在处理分配给其相应DRAM通道的数据，以避免跨通道DRAM访问——考虑到Mono3D DRAM和逻辑裸片之间存在大量的垂直布线，这是一个关键的考虑因素。图7(b)展示了PU的微架构，它由一个近bank处理元件（PE）集群、一个共享内存、一个特殊功能引擎、一个环形路由器和一个归约器组成。近bank PE集群集成了多个为GeMM和GeMV操作优化的PE。在通道内实现的归约器采用并行归约树，根据需要聚合通道内多个PE的部分和（psums）。环形路由器包含一个本地交换机，用于在PU间通信期间进行高效数据路由，以及一个聚合器，用于原地数据归约。传入的数据流可以立即在路由器中累积，而无需经过共享内存。累积的结果可以本地存储在PU中或根据需要转发给相邻的PU。特殊功能引擎执行特殊操作，如注意力机制中的Softmax以及专家层中常见的其他激活函数（如SiLU, GeLU）。它包括一个向量寄存器文件、一个标量寄存器文件和多个算术单元。特殊功能引擎以单指令多数据（SIMD）的方式运行，通过将复杂函数分解为简单原语，并在向量和标量寄存器文件中源取和存储操作数或中间结果，来最大化数据重用。

Bank级处理元件（PE）。在bank级别，如图7(c)所示，每个PE被设计用于执行GeMM和GeMV操作。bank级PE由一个张量核心和专门的内存组件集成而成：一个矩阵寄存器文件、一个psum内存和一个简单的本地内存控制器。该内存控制器直接与其对应的DRAM bank接口，通过一个可编程的分层表动态地将行地址转换为特定的内存层标识符，从而实现自适应的DRAM延迟控制（tRCD）以优化性能。行交换缓冲区存储临时行数据，以支持层间数据移动，而无需显式的外部数据获取。张量核心包含n个并行的k-tap点积引擎和n个本地累加器。双缓冲的psum内存结构同时支持中间结果的累积和输出传输。处理后的输出可以传递给特殊功能引擎进行逐元素函数评估，或返回到通道级共享内存以进行后续的计算步骤。

与HBM-NMP方案的比较。Stratum的架构，特别是为基于混合键合的Mono3D DRAM集成而优化的，与以HBM为中心的NMP方法（如AttAcc【索引67】、Neupims【索引43】和Duplex【索引89】）有所不同。片上环形网络旨在支持MoE推理的通信模式（例如，all-gather, reduce-scatter），取消了Duplex【索引89】中使用的集中式全局缓冲和交叉开关，从而提高了可扩展性并简化了物理设计。与依赖专用Softmax单元的Duplex【索引89】和AttAcc【索引67】不同，我们的基于SIMD的引擎通过编程指令执行通用的非线性运算符。此外，该处理器完全实现在逻辑裸片上，并通过混合键合连接到Mono3D DRAM裸片，避免了在AttAcc【索引67】和Neupims【索引43】中观察到的DRAM制造工艺约束和TSV带宽限制。在电路层面，Stratum引入了Mono3D DRAM特有的原语——包括分层表和行交换缓冲区——以利用分层内存延迟并加速专家迁移，从而服务于MoE模型。


图7：Stratum NMP架构。(a) 芯片级处理器概览。(b) 通道级处理单元（PU）的微架构，以及(c) bank级处理元件（PE）的微架构。

4 Stratum算子映射与执行

4.1 专家处理

MoE层执行流程。一个MoE层的执行流程包括三个主要阶段：令牌路由、专家计算和结果聚合。如图8(a)所示，一个批次中的令牌可能根据在xPU上计算的路由决策被路由到不同的专家。由于路由步骤的计算成本可以忽略不计（通常涉及一个轻量级的线性层，例如4096输入和8输出维度），这是可行的。随后，只有被激活的专家——即那些至少被分配了一个令牌的专家——才会被执行。最后，所有专家的输出通过加权求和的方式合并，以产生最终的输出令牌。专家计算和结果聚合都由Stratum NMP处理器执行。

单个专家内的计算。在MoE模型中，单个专家的计算通常由三个级联的GeMM操作组成【索引4, 51】，如图8(b)所示。设$T$为路由到一个专家的批次中的令牌数，$H$为隐藏维度，$M$为中间维度。首先，大小为$T \times H$的输入隐藏矩阵$X_1$与两个大小为$H \times M$的权重矩阵相乘以产生中间矩阵$Z_1$和$Z_2$（均为$T \times M$大小）。一个非线性的、逐元素的激活函数应用于$Z_1$，其结果与$Z_2$通过Hadamard积结合形成$X_2$。最后，$X_2$与一个大小为$M \times H$的向下投影权重矩阵相乘，产生大小为$T \times H$的输出$Z_3$。

分区策略。在实践中，不同的专家可能会接收到不同数量的令牌。此外，专家可能被映射到Mono3D DRAM层次结构中的不同层，每层具有不同的内存访问延迟，这进一步加剧了负载不平衡。因此，将多个专家分布在不同的PU上可能会导致PU之间严重的负载不平衡问题。为了解决这个问题，多个被选中专家的执行被安排为顺序进行，例如，一次处理一个专家。所有PU协同合作，使用张量并行性一次处理一个专家。这要求将所有三个GeMM操作中涉及的每个矩阵划分为多个块，每个块分配给一个PU进行并行执行。图8(b)展示了Stratum中使用的矩阵分区方案，为简化起见，仅假设有四个PU。沿不同维度进行分区会在输入复制、权重复制和部分和聚合之间产生权衡。我们避免沿$T$维度进行切分，以防止复制专家权重，因为专家权重占用了主要的内存。相反，我们将GeMM1和GeMM2的权重矩阵垂直切分，而将GeMM3的权重矩阵水平切分。这种方法以最初将$X_t$复制到多个PU，然后从多个PU收集$Z_3$的部分结果为代价，消除了向上投影和向下投影阶段之间的数据通信。请注意，复制$X_t$的成本得到了很好的摊销，因为所有活跃专家的输入矩阵$X_1$都源自$X_t$（即批次中令牌的集合）。此外，从多个PU收集并为$Z_3$进行归约的计算可以与下一个专家的处理并行进行，从而有效地隐藏了延迟。


图8：(a) MoE的令牌到专家映射示例。(b) 具有T个路由令牌的专家的计算阶段及矩阵分区，为简化起见假设有四个PU。(c) Stratum中MoE层的分步执行。

执行阶段。图8(c)展示了MoE层的分步执行流程。xPU首先将一批输入令牌连同相应的专家ID和缩放权重发送到Mono3D DRAM，并将其切换到NMP模式（步骤1）。由于采用了矩阵分区策略，每个Mono3D DRAM通道必须接收整个输入令牌矩阵。接下来，Stratum NMP处理器通过步骤2-7顺序执行被激活的专家。在步骤2和3中，所有PE中的张量核心执行两个向上投影的GeMM操作，以计算中间结果$Z_1$和$Z_2$。步骤4和5涉及使用特殊功能引擎应用激活函数并执行Hadamard积。得益于矩阵切分策略，每个PU无需进行PU间通信即可获得第三个GeMM所需的输入切片。第三个GeMM在步骤6中执行，随后是一个reduce-scatter操作，以在PU之间累积最终的输出矩阵$Z_3$。然后对每个剩余的被激活专家重复步骤2-7。在步骤9中，特殊功能引擎对专家输出进行加权求和，以产生最终的输出令牌，这些令牌被写回到指定的DRAM内存空间。最后，在步骤10中，Mono3D DRAM退出NMP模式，xPU通过访问指定的地址空间来检索计算出的令牌。


图9：专家处理的优化时序图。

执行优化。图9展示了一个优化的执行流水线，旨在最大化计算和通信资源的利用率。首先，为减轻xPU到Mono3D DRAM数据传输的延迟，输入令牌矩阵被划分为多个切片，每个切片发送到不同的Mono3D DRAM通道。这减少了输入准备的开销，随后通过高速逻辑裸片环形网络实现的all-gather操作为所有PU重构完整的输入矩阵。其次，GeMM2的计算与激活函数的评估重叠进行，因为它们之间没有数据依赖关系，从而实现了更好的流水线利用。第三，与GeMM3相关的reduce-scatter通信与下一个专家的GeMM1执行并行进行，从而将通信延迟隐藏在计算之后。最后，加权求和操作在每个专家的输出可用时立即由特殊功能引擎执行，从而最小化空闲周期并提高整体吞吐量。

PU内部分区。在每个PU内部，由于高带宽的共享内存，PE之间的通信开销可以忽略不计。因此，PU内部的矩阵分区主要集中在最大化张量核心的映射利用率上。为此，权重矩阵的较长维度被分区，并将得到的子块分布到各个PE上进行并行处理。因此，向上投影的权重切片$W_{1,2}[i]$通常是水平分区的，而向下投影的权重切片$W_3[i]$则是垂直分区的，以优化计算效率。

4.2 注意力处理

动机与挑战。大型语言模型（LLM）的生成任务通常受到对键值（KV）缓存的数据访问的瓶颈制约。Stratum通过利用Mono3D DRAM和基础裸片上的NMP逻辑之间的高带宽来高效解决此问题。然而，要充分利用此带宽，关键在于及时有效地处理从DRAM层垂直获取的数据。否则，由于逻辑裸片内部的计算或通信瓶颈，可用带宽可能会被低效利用。

策略：头级并行。由于注意力头之间没有数据依赖性，Stratum利用头级并行来高效执行注意力操作。图10(a)展示了在逻辑裸片上分配注意力头任务的情况。来自一组请求的多个注意力头可以被分配到不同的Mono3D DRAM设备上。分配的头数量可以根据网络模型（例如MoE模型中常见的分组查询注意力【索引4, 51】）以及在服务延迟要求下的请求并发性而变化。为提供一个适应不同头级并行度的处理架构，逻辑裸片上的PU可以被灵活地划分为多个大小可变的PU组，只要一个组内的PU是邻居并通过片上环形拓扑连接，如图10(a)所示，其中带箭头的PU表示环上的PU。这种安排也允许通过高速双向链接实现高效的组内通信。我们为每个组至少分配两个头，以便在不同计算阶段之间进行交错处理，从而提高吞吐量和硬件利用率——例如，一个头可能在执行线性操作，而另一个在执行Softmax。

PU组内部分区与通信。图10(b)描述了单个头的键（Key）和值（Value）矩阵如何在PU组内的PU之间进行分区。通常，序列长度维度（例如512-32k个令牌）远大于注意力头维度（例如64-128），这促使我们沿序列长度维度进行分区。然而，Softmax操作本质上需要所有令牌的全局信息，即用于归一化的全局最大值（即 $row\_max(QK^T)$）和指数的全局和（即 $\sum exp(QK^T - row\_max(QK^T))$）【索引35】。幸运的是，每个PU可以使用其专用的特殊功能引擎独立计算局部最大值和局部和，仅需在PU之间交换标量值即可推导出全局值。为了在解码阶段平衡PU的工作负载，新生成的键值对以轮询方式分布在一个PU组内的不同PU上。


图10：注意力层的执行。(a) 头（例如八个）在PU组（例如四个）间的分配。(b) PU组内部：注意力算子映射。(c) 多个头（例如两个）的并发处理。

优化执行流程。图10(c)展示了一个PU组内多个注意力头的优化执行流程。最初，xPU将计算出的键值对写入相应的DRAM通道。查询（可能是分组查询矩阵）被划分为多个切片，每个切片分配给PU组内的一个不同DRAM通道。随后，组内的所有PU通过子环的all-gather操作获得完整的查询矩阵，这与MoE层类似。当多个头被分配到同一个PU组时，Softmax操作可以与$QK^T \times V$和$Attn. \times V_{proj}$算子交错进行，以最小化总延迟。注意，Softmax算子被分为三个步骤，并有两轮PU间通信，如图10所示。最后，第一个头的reduce-scatter延迟可以隐藏在第二个头的$Attn. \times V_{proj}$操作中。

总结。总而言之，Stratum通过优化的数据放置、算子映射和调度，最大程度地利用了混合键合所带来的垂直带宽。该系统对专家计算采用跨所有PU的张量并行，对注意力计算采用分组PU的头并行。这两种策略都将大部分内存访问引导至本地Mono3D DRAM bank，通过混合键合I/O进行。剩余的PU间通信，如all-gather、reduce-scatter或标量交换，由片上环形网络高效支持。此外，调度器将矩阵运算（如GeMM和GeMV）与特殊函数计算（如SiLU和Softmax）重叠，协调片上通信和计算，以提高整体并行性。

4.3 设计的物理约束

热与面积约束。通过混合键合集成的Mono3D DRAM和逻辑裸片处理器必须同时满足热和面积约束。在NMP模式下，系统可能会受到一个峰值功率预算$P_{peak}$的限制，该预算由热分析确定（见§6.2.2），从而得出以下功率约束：

此处，$BW_{fast\_tier}$是Mono3D DRAM分层中最快层的峰值带宽，$E_b$是通过混合键合从DRAM层到逻辑裸片数据传输的每比特能量，$N_{mac}$是张量核心中乘法累加（MAC）单元的总数，$f_{logic}$是逻辑裸片的工作频率，$E_{mac}$是每次MAC操作的能量。杂项功率$P_{misc}$包括逻辑裸片SRAM、寄存器文件、路由器、特殊功能引擎、PU内归约器和本地内存控制器，其值根据算子类型和数据流而变化。

面积约束。虽然基于混合键合的数据I/O不消耗逻辑裸片上的有源面积，但TSV仍然是为DRAM和逻辑裸片供电所必需的【索引88】。因此，必须满足以下面积约束：

其中，$A_{PD}$是用于供电的总TSV面积，$N_{mac}$是MAC单元数，$A_{mac}$是在$f_{logic}$频率下工作的每个MAC单元的面积，$A_{PHY}$是xPU-DRAM接口的物理通信层面积，$A_{peri}$是逻辑裸片上的低压Mono3D DRAM外围电路（如D/Q缓冲器、电平转换器等）的面积，$A_{misc}$捕获了与$P_{misc}$类似的杂项逻辑面积组件，而$\alpha$是目标利用率。假设一个面积为$A_{TSV}$的单个TSV可以提供$I_{TSV}$的电流，总TSV面积由下式给出：

其中$V_{dram\_c}$、$V_{dram\_p}$和$V_{logic}$分别表示Mono3D DRAM核心、高压外围电路和低压逻辑裸片的供电电压。

设计指导。方程(1)、(2)和(3)将用于指导逻辑裸片处理器的设计配置（见§6.2.3）。

5 Stratum 算法-系统协同优化

5.1 专家使用预测

利用专家特化。如§2.2所述，预训练的MoE模型在推理时通常表现出领域特定的专家特化【索引87】，如图4所示。鉴于MoE推理的主要挑战之一是处理所有专家的巨大总参数量，这种特化为高效推理和服务提供了宝贵的机会。当专家特化与特定的查询主题一致时，优化MoE专家的放置就成为可能。对于给定的主题，具有较高使用概率（命中率）的专家可以被映射到更快的Mono3D DRAM层，从而减少从DRAM到基础逻辑裸片的数据传输延迟。

主题分类器。Stratum的一个关键组成部分是一个主题分类器，它为传入的查询打上标签。这使得Stratum调度器能够估计每个查询的主题分布。结合每个主题的专家使用表（如图6所示），调度器将专家的权重矩阵分配到适当的专家层。我们的实现训练了一个基于DistillBERT【索引28, 72】的主题分类器，该分类器具有67M参数，可对6个主题进行分类，作为我们基于Stratum的在线服务系统的一部分。为了解决从标准NLP数据集到真实服务查询中观察到的多样化提示风格的分布偏移问题，我们采用了一个数据合成管道，使用基于GPT-4o的重写来增强训练数据。由于其紧凑的尺寸，我们的主题分类器在适度的请求率下（每秒少于四个查询），每个解码步骤引入的延迟开销小于2%，同时在真实世界的服务数据集（Chatbot Arena对话【索引3】）上，对于6主题模型分别达到了85.0%和81.0%的分类准确率。关于数据增强、训练和评估的更多细节在§6.3.1中提供。

5.2 数据放置策略

数据分类与放置原则。Stratum将MoE模型内的数据分为四类：热点专家权重、冷点专家权重、KV缓存和非NMP数据。热点专家包括共享专家以及对给定主题表现出高路由命中概率的其他专家。非NMP数据主要包括杂项参数，如位置嵌入参数、层归一化偏移和缩放参数等。这些通常用于外部处理器而非NMP的计算。通过利用不同内存层之间的异构访问延迟，可以优化数据放置策略以提高服务性能。


图11：针对具有分层内存的Mono3D DRAM-NMP系统的专家放置优化示例。

具体放置策略。如图11所示，Stratum将由xPU处理的非NMP数据分配给最慢的内存层，因为访问它需要穿过中介层瓶颈，这比最慢层的内部DRAM带宽慢一个数量级。这有助于将更快的内存层专用于NMP相关的工作负载。Stratum根据对特定主题请求的离线分析，将专家分为热点和冷点两类，将热点专家分配到更快的内存层，冷点专家分配到较慢的层。这种放置确保了热点专家能够受益于更快的Mono3D DRAM内存层提供的低延迟访问。专家权重的放置在算法1中有详细说明。每个专家权重被分割成多个分片，并根据张量并行策略（见§4.1）分布在Mono3D DRAM的bank中。从算法1获得的物理行地址到逻辑内存层的映射起到一种量化过程的作用，可通过分层表（见§3.2）进行配置。在我们的评估中，我们采用了一种统一的映射策略，为每个内存层分配相同数量的行（见§6.2.1）。KV缓存数据，其容量随着请求生成的进展而动态变化，存储在中间速度的内存中。在完成一个主题（例如主题A）的处理后，Stratum调度器会转换到一个新主题（例如主题B），并根据新主题的专家激活频率启动专家交换。为避免昂贵的主机处理器传输，这种交换使用近内存操作执行，如§3.2所述。具体来说，本地内存控制器通过将两个DRAM行临时缓冲在一个专用的行交换缓冲区（见图7(c)）中，然后再将它们写回到新的行地址，来完成交换。

算法1 专家权重放置

输入: 层数$L$；每层专家数$E$；活跃专家数$A$；使用频率$F = \{f_{le} | l \in [1, L], e \in [1, E]\}$；单个专家权重大小$S_w$ (字节)；DRAM bank数$N_{\text{bank}}$；DRAM行缓冲大小$S_{\text{rb}}$ (字节)；为NMP数据保留的DRAM行数 $\Phi$。
输出: 所有专家权重的DRAM行地址区间 $\{[r_{le}^{start}, r_{le}^{end}] | l \in [1, L], e \in [1, E]\}$。
1: $\Delta \leftarrow S_w / (N_{\text{bank}} \cdot S_{\text{rb}})$ // 一个专家占用的行数
2: $K \leftarrow L \cdot A$ // 指定的快速专家数量阈值
3: 将$F$降序排序得到 $\langle f_{l_1e_1}, \dots, f_{l_{LE}e_{LE}} \rangle$
4: for $i = 1$ to $L \cdot E$ do
5: if $i \leq K$ then
6: $r_{l_ie_i}^{\text{start}} \leftarrow (i - 1) \cdot \Delta$
7: else
8: $r_{l_ie_i}^{\text{start}} \leftarrow \Phi - (L \cdot E - i + 1) \cdot \Delta$
9: end if
10: $r_{l_ie_i}^{\text{end}} \leftarrow r_{l_ie_i}^{\text{start}} + \Delta - 1$
11: end for
12: return $\{[r_{le}^{\text{start}}, r_{le}^{\text{end}}] | l \in [1, L], e \in [1, E]\}$

A4 实验环境

Mono3D可堆叠DRAM配置

  • 技术: 采用垂直位线连接的3D可堆叠水平1T1C DRAM技术。
  • 规模: 设计扩展至1024层，特征尺寸35nm。
  • 结构: Bank结构如图12所示，一个MAT由1024个BL×1024个WL组成，一个bank由1024个MAT组成。芯片容量为32GB，内存密度为2.156 Gb/mm²，比最新的32Gb DDR5裸片高5.2倍。
  • 性能: 内部带宽根据内存层的不同，范围从19.01 TB/s到30.34 TB/s。设计了8个具有不同访问延迟的内存层级。
  • 建模: 使用Coventor过程模型【索引23】提取3D DRAM阵列的RC参数，并结合NeuroSim【索引56】提取的外围电路结果和DDR5标准【索引2】的时序进行仿真。1T1C模型基于【索引36】中的3D DRAM结构规范，使用Coventor SEMulator3D过程模拟器【索引23】构建。详细参数见表1。


图12：Mono3D DRAM bank配置。性能通过NeuroSim [56]和Coventor过程模拟器 [23]进行模拟。


表1：单片3D可堆叠DRAM参数

逻辑裸片处理器配置

  • 硬件实现: 使用SystemVerilog实现，并采用Cadence Genus【索引7】和7nm预测性工艺设计套件ASAP7【索引19】进行综合。采用IEEE754 FP-16算术数据格式。
  • SRAM建模: 逻辑裸片上的本地psum内存和共享内存使用FinCACTI【索引73】建模，并根据公开的SRAM规范【索引8, 47】进行校准。
  • 性能: 峰值性能为128 TFLOPS，包含64k个MAC单元，工作频率1GHz。
  • 仿真: 使用自研模拟器推导执行周期、片上通信周期和相关能耗指标。


表3：Stratum逻辑裸片处理器规格

系统与模型配置

  • 模型: OLMoE-1B-7B (7B), Mixtral 8×7B (47B), Qwen2.5-32B (32B), Llama-4-Scout (109B)。

  • 硬件平台:

    • GPU基线: NVIDIA RTX A6000 或 H100 SXM5 HBM3 GPU。
    • Stratum: Stratum-S, Stratum-L, Stratum-XL，分别对应不同的模型和GPU基线。

  • 软件与仿真:

    • GPU: 使用vLLM 0.8.1【索引55】在基准吞吐量模式下进行评估。GPU能耗由NVIDIA-SMI工具获取。
    • Stratum: 系统级模拟器包含请求生成器（泊松过程）、SLO感知调度器、内存与计算映射器，并与Stratum NMP模拟器接口。

  • 主题分类器:

    • 模型: 基于DistilBERT【索引72】的67M参数模型，用于6个主题的文本分类。
    • 数据集: 混合了多个标准NLP数据集（如Pile of Law, Atlas, SciQ, MATH等）和通过GPT-4o从Chatbot Arena【索引3】合成的真实风格查询数据。


表2：评估工作负载设置

A4 实验结果

6.2 硬件评估

  • 3D-DRAM中的分层机制：如图14所示，Mono3D DRAM的访问延迟随着WL层数的增加几乎呈线性增长，这是由于WL阶梯结构延伸导致的RC延迟增加。最快层的访问速度比最慢层快1.6倍。基于此，本文定义了8个时序层级。


图14：Mono3D DRAM跨WL层的延迟。插图展示了因WL RC延迟增加而导致的各种访问延迟。

  • 功率和面积预算：通过HotSpot【索引75, 76】进行热仿真，并考虑高端液冷方案，确定逻辑裸片的功率上限约为每个芯片45W。在121mm²的裸片总面积中，除去HBM3 PHY模块（23.94mm²）、低压外围电路（14.80mm²）和供电TSV（0.21mm²）后，逻辑裸片处理器的可用面积预算为82mm²。

  • 逻辑裸片处理器：根据功率和面积预算设计的处理器规格如表3所示。其峰值性能为128 TFLOPS，总活动逻辑面积为76.63mm²（利用率63%），峰值功率为42.67W，均在预算范围内。图15展示了面积和功率的详细分解，其中PEs占用了大部分面积，而分层表（Tiering Table）的面积开销极小（PE面积的0.1%）。


图15：(a) 逻辑裸片处理器的面积分解；(b) Mono3D DRAM-逻辑裸片在峰值性能下的功率分解。

6.3 系统评估

  • 算法评估（主题分类器）：训练的DistilBERT分类器在MMLU测试集上准确率为94.5%，在Chatbot Arena测试集上为85.0%，性能接近OpenAI O3-mini-high。其推理开销小于10ms，且定义的6个粗粒度主题覆盖了93%的真实世界查询，验证了其有效性和通用性。

  • 系统性能：如图16所示，在所有测试案例中，采用分层（tiering）的Stratum系统性能始终优于GPU基线。对于OLMoE、Mixtral、Qwen2.5和Llama-4模型，解码吞吐量平均分别提升了8.29倍、5.39倍、6.13倍和4.48倍。能效方面，分别提升高达7.66倍、2.74倍、3.51倍和4.87倍。与不分层（no-tiering）的Mono3D DRAM设计相比，分层和数据映射优化平均带来了1.3-1.45倍的吞吐量提升。与之前的工作Duplex【索引89】相比，Stratum也展现出最高达2.9倍的吞吐量和2.7倍的能效优势。


图16：系统解码吞吐量和能效的评估与比较。

  • 专家放置优化的有效性：图17显示，提高热点专家的命中率能显著提升性能。与均匀分布的专家使用情况相比，准确的预测能带来1.32倍至1.51倍的吞吐量提升。使用本文的主题预测模型，在服务Mixtral、OLMoE和Llama-4时，可实现31.6%、48.5%和68.9%的聚合热点专家命中率。


图17：热点专家命中率对(a) MLP (MoE层) 延迟和(b) Stratum-L整体系统吞吐量的影响。

  • 专家放置的开销：如表4所示，在批次之间进行专家交换的时间和能量开销都远低于1%，可以忽略不计。这是因为交换在bank内部进行，并由专用的行交换缓冲区支持，速度快且不经过DRAM-xPU接口。


表4：跨Mono3D DRAM层进行专家交换的开销

  • 随批次大小的性能扩展：如图18(a)所示，Stratum在不同批次大小下始终比GPU基线快4.7-9.8倍。但随着批次增大，相对性能优势有所减小，尤其是在短序列长度时，因为GPU裸片具有更高的计算带宽比。

  • 随Mono3D DRAM层数的性能扩展：如图18(b)所示，与256层和64层的设计相比，1024层的设计在单位面积吞吐量上分别高出1.21倍和2.96倍，显示了采用超过1000层Mono3D DRAM的成本效益优势。


图18：(a) 批次大小和(b) Mono3D DRAM层数对系统级指标的影响，评估使用Llama-4-Scout在Stratum-XL上进行。

  • 在较少层数Mono3D DRAM上的分层机制：即使在512层的配置下（最快和最慢层有1.3倍的延迟差异），所提出的主题感知分层放置策略依然能在多个基准测试上带来约18%的整体性能提升，验证了该策略的广泛适用性。

A7 补充细节

7 相关工作

3D可堆叠DRAM。单片3D可堆叠DRAM作为HBM的一种有前途的替代方案，通过在同一晶圆上顺序制造多个DRAM层而出现。与依赖TSV和昂贵的裸片堆叠的HBM不同，Mono3D DRAM采用精细间距的混合键合，以实现更高的内部带宽和集成密度【索引10, 11, 36, 46, 48, 83】。领先的Mono3D DRAM技术包括水平1T1C【索引36, 48】（重新定向并堆叠1T1C DRAM单元）和栅控晶闸管【索引10, 11】（利用雪崩机制）。最近的工作进一步表明，Mono3D DRAM约1µm的键合间距【索引9】使其垂直互连密度比HBM高出5倍【索引88】。

用于Transformer的内存内/近内存处理加速。尽管内存内/近内存处理（PIM/PNM）是一个由来已久的概念，但MAT【索引91】首次将PIM应用于Transformer模型，通过内存高效的流水线子序列流来处理单个编码器块。TransPIM【索引92】通过混合PIM-PNM架构将其扩展到全模型执行。Neupims【索引43】和AttAcc【索引67】专注于仅解码器Transformer模型，将解码阶段的注意力层卸载到xPU-PNM混合处理系统上的PNM。Duplex【索引89】进一步扩展了对MoE、GQA和连续批处理的支持，并采用了动态计算分区。然而，所有这些设计都依赖于2D DRAM或裸片堆叠的HBM，这限制了它们在基于Mono3D DRAM的系统中的应用效果。

A5 结论

本文提出了Stratum，一种用于高效MoE服务的新型系统-硬件协同设计。它首次利用了通过3D混合键合与逻辑单元集成的高密度Mono3D DRAM裸片，并进一步通过2.5D硅中介层与GPU连接。该架构为传统的基于GPU-HBM的系统提供了一种成本效益高且吞吐量大的替代方案。在硬件层面，Stratum引入了内存内分层机制，以利用Mono3D DRAM中的垂直访问延迟差异，并设计了一个为专家和注意力执行优化的近内存处理器（NMP）。在系统层面，我们利用与主题相关的专家激活模式，对专家进行分类并将其映射到不同的内存层，并设计了一个由轻量级分类器指导的主题感知调度器，以满足服务等级目标。跨越设备、电路、算法和系统层面的评估表明，与GPU基线相比，Stratum的解码吞吐量提高了高达8.29倍，能耗降低了高达7.66倍。