RED-PIM: 使用处理内存减少Transformer的数据移动

# RED-PIM: 利用处理中内存减少Transformer的数据移动
来源: https://arxiv.org/html/2607.21731
###### 摘要

Transformer 广泛应用于许多领域，包括自然语言处理、计算机视觉、网页搜索和DNA序列分析。鉴于其广泛适用性，提升Transformer模型的性能至关重要。然而，在注意力操作期间，处理单元与内存之间的大量数据移动严重限制了其效率。处理中内存（PIM）通过在内存内部直接执行计算来缓解这一问题。虽然先前的工作提出了基于PIM的Transformer实现，但它们面临着昂贵的bank间通信成本，并且由于内存bank容量有限而难以扩展。因此，注意力相关数据必须跨bank分割，削弱了PIM的潜在优势。

在这项工作中，我们提出了RED-PIM，一种算法-架构协同设计，通过将bank间数据移动从O\(N^2\)减少到O\(N\)，并将中间注意力矩阵从N×N缩小到d×d，从而降低注意力延迟。通过重新组织矩阵操作、本地执行计算以及采用优化的数据传输策略，RED-PIM显著降低了计算成本和互连流量。与基线PIM实现相比，RED-PIM实现了16.05%到99.99%的推理时间减少（几何平均值为66.42%），其中较长序列的增益最大。在真实数据集上，RED-PIM对长文档的性能提升为99.60%，对短文档为13.44%，同时保持或提高了准确性。这些结果证明了RED-PIM在可扩展且高效的Transformer推理中的有效性。

## I 引言

传统的CPU和GPU架构在满足大规模神经网络的计算需求方面面临重大挑战，这是由于内存与处理单元之间的数据移动导致的低效率。这种大量数据的来回传输造成了瓶颈，降低了性能并增加了能耗。

在各种深度学习模型中，这个问题在基于Transformer的架构[42 (https://arxiv.org/html/2607.21731#bib.bib1)]中更为显著，主要是因为它们的自注意力模块。原因是自注意力涉及计算所有输入token对之间的注意力分数。为此，在从输入token初始计算查询、键和值向量之后，模型必须计算所有查询与键之间的相似度以生成注意力图。这需要O\(N^2\)次数据移动用于（查询，键）比较，其中N是序列长度。随后，将此注意力图与值向量相乘以获得最终结果也需要O\(N^2\)次移动。

由于大规模Transformer所需的数据量很大，典型的CPU或GPU缓存效率低下。自注意力计算生成极其庞大的中间数据结构，远远超出典型缓存的容量，导致频繁的缓存逐出和延迟增加。举例来说，考虑一个具有8,000个token的序列。生成的注意力图由6400万个值组成，即使使用最低精度格式（即FP16）也需要约32MB内存，远大于典型CPU或GPU缓存的容量。

处理/近内存（PIM）[48 (https://arxiv.org/html/2607.21731#bib.bib10),49 (https://arxiv.org/html/2607.21731#bib.bib11),40 (https://arxiv.org/html/2607.21731#bib.bib15),23 (https://arxiv.org/html/2607.21731#bib.bib16),27 (https://arxiv.org/html/2607.21731#bib.bib17),26 (https://arxiv.org/html/2607.21731#bib.bib18),18 (https://arxiv.org/html/2607.21731#bib.bib20),25 (https://arxiv.org/html/2607.21731#bib.bib21)]提供了一种替代方案，以减少内存与处理单元之间的数据移动。PIM通过将计算直接集成到内存阵列中来解决这一挑战，从而消除了处理器与内存之间频繁数据传输的需要。然而，即使有了PIM，内存bank的有限容量仍然是扩展Transformer模型的主要障碍。例如，一个典型的HBM bank只有32MB存储容量，在FP16精度下最多支持8,000个token。相比之下，现代大型语言模型（LLMs），如GPT‑4.1[31 (https://arxiv.org/html/2607.21731#bib.bib56)]、Gemini 2[7 (https://arxiv.org/html/2607.21731#bib.bib57)]和Grok‑3[47 (https://arxiv.org/html/2607.21731#bib.bib58)]，支持更长的序列，最高可达100万个token。尝试将完整的注意力图存储在多个bank中会引入显著的bank间通信开销，这削弱了PIM旨在提供的性能和能效优势。

在本文中，我们提出RED-PIM，通过利用PIM并减少bank间数据移动来降低注意力操作的延迟。RED-PIM通过采用一种替代的注意力计算算法来实现这一点，该算法消除了构建完整N×N注意力图的需要。相反，它只操作并存储一个紧凑的d×d矩阵，其中d是自注意力模块中的潜在维度，且d≪N。具体来说，RED-PIM采用高效注意力[36 (https://arxiv.org/html/2607.21731#bib.bib29)]来显著减少内存使用，同时保持模型性能。我们重新组织注意力操作，使计算更局部化，减少数据移动，并创建可放入单个内存bank的更小中间结果。

除了重新组织注意力操作外，我们还引入了架构优化，包括近bank处理控制单元（PCU），它们使用优化的实现执行本地计算。我们提出架构增强，通过内存bank间的分层聚合策略来优化bank间数据流，从而降低数据传输和延迟。

在本文中，我们做出以下贡献：

- • 标准自注意力的瓶颈分析：我们表明，自注意力中的标准查询-键乘法引入了两个关键的低效性：(1) 它导致输入序列中所有token的全对全比较，从而造成显著的数据移动；(2) 它产生一个大的中间矩阵（注意力图），通常超过单个HBM bank的存储容量。因此，随后的注意力-值乘法需要额外的bank间通信来访问完整的分数矩阵。
- • 减少数据移动的算法优化：通过重新构造缩放点积注意力（SDPA），我们优化了数据局部性，最小化了冗余内存访问，并将HBM bank之间的数据移动从O\(N^2\)减少到O\(N\)。
- • HBM内处理的架构设计：我们将我们的架构与FIMDRAM[21 (https://arxiv.org/html/2607.21731#bib.bib30)]集成，利用PIM能力在内存bank内直接执行自注意力计算，并采用优化的实现。专用网络拓扑促进了高效的bank间通信。我们引入了一种分层聚合策略，逐步细化注意力计算并减少内存bank之间的冗余数据传输。我们设计了一个针对硬件优化的softmax函数，使用位操作和查找表来加速指数计算。这种方法在保持精度的同时减少了Transformer模型中的计算复杂度和内存使用。
- • 用于数据移动分析的定制仿真框架：基于DAMOV模拟器[30 (https://arxiv.org/html/2607.21731#bib.bib49)]，我们开发了一个定制的仿真环境来评估和建模HBM架构中的bank间数据移动，从而能够详细分析注意力机制的通信模式。
- • 广泛评估展示显著性能提升：我们的评估显示，与基线PIM实现相比，推理时间减少从16.05%到99.99%（几何平均值为66.42%），在较长序列上观察到最大改进。在真实数据集上，我们对长文档（IMDB, PubMed, Arxiv, WikiHop, 和 GovReport[6 (https://arxiv.org/html/2607.21731#bib.bib45),28 (https://arxiv.org/html/2607.21731#bib.bib54),41 (https://arxiv.org/html/2607.21731#bib.bib53),15 (https://arxiv.org/html/2607.21731#bib.bib52)]）实现了99.56%的几何平均改进，对短文档（GLUE基准[43 (https://arxiv.org/html/2607.21731#bib.bib38)]）为13.44%，同时保持或提高了准确性。

在本文的其余部分，我们在第II节 (https://arxiv.org/html/2607.21731#S2)中讨论Transformer模型、处理中内存和高带宽内存的一些背景知识。在第III节 (https://arxiv.org/html/2607.21731#S3)中介绍我们提出的算法变更和架构。在第IV节 (https://arxiv.org/html/2607.21731#S4)中讨论我们的方法和结果。第V节 (https://arxiv.org/html/2607.21731#S5)讨论相关工作，我们在第VI节 (https://arxiv.org/html/2607.21731#S6)中总结。

## II 背景

### II-A Transformer

Transformer由Vaswani等人在2017年提出[42 (https://arxiv.org/html/2607.21731#bib.bib1)]，是一类在各种应用中广泛使用的深度学习模型。它们的流行是因为能够并行处理序列数据，并有效捕获长距离依赖。基于Transformer的架构通常使用编码器块、解码器块或两者的堆叠构建。每个块由几个关键组件组成：多头注意力（MHA）、前馈神经网络（FFN）和层归一化，如图1 (https://arxiv.org/html/2607.21731#S2.F1)所示[42 (https://arxiv.org/html/2607.21731#bib.bib1)]。

MHA块旨在捕获序列中token之间的复杂关系。每个注意力头计算注意力分数，量化一个token对另一个token的影响。然后将多个注意力头的输出组合并处理，以生成输入序列的综合上下文表示。为了生成这样的表示，MHA块首先通过将线性变换应用于输入序列X = [x₁, x₂, ..., x_N]来推导查询(Q)、键(K)和值(V)矩阵。这些变换使用可训练的权重矩阵：查询矩阵W_Q、键矩阵W_K和值矩阵W_V：

Q = X·W_Q, V = X·W_V, K = X·W_K(1)
然后，缩放点积注意力机制按如下方式计算注意力分数：

Attention(K,Q,V) = softmax(Q·K^T / √d)·V (2)
其中d是键向量的维度。缩放因子√d防止点积变得过大，否则可能将softmax函数推入梯度非常小的区域，从而阻碍有效学习。多个注意力头的输出被拼接并传递到另一个线性层，以组合它们的结果并生成最终表示。

FFN层接收MHA块的输出，并将其转换为更具表现力的表示。这种转换通常涉及应用两个线性层，中间带有激活函数（例如ReLU），如下所示：

FFN(x) = max(0, x·W₁ + b₁)·W₂ + b₂ (3)
层归一化和残差连接对于增强Transformer的稳定性和提高训练效率至关重要。残差连接绕过注意力和前馈层，而层归一化确保一致的输入缩放，从而改善梯度流。

参考图注图1: Transformer模型架构[42 (https://arxiv.org/html/2607.21731#bib.bib1)]。近年来，Transformer已经显著增长。当前的Transformer模型包含数十亿个参数[44 (https://arxiv.org/html/2607.21731#bib.bib22)]。除了模型规模之外，序列长度也在增加，以适应更复杂的任务，进一步放大了这些模型的计算和内存需求[8 (https://arxiv.org/html/2607.21731#bib.bib23),11 (https://arxiv.org/html/2607.21731#bib.bib24)]。尽管规模和序列长度的增加提高了各个领域的性能，但也导致了显著更长的执行时间和更高的计算成本。

### II-B 处理/近内存架构

PIM的概念是为了应对传统冯·诺依曼架构中内存与处理单元之间数据移动日益增长的计算瓶颈而出现的[12 (https://arxiv.org/html/2607.21731#bib.bib25)]。在这些架构中，内存和计算单元的分离导致了显著的能量和时间成本，特别是在数据密集型应用中，如机器学习、图处理和高性能计算[1 (https://arxiv.org/html/2607.21731#bib.bib26),3 (https://arxiv.org/html/2607.21731#bib.bib27),12 (https://arxiv.org/html/2607.21731#bib.bib25)]。

PIM通过将计算能力直接集成到内存阵列中来解决这一挑战，使得数据可以在其存储位置附近进行处理。这种方法消除了处理器与内存之间频繁数据传输的需要，显著降低了执行时间和能耗[29 (https://arxiv.org/html/2607.21731#bib.bib28)]。

### II-C 先前基于PIM的Transformer加速器

先前关于基于PIM的Transformer加速器的研究大致可分为两类。

第一类包括将基于主机的加速器与近内存处理器相结合的混合架构，以提升Transformer性能[9 (https://arxiv.org/html/2607.21731#bib.bib12),14 (https://arxiv.org/html/2607.21731#bib.bib14),37 (https://arxiv.org/html/2607.21731#bib.bib13)]。虽然这些方法旨在结合传统系统和PIM系统的优势，但它们常常遭受主机与内存之间高通信开销的困扰。此外，它们引入了大量的架构复杂性，限制了可扩展性和能效。

第二类包括完全基于PIM的Transformer加速器，完全在内存内执行计算。这些包括模拟PIM设计[48 (https://arxiv.org/html/2607.21731#bib.bib10),40 (https://arxiv.org/html/2607.21731#bib.bib15),23 (https://arxiv.org/html/2607.21731#bib.bib16),27 (https://arxiv.org/html/2607.21731#bib.bib17)]和基于DRAM的解决方案[49 (https://arxiv.org/html/2607.21731#bib.bib11),37 (https://arxiv.org/html/2607.21731#bib.bib13)]。尽管有前景，这些设计仍然面临关键挑战：它们难以应对自注意力生成的巨大中间数据大小，并且bank间通信仍然是一个主要瓶颈，特别是对于长序列。此外，许多模拟PIM方法在精度、可靠性和非理想模拟行为方面面临限制，这阻碍了它们在实际部署中的实用性和准确性。

我们将在第V节 (https://arxiv.org/html/2607.21731#S5)中进一步讨论这些提议及其局限性。

### II-D 高带宽内存

HBM，如图2 (https://arxiv.org/html/2607.21731