英伟达颠覆英伟达

为了解决这些效率低下的问题，由加州大学圣地亚哥分校、哥伦比亚大学、延世大学、英伟达和三星的研究人员发表了一篇新的技术论文，并提出了 AMMA（a multi-chiplet, memory-centric architecture for low-latency long-context attention），一种面向低延迟长上下文注意力机制的多芯片、内存中心架构。

AMMA 用 HBM-PNM 立方体取代了 GPU 计算芯片，使可用内存带宽大致翻倍，从而更好地满足内存密集型注意力工作负载的需求。为了将带宽提升转化为相应的性能提升，该团队引入了以下方案：(i) 一种逻辑芯片微架构，可在最小的功耗和面积预算下充分利用每个立方体的内部带宽进行注意力解码；(ii) 一种两级混合并行方案；(iii) 一种重排序集体流，可降低芯片内芯片间通信开销。

此外，该团队还对每个立方体的计算能力和芯片内 D2D 链路带宽进行了设计空间探索，为硬件设计人员提供了切实可行的指导。评估结果表明，与 NVIDIA H100 相比，AMMA 的注意力延迟降低了 15.5 倍，能耗降低了 6.9 倍。

现有的LLM服务系统是围绕以GPU为中心的范式构建的。如图 1（a）所示，诸如MegaScale-Infer 和Step-3 等生产系统采用了基于GPU的注意力-FFN分离方案，将注意力层和FFN（MoE）层分别放置在不同的GPU池上，以实现独立扩展。NVIDIA的下一代异构平台将这种分离扩展到硬件层面，将FFN层卸载到专用的LPU上，从而实现超低延迟，同时保留Rubin GPU用于注意力处理。即使是AttAcc和NeuPIMs等学术方案，尽管将内存处理（PIM）单元置于DRAM内部以加速解码注意力过程中内存密集型GEMV的执行，仍然依赖GPU作为跨设备通信的中心枢纽。无论在工业界还是学术界，GPU 仍然是注意力服务管道中无可争议的核心。

GPU 强大的计算能力架构与解码注意力机制受限于带宽的特性存在根本性的不匹配。随着推理和智能体工作负载将上下文长度推至数百万级别，这种不匹配变得日益严重，解码注意力延迟成为主要的服务瓶颈。例如，NVIDIA Rubin 为每字节内存带宽提供 795 FLOPs 的计算能力，而 GQA 注意力机制在 FP8 精度下仅需每字节 32 FLOPs 的计算能力，计算能力盈余高达 25 倍。Rubin 的两个计算芯片占据了封装面积的 67% 和设计功耗的 73%，但注意力机制仅利用了其峰值吞吐量的 4%。鉴于 LLM 服务需求已经给数据中心的电力预算带来巨大压力，这种低效是不可接受的。

这种不匹配促使我们对封装架构进行彻底的重新思考：为什么不完全用额外的 HBM 立方体替换 GPU 的计算芯片，并为每个立方体配备 PNM 功能呢？这样做可以将芯片内聚合的 HBM 带宽大致翻倍，同时将每个 HBM 立方体变成一个专用的 PNM 加速器。与之前 PIM/PNM 工作中使用的简单 GEMV 单元不同，这里的 PNM 单元是一个功能齐全的加速器，集成了数据加载、GEMM、GEMV 和芯片间通信，使得同一封装内的 HBM 立方体无需依赖主机 GPU 即可直接通信。因此，AMMA 是一个独立的、完全可编程的、以内存为中心的加速器，具有一流的系统性能，这标志着与目前以 GPU 为中心的模式的根本性转变。它还可以参与解耦服务，将 FFN 层交给 LPU 用于对延迟敏感的部署，或交给 GPU 用于对吞吐量要求高的部署，如图 1(a) 所示。

幸运的是，我们正处于一个技术转折点，这使得这一愿景在商业上成为可能。如图 1(b) 所示，从 HBM4 开始，存储器厂商采用先进工艺节点（≤5 nm）制造基础逻辑芯片，从而为专用微架构提供足够的晶体管密度，充分利用每个立方体的内部带宽。通过高速封装内 D2D 链路连接多个这样的立方体，即可形成专为长上下文注意力机制设计的节能型处理器。

然而，简单地用 HBM-PNM 立方体替换计算芯片并不能自动将额外的带宽转化为性能提升。逻辑芯片架构尚属未知，而 HBM 立方体固有的分布式特性使得充分利用封装内的大总带宽成为一项挑战。由此产生了四个具体挑战。

(C1) 计算内存比的根本差异需要一种全新的架构。由于目标工作负载受限于内存，数据重用极少，因此每个 HBM 立方体的逻辑芯片必须在严格的面积和功耗预算下充分利用其巨大的内部带宽。这与 GPU 的架构截然相反，在 GPU 中，计算能力远超带宽，并且深层的内存层次结构能够最大限度地提高数据重用率。因此，GPU 的架构原则无法直接应用于 HBM，需要一种截然不同的设计理念。

(C2) 并行策略引入了长距离数据传输。张量并行 (TP：Tensor Parallelism) 是降低延迟的标准多 GPU 方法，它将注意力分配到所有设备上。然而，如果简单地将 TP 应用于 16 个封装内的立方体，则会强制远距离立方体之间进行通信，而这些立方体之间需要经过多次网格跳跃，其开销很容易抵消聚合内存带宽的提升。因此，需要一种拓扑感知并行策略，将通信限制在局部范围内。

(C3) 多次集体操作会增加通信开销。即使采用合适的并行策略，传统的注意力流程也需要多轮 AllReduce 和 AllGather 操作，其累积延迟会成为显著的瓶颈，因此简化通信流程至关重要。

(C4) 关键硬件参数仍有待探索。在面积和功耗受限的情况下，两个关键的设计参数——单立方体计算吞吐量和芯片间链路带宽——会争夺共享的片上网络 (NoC) 资源。要找到实现最佳端到端注意力延迟的平衡点，需要对设计空间进行全面探索。

为了应对这些挑战，我们提出了AMMA，一种以多芯片内存为中心的架构，它突破了当前以GPU为中心的服务系统范式。如图1所示，AMMA将HBM和PNM置于核心位置，形成了一种新型的LLM加速器，并为未来高度异构、基于解耦的服务系统开辟了新的方向。具体而言，AMMA通过封装内的D2D链路将16个HBM芯片连接起来，形成一个单芯片。我们在每个HBM的逻辑芯片内设计了一个专用的微架构，以充分利用其内部带宽（C1）。此外，我们还引入了一种两级并行方案（C2）和一种重排序的集体通信流（C3），以降低D2D通信开销。最后，我们进行了全面的设计空间探索，为硬件设计人员提供可操作的指导（C4）。我们希望这项工作能够展示以内存为中心的架构作为GPU之外的独特架构类别的潜力，并启发人们对下一代异构平台进行更深入的研究，在这些平台上，以内存为中心的加速器将发挥核心作用。我们的贡献总结如下：

• 我们提出了一种面向低延迟注意力服务的多芯片内存中心架构，每个立方体在其逻辑芯片中集成了一个专用的微架构，以充分利用聚合的HBM带宽。

• 我们设计了一种两级混合并行方案，该方案将注意力映射到分布式立方体上，并限制通信范围以减少立方体之间的数据传输。

• 我们重新设计了集体通信流程，以减少集体操作的数量和全局同步开销，并为重新排序后的流程提供了形式化的正确性证明。

• 我们通过扫描每个立方体的计算吞吐量和芯片间链路带宽来探索硬件设计空间，分析它们对注意力延迟的影响，并为硬件设计人员提供实用指导。

每个注意力层包含四个操作。QKV 投影通过权重矩阵 B 和 C 将输入的隐藏状态A线性映射到查询、键和值，生成 HQ 查询头和 H KV 键值对。然后，对每个头计算核心注意力。

A：

B：

C：

其中 D 和 K，E。输出投影 F将连接后的头部输出映射回模型维度 Dm。分组查询注意力机制 (GQA) 减少了 KV 头部的数量，使得 G=HQ/HKV 个查询头部共享一个 KV 头部，从而降低了内存开销。GQA 在现代 LLM 中很常见。

D：

E：

F：

LLM推理分两个阶段进行。预填充阶段一次性处理整个提示信息，生成GEMM形状的计算，这种计算受限于计算资源，且对GPU友好。解码阶段自回归地逐个生成词元。在每个解码步骤中，新的查询会关注所有前S个词元，方法是从存储的KV缓存中读取它们的键值对，该缓存的大小为每层2HKVsdh。解码过程中只有一个查询词元时，公式1的算术强度会降低到32 FLOPs/byte（Qwen3 G=16，FP8），远低于现代GPU的计算带宽比。QKV和输出投影在较小的B值下也表现出同样低的算术强度，使得整个解码过程受限于内存。

对于推理和智能体工作负载，上下文长度可能达到百万级。因此，为了满足延迟目标（这对用户体验至关重要），批处理大小通常保持在适中水平。

如图 2 所示，将计算单元集成到 HBM 堆栈中有两种方式，各有优缺点。

内存内处理 (PIM：Processing in memory) 将算术单元直接放置在 DRAM 芯片上。通过将计算单元与各个存储体或存储体组共置，PIM 绕过了 TSV 总线瓶颈，并充分利用了 DRAM 阵列带宽，相比外部 HBM 接口，带宽提升了一个数量级。然而，计算单元会占用单元面积，直接降低存储容量。此外，DRAM 芯片是在成熟的工艺节点上制造的，这些节点针对单元密度（例如 1β、1γ、1c）进行了优化，而非逻辑性能，而严格的面积和功耗限制使得 DRAM 内计算仅限于简单的 GEMV 运算。三星 HBM2-PIM 和 SK 海力士 AIM 等商用产品采用了这种方法，但其有限的计算能力和较低的容量使其无法满足现代注意力机制工作负载的需求。

近内存处理 (PNM：Processing near memory) 将计算单元集成到逻辑芯片中，且不会造成容量损失。由于逻辑芯片的工艺节点独立于 DRAM 阵列，PNM 允许实现比 PIM 更复杂的逻辑，尤其是在 HBM4 及以后，逻辑芯片的制程将达到 ≤ 5 nm，从而为复杂的逻辑提供了足够的晶体管密度。PNM 本身并不能提高每个立方体的带宽，但用配备 PNM 的立方体替换 GPU 计算芯片，可以使封装的总带宽大致翻倍。然而，面积和功耗仍然是关键的限制因素：逻辑芯片已经容纳了内存控制器和 PHY，并且位于会散发大量热量的 DRAM 阵列下方。过高的功耗会导致过热和 DRAM 错误，因此必须仔细控制新增的计算能力，这也解释了下文中微架构选择的原因。

1、GPU 的局限性

大多数现有的 LLM 服务系统依赖 GPU 来解码注意力机制。然而，GPU 的设计目标是计算吞吐量，这使得它们并不适合内存密集型的注意力机制。如图 3(a) 所示，HBM 带宽利用率超过 90%，几乎达到饱和，而 GPU 计算单元的利用率却低于 5%，这造成了巨大的面积和功耗浪费。

图 4 中的屋顶线分析证实了这种不匹配。注意力内核的算术强度远低于 GPU 的计算带宽比，使其处于内存密集型状态，此时额外的计算能力并不能带来性能提升。基于此，AMMA 用 HBM-PNM 立方体取代了 GPU 计算单元，从而提供了更高的聚合内存带宽，展现出加速需要极长上下文的推理和智能体工作负载的巨大潜力。

2、GPU功耗分析

为了解GPU的功耗分配情况，我们对运行Qwen3-235B注意力机制、批大小为1的H100进行了功耗分析。我们使用NVML进行测量，并使用CACTI进行建模。图3(b)的结果揭示了两个反直觉的观察结果。

即使在内存密集型工作负载下，GPU的功耗也会达到TDP。长上下文注意力机制仅使用5%的可用计算能力，因此预期功耗远低于TDP。然而，实际测量结果显示平均功耗为693W，与H100的700W TDP相差不到1%。考虑到静态功耗为125W，HBM总功耗低于200W，大约还有370W的功耗无法解释。我们的分析表明，功耗主要来自片上网络（NoC）、逻辑控制器（LLC）和非计算型流水线组件（SM）逻辑。随着上下文长度的增加，越来越多的SM被激活用于并行处理，它们的功耗也随之增长。虽然它们的计算单元处于空闲状态，但周围的数据移动和控制架构仍然需要供电。由于AMMA的功耗远低于独立GPU，因此采用类似GPU的微架构是不可行的。

逻辑控制器消耗大量功耗，但性能贡献却很小。在长时间上下文注意力任务中，仅逻辑控制器就消耗130瓦的功耗，同时几乎达到100%的未命中率。这是因为工作集远远超过了逻辑控制器的容量，而低运算强度也限制了数据重用的机会。寄存器文件和共享内存也存在同样的问题，它们原本是为了利用注意力工作负载根本不具备的数据重用能力而设计的，因此它们反而成了面积和功耗的负担，而不是性能优势。

这些发现证实，GPU微架构的惯例并不适用于AMMA。因此，采用专用的逻辑芯片架构，将其面积和功率预算用于充分利用 HBM 带宽，是必不可少的。

3、先前 PIM/PNM 方案的局限性

先前的 PIM/PNM 方案存在两个根本性的局限性，使其不适用于现代 GQA/MLA 注意力机制。

以 GPU 为中心的范式引入了通信开销。现有方案，例如 AttAcc将 GPU 视为立方体间通信的中心设备。这对于 MHA 来说效果良好，因为足够的键值对 (KV) 头允许每个 PIM 设备通过 TP 独立处理一个键值对。然而，GQA 和 MLA 将键值对压缩了 16-128 倍。

因此，仅靠 TP 无法将工作分配到各个 PIM 设备（例如，64 个设备 vs. 4 个键值对），从而迫使在序列维度上进行上下文并行 (CP)。CP 引入了 AllReduce，这对以 GPU 为中心的 PIM 来说是灾难性的。如图 5 所示，AttAcc 通过 NVLink 连接了 8 个 GPU 和 64 个 PIM 设备。在处理 AllReduce 时，每个 GPU 必须先从 8 个 PIM 设备收集数据，然后执行跨 GPU 归约。仅此通信开销就比 64K 数据量下的注意力计算耗时 1.5 倍，一旦增加足够的计算能力来缓解下文讨论的计算瓶颈，耗时将增加到 24 倍。这些结果还假设 NVLink 的理想延迟为 900 ns，且无内核启动开销，而我们的实际 GPU 分析表明，单次 8B 数据传输的端到端耗时超过 12,000 ns。

AMMA 通过摒弃以 GPU 为中心的范式来消除这一瓶颈。每个 HBM 立方体都拥有自己的逻辑芯片，并作为一个完全独立的加速器运行。多个立方体通过高速 D2D 链路互连，形成一个独立的、无需主机 GPU 的多芯片注意力处理器。

计算能力不足是 GQA 的瓶颈。大多数先前的 PIM 设计将 GEMV 单元放置在 DRAM 芯片中，这足以满足 MHA 的需求，但不足以满足 GQA 的需求。如图 4 所示，GQA 通过 KV 磁头共享将算术强度提高了 16-32 倍，使得之前的架构受限于计算能力。图 5 证实，尽管 GQA 比 MHA 减少了 16 倍的内存流量，但由于这一瓶颈，总注意力延迟并未得到改善。这种转变需要尺寸远超 DRAM 芯片容量的 GEMM 单元，从而推动了向逻辑芯片 PNM 的迁移。Duplex探索了这一方向，但其约 28 纳米的逻辑芯片工艺严重限制了设计空间。我们提出 AMMA 正值这一转折点，先进的逻辑芯片工艺使得复杂的微架构成为可能，并释放了以内存为中心的加速的新潜力。

1、概述

AMMA 架构如图 6 所示。每个芯片封装将 16 个 HBM-NMP 立方体排列成 4×4 的二维网格，每个立方体集成了一个 HBM 堆栈和一个支持 PNM 的逻辑芯片，并通过高速 D2D 链路进行通信。由于注意力工作负载在访问模式和资源平衡方面与传统的 GPU 工作负载存在根本差异，AMMA 摒弃了 GPU 式的设计，采用了一种遵循三个原则的微架构。

(P1) 大量小型 SA（systolic arrays ） 而非少量大型 SA。解码注意力的 M 维度很小（1-32，由批处理大小或 GQA 中每个 KV 头的 Q 头数量决定），因此 TPU 式的 128×128 阵列会导致超过 87% 的 PE 闲置。我们转而部署 96 个 16×16 脉动阵列，其大小与工作负载较小的 M 值相匹配，确保每个 PE 都被充分利用，同时提供相同的总吞吐量。我们选择脉动阵列（SA）而非向量单元，是因为 SA 可以在 16 个 PE 中复用每个操作数，从而将所需的 SRAM 读取带宽降低 16 倍，进而显著降低面积和功耗开销。

（P2）无 LLC 架构。从架构和工作负载的角度来看，LLC 都是不必要的。从架构角度来看，GPU LLC 可以提升有效带宽，从而弥合计算带宽和内存带宽之间的差距。然而，AMMA 的计算规模与原始 HBM 带宽相匹配，自然消除了这一差距。从工作负载角度来看，LLC 利用数据复用来减少冗余的 DRAM 访问，但在低延迟解码注意力机制下，小批量大小（B=1–32）时，键值缓存和输入请求仅通过一次，几乎没有可利用的复用空间。移除 LLC 可以节省 20% 的功耗预算和大量的芯片面积，我们可以将这些资源重新分配给有用的计算。

（P3）用于数据共享的两级交叉开关。移除 LLC 也消除了其在核心间进行数据交换的功能。我们用一个两级交叉开关来替代它，该交叉开关广播共享输入并收集96个SA的部分输出。在输入端，所有处理不同图块的SA都需要相同的查询向量，而交叉开关无需冗余的DRAM读取即可提供该查询向量。在输出端，来自同一SA下多个SA的部分和通过交叉开关收集，从而避免了昂贵的DRAM流量。我们采用分层结构组织交叉开关，每个8SA核心内有一个本地交叉开关，并在12个核心上有一个全局交叉开关，与所有SA的扁平化设计相比，面积从O(N²)降低到O(N)。

2、微架构：立方体和核心

16 个立方体构成一个二维网格，每个立方体通过芯片间 (D2D) 链路与最多四个相邻立方体相连。如图 6(a) 所示，每个立方体内部，DRAM 芯片位于逻辑芯片之上，并通过微凸点或混合键合焊盘连接。逻辑芯片在其标准区域内包含 HBM PHY 和内存控制器，而我们添加的计算逻辑用于实现 PNM（部分和最小化）。

如图 6(b) 所示，每个立方体包含 12 个计算核心、一个内存控制器和四个 D2D 端口，所有这些都通过片上网络 (NoC) 连接。这 12 个核心通过 P3 中描述的两级交叉开关进一步互连，从而实现跨核心的查询广播（这些核心计算不同的 N-tile）以及跨核心的部分和归约。

每个核心包含一个指令前端、一个 DMA 引擎、两个输入缓冲区组、一个包含 8 个 SA（状态代理）的集群、一个输出缓冲区和一个向量单元。 DMA引擎直接从HBM将数据流式传输到每个核心的缓冲区。输入缓冲区A保存A矩阵的行，默认情况下将每个切片馈送到其对应的SA，但当出现数据重用机会时，也可以通过本地交叉开关将共享行广播到所有8个SA。分布式输入缓冲区B在SA集群上进行分区，以便每个SA都能以最小的线路长度接收其自身的B矩阵切片。两个输入缓冲区都采用双缓冲设计，以隐藏200纳秒的HBM访问延迟，DMA引擎会在SA计算当前块时预取下一个块。每个核心总共有32KB的SRAM（输入缓冲区A为2×6.4KB，分布式输入缓冲区B为2×6.4KB，输出缓冲区为6.4KB），每个立方体3MB，所有16个立方体共48MB。这比 H100 小两倍多，比 Rubin 小六倍多，功耗和面积也相应降低。向量单元处理逐元素的后 SA 操作，例如 softmax 和线性范数，它从输出缓冲区读取数据，并以流水线方式运行，与下一个图块的 SA 计算重叠。

3、脉动阵列数据流

AMMA 部署了 96 个小型 16×16 的 SA，它们并行运行，且 M 维度较小（M=1–16，由批处理大小或 GQA 中每个 KV 头的 Q 头数量决定）。TPU 式架构将几个大型 128×128 的 SA 与权重静态数据流配对，依靠大批工作负载来提供较大的 M 值，从而使每个 PE 保持繁忙状态。我们采用的大量小型阵列和较小的 M 值，需要不同的数据流和分块策略。

如图 7(a) 所示，每个规范的 SA 数据流保持不同的操作数静态，并将一个 GEMM 维度的数据流通过阵列，因此该维度必须足够大才能实现高利用率：WS 流（streams） M，IS 流 N，OS 流 K。

我们为 AMMA 选择 OS 数据流，因为它是唯一适合我们长上下文、小 M 值环境的数据流。由于仅通过 16×16 权重块流式传输 M=1-16 行数据会浪费大部分 PE 周期，因此 WS 方案立即被排除。IS 方案看似有吸引力，因为序列长度 N 本身就很大，但在 IS 方案下，每个 SA 都会累积一个随 N 增长的部分输出，而跨越 96 个并行 SA 收集这些部分输出会产生 O(N) 的通信开销，并且随着序列长度的增加而线性恶化。OS 方案避免了这两个问题，因为 K 足够大，可以进行流式传输，并且每个完成的输出块都是一个固定的 16×16 块，使得跨 SA 收集的成本与序列长度无关。

4、高利用率的分块策略

在操作系统 (OS) 建立之后，剩下的问题是如何将 GEMM 的维度分块到 96 个空间阵列 (SA) 上以最大化利用率。我们采用两种技术来解决这个问题：一种是将 K 和 N 划分到各个 SA 上的分块方法，另一种是连续分块方法，该方法在每个 SA 内流水线化处理连续的图块，以消除图块间的空闲周期。

考虑一个 GEMM 𝐂=𝐀𝐁，其形状分别为 M×K 和 K×N。我们将 N 分块为 N/16 个列图块，并可选择将 K 分块为深度为 k=K/SK 的 SK 个段，从而得到 T=SK⋅N/16 个大小为 16×16 的输出图块。令 P=96 表示可用 SA 的数量。总体利用率可以分解为两个因素的乘积：

其中第一个因子统计了 96 个 SA 中有多少处于忙碌状态，第二个因子则反映了填充/排空管道的开销，即 M SA−1=15 个周期，分别对应每个 tile 执行周期的开始和结束。

分块方法：将公式 2 扩展到 T≤P 的情况（并非所有 SA 都处于忙碌状态），可以揭示一个关键信息：

分子由问题规模决定，与分块选择无关。分母随着 k 的减小而缩小，因此当 SA 未充分利用时，较小的单块深度可以带来更高的利用率。拆分 K 会创建更多块，激活空闲的 SA，并行性的提升超过了单 SA 效率的损失。一旦所有 SA 都达到饱和（T≥P），第一个因子锁定为 1，利用率变为 k/(k+30)，并随 k 的增大而增长。进一步拆分只会降低利用率。由此得出一个简单的原则：拆分 K 到足以使每个 SA 至少获得一个块的程度，然后停止。

我们在图 7(b) 中用两个例子 (P=96) 说明了这一点。当 N=768 时，只有 768/16=48 个列块，少于 96 个 SA。设置 SK=2 使块数翻倍至 T=96，使每个 SA 恰好获得一个块，同时保持单块深度 k=K/2 远高于 MSA。当 N=3072 时，共有 3072/16=192 个列图块，已超过 96。无需进行 K 分割 (SK=1)，每个 SA 以全深度 k=K 处理两个连续的图块。

连续分块：当一个 SA 处理多个连续的图块时，默认调度会串行执行一个图块的清空和下一个图块的填充，导致在 15 个周期的清空阶段，逐渐释放的 PE 处于空闲状态（图 7(c)，“默认”）。我们提出的连续分块方法通过立即将下一个图块的数据馈送到当前图块刚刚释放的 PE 中来消除这种浪费。当图块 i 从阵列的一端清空时，图块 i+1 从另一端填充，两个阶段完全重叠。双缓冲 DMA 预取确保操作数数据在每次清空开始之前都已准备就绪。

如果一个 SA 处理 n 个连续的图块，则填充/排空开销为2(MSA−1) 个周期，并且整个运行过程只需支付一次，而不是每个图块支付一次：

对于 N=3072 且每个 SA 处理 n=2 个图块的情况，连续平铺可将每个 SA 的效率从 52% 提高到 67%（当 k=32 时），当 n=4 时则提高到 81%。在每个 SA 处理数十个图块的更长场景中，所有 SA 的利用率都接近完美。

上一节描述了每个立方体内的计算工作流程。现在，我们将探讨如何在单个芯片内的所有 16 个 HBM-NMP 立方体上分配完整的注意力流水线。

1、默认的类 GPU 并行

现有的多 GPU 系统也解决了类似的并行设计问题，但它们根本不同的约束和目标使得它们的解决方案不适用于 AMMA。GPU 集群通常将 TP 度设置为低于 KV 头的数量，并依靠数据并行 (DP) 来扩展跨请求的聚合吞吐量（>10,000）。TP 度保持在低于头数，是为了确保每个 TP 组至少覆盖一个完整的 KV 头，因为将一个头拆分到多个 TP 设备上需要在每个注意力阶段之后执行 AllReduce，这会在高 TP 度下显著增加端到端延迟。

DP 和 TP 都无法迁移到 AMMA。 AMMA 的设计目标是在少量请求（1-32 个）上实现低延迟，而非在多个并发请求上实现聚合吞吐量，因此 DP 对此并无益处。为了满足低延迟目标，所有 16 个立方体必须同时协作处理同一个注意力请求，这就使得在所有立方体上高效地分配权重和键值缓存至关重要。简单地在所有 16 个立方体上应用 TP（图 8(a)）会在每个注意力阶段触发 16 路 AllGather 和 AllReduce 操作，从而将通信扩展到整个网格直径。更糟糕的是，TP16 要求数据传输量与序列维度成正比，考虑到 1M 的上下文长度，这将显著增加延迟。

2、两级混合并行

为了解决这个问题，我们提出了一种两级混合并行方案，如图 8(b) 所示。为了清晰起见，我们使用 4 个 KV 头和 4 个 Q 头的 MHA 配置进行演示，但该设计可推广到任意数量的头和注意力机制，例如 GQA。我们首先根据 KV 头的数量将 16 个立方体划分为若干组。对于 4 个 KV 头，这会产生 4 个立方体组，每组 4 个立方体，其中每个组的形状为 2×2 子网格，而不是 1×4 条带。

这样做是为了最大限度地减少组内的最大跳数，并更好地支持集体操作。

第一层：使用 TP 将 KV 头映射到立方体组。每个立方体组分配一个 KV 头，并存储相应的投影权重和 KV 缓存条目。查询也相应地进行划分，每个组处理与其分配的 KV 头关联的所有 Q 头。然后，这四个组并行运行以覆盖所有头。由于在 QKV 投影和注意力阶段，不同的 KV 头自然是相互独立的，因此这种分配方式避免了在主要的 KV 缓存读取过程中发生任何跨组通信，从而在输出投影结束时仅需在立方体组之间执行一次 AllReduce 操作。

第二级：使用 CP 在每个组内拆分 KV 缓存。在每个四立方体组内，各个立方体协作计算一个 KV 头的注意力。为了充分利用聚合的 HBM 带宽，我们将 KV 缓存沿序列维度拆分到四个立方体上。在解码过程中，查询在组内广播，每个立方体计算其本地序列分片的注意力，然后组将部分结果合并以形成每个头的最终输出。这引入了一次 AllGather 和一次 AllReduce 操作，但与朴素 TP-16 中的 16 路集合不同，这两者都被限制在四立方体子网格内，通信量与序列长度无关，从而大大降低了开销。

这两个层级共同构成了一种TP+CP混合模式，它将大部分数据传输限制在局部邻域内，确保跨立方体通信不会随序列长度而扩展，同时还能利用所有16个立方体进行聚合带宽。这种TP+CP组合在GPU集群上的分布式训练中很常见，但在推理过程中却很少使用，因为CP操作需要在通过NVLink连接的独立芯片之间进行，而芯片间的延迟使其不适用于对延迟敏感的单请求服务。然而，在AMMA架构中，同一组内的立方体位于同一芯片上，并通过D2D链路进行通信，延迟更低，带宽更高，这使得CP在推理阶段既实用又有利可图。

第 5 节介绍的两级混合并行机制将每个集体操作限制在局部邻居之间，从而消除了与序列长度相关的流量。现在，我们更进一步，通过重新设计输出投影阶段周围的通信流程，减少关键路径上的集体操作数量。

1、重排序的集体流程

图 9 展示了我们提出的两级混合并行机制的默认集体操作流程和改进后的流程。本节通篇使用以下符号：上标 m 表示立方体组（一级索引），n 表示该组内的立方体（二级索引）。

Smn 是一个在两级都进行切片的张量，而 Sm(n) 表示仅在组级别切片的立方体 n 的部分和。对于权重矩阵，WOmn[x] 表示一级分区沿 y（输入）维度，二级分区沿 x（输出）维度。图中，阴影方块表示需要跨立方体进行归约的部分和，实心方块表示跨立方体重复的数据。

默认的集体处理流程如图 9(a) 所示。在核心注意力计算之后，组 m 中的每个立方体 n 都包含一个部分和激活值 Am(n)，该激活值需要在其组内进行归约。默认方法应用完整的组内 AllReduce（先执行 ReduceScatter，再执行 AllGather）来重建每个立方体的完整注意力输出 Am。然后，每个立方体将 Am 乘以其局部权重分片 WOmn[yx]。最后，跨组的第二次 AllReduce 将每个头的结果累加到最终输出中。

此默认流程包含两个冗余来源。首先，组内 AllGather 将完整的注意力输出复制到每个立方体，而随后的跨组集体处理会立即分散并重新归约这些输出。其次，跨组 AllReduce 会将最终结果广播到包中的每个立方体，即使 AMMA 采用的是解耦服务模型（图 1），其中只需要一个目标立方体将结果转发到远程 FFN 加速器。

我们通过两项协调的改进来解决这两个效率低下的问题，如图 9(b) 所示。首先，我们将组内 AllReduce 替换为单独的 ReduceScatter，这样每个立方体 n 就只保留一个独立的散射切片 Amn，而不是完整的输出。为了匹配这种更窄的输入，我们将投影权重从 WOmn[yx] 重新切片，

到 WOmn[yy]，将二级分区轴从输出维度切换到输入维度。然后，每个立方体独立计算 Amn WOmn[yy]，生成一个部分和输出切片 O(m)(n)，没有权重重复，这也消除了默认流程中所需的投影后 AllGather。其次，我们将跨组的 AllReduce 操作替换为指向单个目标多维数据集的点对点 Reduce 操作，从而节省了大约一半的流量。这些更改总共移除了两次 AllGather 操作，并将一次 AllReduce 操作降级为 Reduce 操作，同时保持每个多维数据集的计算和内存占用不变。

方法论。我们结合仿真和真实 GPU 性能分析来评估 AMMA 的性能。对于单立方体性能，我们使用开源的脉动阵列模拟器 ScaleSim（Samajdar 等人，2018）对每个 HBM-NMP 立方体进行建模。对于多立方体性能，我们使用 AstraSim（Rashidi 等人，2020；Won 等人，2023），将其默认的 GPU 参数替换为我们自定义的 HBM-NMP 立方体规格，并导入 ScaleSim 生成的每个立方体的测试结果。对于 GPU 基准测试，我们在 8 台 H100 服务器上，针对不同的批处理大小和序列长度，收集端到端延迟。由于 Rubin 尚未公开，我们通过将 H100 的测量结果与 Rubin 已发布的带宽和计算规格进行缩放，同时保持测量的利用率不变，来预测其性能。

指标。我们使用延迟作为主要性能指标，因为它直接反映了面向用户的服务质量。我们还报告了能效（Token/J）和功耗（W），这两项指标是数据中心运营商最为关注的。

硬件配置。表 1 总结了我们评估中使用的硬件参数。对于 NuePIM（Heo 等人，2024），为了公平比较，我们将 NPU 参数缩放至 Rubin，将 PIM 参数缩放至 HBM4。配备 DRAM 芯片上 GEMV 单元的 on-bank PIM 提供的带宽是 HBM 接口的 9 倍，但计算能力有限。对于 AMMA，我们根据 UCIe 3.0 协议（Universal Chiplet Interconnect Express, Inc.，2025）将立方体间的 D2D 延迟建模为固定的每跳 15 ns。每个立方体使用 96 个 16×16 脉动阵列，运行频率为 2 GHz，可提供 96 TFLOPS 的计算能力，整个芯片的总吞吐量为 1536 TFLOPS。我们还将其与映射策略的变体进行了比较，所有变体均在 AMMA 架构上实现。我们使用 TP16 表示 HBM 立方体之间默认的类 GPU 张量并行性，HP 表示仅使用两级混合映射，HP_RO 表示结合混合映射和重排序调度的完整设计。

模型和工作负载。我们在 GQA 和 MLA 模型上进行评估：GQA 模型使用 Qwen3-235B 和 Llama4-Maverick，MLA 模型使用 DeepSeek-V3。我们建模的工作负载包括 QKV 投影、核心注意力机制和输出投影。我们排除了 MoE 操作，因为当前系统和新兴工业设计都表现出明显的注意力-MoE 分离趋势，其中 LPU 或 GPU 等专用硬件处理 MoE，而 AMMA 则在异构系统中充当注意力引擎。

1、延迟

GQA 模型测试结果：图 10 展示了两种 GQA 模型 Qwen 和 Llama 的逐层解码延迟加速比，所有结果均以单个 H100 模型为基准进行归一化。AMMA 模型在所有配置下均实现了最高的加速比：在 BS=1 时比 H100 模型快 12.0–16.3 倍，在 BS=32 时仍保持 13–20 倍的加速比。

这得益于每个封装集成了更多的 HBM4 立方体（16 个 HBM3e 立方体对比 5 个 HBM3e 立方体）以及聚合了 40 TB/s 的带宽——是 H100 模型的 11.9 倍。

与 Rubin 模型（8 个 HBM4 立方体）相比，我们保持了稳定的 1.8–2.5 倍领先优势。即使与带宽与我们相同的 Rubin TP2 相比，我们在短到中等序列中仍保持 1.5–2.4 倍的性能优势，仅在 100 万个令牌时缩小至 1.1 倍，这是因为较大的 KV 缓存容量分摊了 Rubin 的芯片间 NVLink 开销。尽管在 100 万个令牌时差距不大，但我们的设计在能效方面仍提高了 2.8 倍（图 11）。

与 PIM 基准 NeuPIM 相比，我们在 Qwen3-235B 上的性能优势为 3.4 倍，在 Llama4-Maverick 上的性能优势为 1.4 倍。NeuPIM 仍然依赖 GPU 来处理投影层，因此其性能受限于内存，并且其 bank 级 PIM 单元虽然提供了较高的原始聚合带宽（比我们的高 4.5 倍），但计算吞吐量不足，使得 GQA 成为计算密集型算法，并浪费了大部分带宽。

MLA 模型上的结果。在采用 MLA 的 DeepSeek V3 上，我们观察到了截然不同的趋势。在短序列（4K）下，投影算法占据主导地位，AMMA 凭借其带宽优势，性能比 Rubin 高出 1.9 倍。然而，随着序列长度的增加，注意力机制逐渐占据主导地位，Rubin 的性能比 AMMA 高出 2.9 倍。

MLA 的算术强度大约是标准 GQA 的 8 倍，这使得在 AMMA 的单立方体计算预算下，注意力机制的性能受限于计算能力。将单立方体计算能力提升至 512 TFLOPS 后，AMMA 在所有配置下的领先优势恢复到 1.8-2.1 倍，这表明单芯片 MLA 部署需要更高的计算密度。实际上，可以将面积和功耗预算从 D2D 带宽重新分配给额外的计算单元。

值得注意的是，当 MLA 部署在多个设备上时，计算不再是瓶颈。由于 DeepSeek V3 只有一个 KV 头，因此像 SGLang 这样的框架会将 Q 头拆分，并在 TP2-4 模式下将 KV 缓存复制到多个设备上。这降低了每个设备的算术强度，使其完全处于 AMMA 的计算带宽最佳范围内，并保持了我们的延迟优势。

2、能量与功耗

图 11 上半部分显示，AMMA 的能量效率（Token/J）始终比 H100 高 5.6 至 6.6 倍，并且比 Rubin 稳定高出 2.6 至 3.1 倍。与 Rubin TP2 相比，差距在 1M 时为 2.8 倍，在 4K 时为 4.8 倍，这是因为 TP2 的静态功耗翻倍，且 NVLink 开销导致短序列的效率降低。这一结果主要归功于两个因素：我们专用的微架构仅提供 Rubin 10% 的计算吞吐量，从而大幅降低了静态功耗；此外，我们采用无 LLC 设计并最大限度地减少了片上 SRAM，从而消除了冗余数据传输，降低了动态功耗。

图 11 下半部分从功耗角度证实了这些优势：AMMA 的功耗低于 1500 瓦——比 Rubin 低 1.4 倍，比 Rubin TP2 低 2.6 倍——同时仍然提供比两者都更高的绝对性能。

3、消融实验

如图 12(a) 所示，HP_RO 相对于 TP16 基线始终实现了最高的加速比，在 8K、256K 和 1M 序列上分别达到 1.1 倍、1.5 倍和 1.6 倍。HP 的性能也优于 TP16，但增益略低（分别为 1.1 倍、1.4 倍和 1.5 倍）。由于 TP16 引入了与序列长度成正比的通信量，因此 HP_RO 相对于 TP16 的改进随着序列长度的增加而增大。与 HP 相比，HP_RO 在短序列上表现出更显著的优势，因为 RO 带来的固定延迟降低会随着序列长度的增加而被注意力计算逐渐稀释。

由于大部分性能提升源于通信量的减少，图 12(b) 仅展示了通信加速比。与单独使用 HP 相比，HP_RO 在 8K、256K 和 1M 时分别实现了 2.7 倍、17.7 倍和 65.4 倍的通信速度提升，而单独使用 HP 时则分别为 1.3 倍、10.0 倍和 36.8 倍。

4、批次大小探索

由于低延迟服务通常采用 1 到 32 的批次大小，我们研究了批次大小如何影响延迟与吞吐量之间的权衡。如图 14(a) 和 (b) 所示，当批次大小从 1 增加到 32 时，系统总吞吐量从 0.223 tok/μs 提高到 0.478 tok/μs (2.14 倍)，但延迟也增加了 30 倍。吞吐量的提高是因为批处理增加了投影操作的运算强度，从而可以更充分地利用 AMMA 的计算能力。然而，延迟增加 30 倍的原因在于注意力机制无法从批处理中受益，因为每个请求的查询仅关注其自身的键值缓存。因此，整个批次的注意力计算必须逐个请求处理，并且输出投影只能在所有请求的注意力计算完成后才能开始，从而导致更长的延迟。

图 14(a) 中红线的趋于平缓的趋势进一步表明，当批处理大小达到 16 时，AMMA 的计算能力已经饱和。超过此点后，继续增加批处理大小不会带来吞吐量的进一步提升。这与 GPU 有着本质区别，GPU 通常需要 200-400 的批处理大小才能充分利用所有可用的计算能力。

为了更好地展现 AMMA 的整体性能，我们在图 12(c) 中展示了帕累托前沿，并将其与 NVIDIA GPU 进行了比较。结果表明，对于长上下文注意力工作负载，无论批处理大小如何，AMMA 的性能始终优于 GPU，在功耗相当或更低的情况下，其吞吐量比 H100 高出 14-16 倍。

5、时间分解

我们在图 13 中提供了逐层的解码延迟分解，以分析 Proj QKV、注意力机制、Proj O 和通信机制各自对总延迟的贡献。

在短序列中，投影操作占据了主要的计算时间：当 BS=1 且 Seq=8K 时，Proj QKV 和 Proj O 合计占 AMMA(HP+RO) 计算量的 85%，而注意力机制仅占 15%。随着序列长度增长到 128K，注意力机制的计算量也成比例增长并占据主导地位，在 BS=1 时达到 60%，在 BS=4 时达到 86%，这是因为每个批次项的键值缓存必须独立流式传输。通过将权重和键值缓存分布到 16 个 HBM 立方体上并采用 PNM，AMMA 实现了 40,000 TB/s 的总带宽（Rubin 为 8,000 TB/s），Proj QKV 的计算量减少了 1.76 倍，Proj O 减少了 2.57 倍，注意力机制减少了 3.89 倍（与 Rubin 相比）。

6、硬件参数探索

我们探索了每个立方体的计算能力和 D2D 带宽如何影响总延迟，为硬件设计人员在逻辑芯片上权衡面积和功耗提供了指导。我们将计算能力从 8 到 256 TFLOPS 进行扫描，并将 D2D 链路带宽从 0.5 到 2.5 TB/s 进行扫描。结果如图 15 所示，其中蓝色表示低延迟，红色表示高延迟。

我们的探索表明，计算能力比 D2D 带宽更为关键，原因有二。首先，我们的混合并行和重排序集合已经最大限度地减少了立方体之间的通信，因此带宽提升的空间有限。其次，在小批量大小（1-32）下，D2D 传输量较小，延迟主要由固定的启动延迟而非传输时间决定，因此提高链路带宽带来的收益甚微。

我们还观察到明显的计算能力饱和上限。在 Qwen3 工作负载下，每个立方体的计算能力超过 96,000 TFLOPS 后，性能几乎没有提升，因为此时所有操作（包括投影和注意力机制）都受限于内存。在这种情况下，将剩余的功率预算重新分配给更高的 HBM 频率比增加计算单元更为有效。

结论

我们提出了 AMMA，一种多芯片 HBM-PNM 架构，它用支持 PNM 的 HBM 立方体取代了 GPU 计算芯片，作为独立的、以内存为中心的长上下文解码注意力加速器。通过协同设计专用逻辑芯片微架构、两级混合并行以及重排序的集体通信，AMMA 充分利用了聚合的内部 HBM 带宽，同时最大限度地减少了立方体之间的数据传输。我们的设计空间探索进一步揭示了在面积和功耗约束下，计算带宽和 D2D 带宽之间可行的权衡方案。我们希望这项工作能够鼓励业界认识到以内存为中心的架构在未来异构和基于解耦的服务系统中的重要性和潜力。

*免责声明：本文由作者原创。文章内容系作者个人观点，半导体行业观察转载仅为了传达一种不同的观点，不代表半导体行业观察对该观点赞同或支持，如果有任何异议，欢迎联系半导体行业观察。