AttAcc：基于PIM的Transformer推理注意力层专用加速器设计与实践-平芜编程栈

1. 项目概述：当Transformer推理遇上内存墙

如果你最近部署过像GPT-3/4、LLaMA这样的大语言模型进行在线推理服务，大概率会为一个问题头疼：随着用户请求的增多，你希望通过增大批处理规模来提升GPU利用率和服务吞吐量，但很快就会发现，响应延迟开始不受控制地飙升，甚至触发服务等级目标的违约。问题的核心，往往不在大家通常关注的矩阵乘计算上，而在于那个让Transformer如此强大的“注意力机制”。在批量推理时，每个请求都会生成自己独有的、庞大的键值矩阵，这些矩阵需要被反复读取，但计算强度却很低，导致整个系统卡在了内存带宽上，GPU强大的算力只能干等着数据搬运。

这正是我们团队在研究和实践中遇到的核心瓶颈。传统的以GPU为中心的计算架构，在处理Transformer生成模型的大批量推理时，出现了明显的“水土不服”。模型参数可以共享，但注意力机制中的上下文记忆（KV Cache）却是个“私人订制”的庞然大物，其总规模随着批次大小和序列长度线性增长，轻松就能超过模型权重本身，挤爆显存。更棘手的是，注意力层的计算属于典型的内存带宽受限操作，即使批量再大，其“计算操作数/访存字节数”的比例也极低，GPU的SM单元大量时间处于空闲状态。

基于此，我们设计并验证了AttAcc——一个专为Transformer生成模型大批量推理中的注意力层量身定制的处理内存加速器。它的核心思想非常直接：既然数据搬运是瓶颈，而KV矩阵具有“一次写入，多次读取”的鲜明特征，那我们就把计算搬到数据旁边去。通过将专用的处理单元嵌入到存储设备内部，AttAcc能够以极高的内部带宽处理注意力计算，同时仅需极低的外部带宽与主机GPU交换必要的输入输出数据。实测表明，在GPT-3 175B模型上，AttAcc与GPU组成的异构系统，能在满足严格延迟目标的前提下，实现相比顶级GPU系统最高3.24倍的吞吐量提升和3.22倍的能效改进。这篇文章，我将为你深入拆解AttAcc的设计思路、架构细节以及我们趟过的一些坑，希望能为面临类似推理性能瓶颈的工程师和研究者提供一个新的解题视角。

2. Transformer推理瓶颈的深度剖析

要理解AttAcc为什么这么设计，首先得把Transformer，特别是GPT这类自回归生成模型在推理时的行为特征和瓶颈点掰开揉碎看清楚。很多人对Transformer的训练很熟悉，但推理，尤其是大批量在线推理，有着截然不同的计算特征和系统挑战。

2.1 生成式推理的两阶段模式与KV Cache

以GPT为例，处理一个用户请求（如“写一首关于春天的诗”）的推理过程，清晰地分为两个阶段：

摘要阶段：模型读取并理解整个输入提示（Prompt）。这个阶段是标准的Transformer前向传播，主要进行的是计算密集型的通用矩阵乘法。更重要的是，在这个过程中，模型会为输入序列中的每个token生成并保存一对键向量和值向量，共同构成该请求的KV矩阵。你可以把它理解为模型为这段对话创建的“短期工作记忆”。
生成阶段：模型开始自回归地逐个生成输出token。每生成一个新token，都需要执行一次完整的解码器前向传播。但这里有个关键：在注意力层，新token的查询向量需要与历史上所有已生成token（包括输入提示）的KV矩阵进行交互，以决定下一个token应该是什么。这意味着，在生成第一个输出token时，KV矩阵的大小是Lin x demb；生成第二个时，大小变为(Lin+1) x demb，以此类推。

这个不断增长的KV矩阵，就是所谓的KV Cache。它是每个请求独有的上下文状态，无法在批次内共享。当你有成千上万个并发请求，且每个请求的对话历史都很长时，所有请求的KV Cache总量会成为一个天文数字。

2.2 批处理的收益与代价

为了提高硬件利用率，我们自然想到批处理——将多个请求打包一起处理。这对于模型中的全连接层（即线性变换层）是福音。批处理将多个独立的矩阵-向量乘合并为一个更大的矩阵-矩阵乘，显著提升了计算密度，让GPU的Tensor Core能吃饱，同时权重数据只需从显存读取一次，即可服务整个批次，访存效率大增。

然而，批处理对于注意力层却是一把双刃剑。

收益有限：注意力层的核心计算是查询向量与KV矩阵的矩阵-向量乘。即使批量增大，这个操作的本质依然是内存带宽受限的。因为对于每个请求，你仍然需要从庞大的KV矩阵中读取对应的键和值向量，与一个很小的查询向量做计算。计算量增长远跟不上访存量增长。
代价显著：
1. 内存容量压力剧增：KV Cache的总大小 =批量大小 x 序列长度 x 模型隐藏维度 x 2（K和V）x 精度（如2字节 for FP16）。以GPT-3 175B（demb=12288）为例，Lin=Lout=2048时，单个请求的KV Cache就高达约18GB。64个这样的请求，仅KV Cache就需要约1.15TB！这远远超过了当前顶级GPU（如A100 80GB）的显存容量。
2. 成为延迟主导：随着批量增大，全连接层的处理时间几乎不变（因为算力饱和），但注意力层的处理时间却线性增长，因为它严重受限于内存带宽。我们的仿真和实测都表明，在大批量下，注意力层的耗时可以占到整个生成阶段耗时的60%以上。
3. 违反服务等级目标：在线服务通常有严格的延迟要求，例如每个token的生成需要在50毫秒内完成。当注意力层延迟因批量增大而飙升时，为了满足SLO，你不得不大幅限制批量大小，这又反过来限制了吞吐量，陷入两难。

注意：这里存在一个常见的误解，认为用更快的GPU或更宽的内存总线就能解决问题。但问题在于，注意力层的计算强度（Ops/Byte）本身极低（约1），这意味着无论内存带宽多高，大部分时间依然花在等待数据上，计算单元利用率难以提升。单纯堆硬件带宽，成本效益极差。

3. AttAcc架构设计：将计算植入数据腹地

面对上述瓶颈，我们的设计哲学是“解耦”与“异构”。让擅长高计算密度任务的GPU继续处理全连接层，而将那个“难啃”的、内存密集型的注意力层，交给一个专门为此优化的加速器——这就是AttAcc。它的核心是利用了处理内存技术。

3.1 为什么PIM是注意力层的“天作之合”？

PIM不是新概念，但在Transformer推理的特定场景下，其优势被放大到了极致。关键在于注意力层数据的两个独特访问模式：

极高的数据复用率：KV矩阵在摘要阶段写入一次，但在后续数十、数百个生成阶段中被反复读取。数据一旦生成，就长期驻留在计算单元附近被频繁使用。
极低的内外数据交换比：我们仔细分析一下注意力层中一次矩阵-向量乘的数据流。以GPT-3为例，对于一次注意力头的计算，输入是一个1x12288的查询向量，需要与一个Lx12288的K（或V）矩阵相乘，产生一个1xL的中间结果（或最终输出）。这里，需要从外部传入AttAcc的数据量（输入向量+输出结果）与AttAcc内部需要读取的KV矩阵数据量之比，大约是(demb + N_head*L) / (L*demb)。当序列长度L很大时（如2048），这个比值趋近于N_head/demb = 96/12288 ≈ 1/128。

这个1/128的比值是AttAcc设计的基石。它意味着，如果我们能把KV矩阵存储在AttAcc内部，并将矩阵-向量乘的计算单元也放在里面，那么对于每次计算，我们只需要从外部（GPU）传入一个很小的查询向量，并在计算完成后传回一个同样很小的结果。而那个庞然大物般的KV矩阵，完全在AttAcc内部的高带宽总线上移动，避免了穿越功耗高昂、带宽受限的片外接口（如PCIe或NVLink）。外部带宽需求因此降低了两个数量级！

3.2 AttAcc硬件架构详解

我们的AttAcc设计以三星的HBM-PIM为蓝本，并针对注意力计算进行了定制化。下图勾勒了其核心架构思想：（想象一个由多个HBM堆栈组成的系统，每个堆栈对应一个或多个注意力头）

3.2.1 存储与计算单元的紧耦合传统的HBM由多个DRAM核心堆叠在一个缓冲芯片上。在HBM-PIM中，我们在每个DRAM bank的I/O边界集成了轻量级的处理单元。AttAcc沿用了这一思想，并将多个这样的“智能存储单元”并行工作。

KV矩阵的存储与划分：我们将整个模型的KV矩阵（按注意力头划分）分布存储在不同的AttAcc设备中。具体来说，每个注意力头对应的K和V矩阵，被列向（对K）和行向（对V）切片，分布到单个HBM堆栈内的多个bank中。这样，当处理一个查询向量时，所有bank可以并行地计算自己那部分K或V切片与查询向量的点积。
计算任务的映射：这种数据划分方式完美匹配了注意力计算的并行性。并行性存在于三个维度：不同用户请求、同一请求内的不同注意力头、同一注意力头内矩阵乘的向量化计算。AttAcc的bank级并行可以充分利用这三个维度的并行性。

3.2.2 计算流水线与操作卸载注意力层的计算包含几个步骤：QK点积（Score）、Softmax、与V点积（Context）。并非所有操作都同样适合在DRAM bank旁执行。

Bank-side PEs（处理单元）：我们将矩阵-向量乘的核心计算单元直接放在DRAM bank旁边。这些PE设计相对简单，专注于高效的乘加运算，能直接以DRAM内部的高带宽（通常是外部接口带宽的数十倍）访问KV矩阵数据。这是性能提升的关键。
Buffer-die PEs：Softmax操作中的指数计算等非线性函数，需要更复杂的算术逻辑单元，放在DRAM工艺的bank旁可能面积和能效不佳。因此，我们将其放在HBM的缓冲芯片上。缓冲芯片通常采用更先进的逻辑工艺，适合部署这类计算单元。Bank-side PE计算出的点积结果被送到Buffer-die进行Softmax规约，结果再分发回各bank用于后续的context计算。

3.2.3 与主机的协同工作流

摘要阶段：GPU正常执行，计算生成每个请求的KV矩阵。生成完毕后，GPU将这些KV矩阵通过高速互联（如CXL）写入到AttAcc的对应存储区域。这是主要的写入开销，但每个请求仅此一次。
生成阶段：对于每个生成步，GPU计算得到当前token的查询向量，将其广播给所有相关的AttAcc设备。每个AttAcc设备内部并行完成自己负责的那部分注意力计算（包括bank-side GEMV和buffer-die softmax），然后将部分结果归约后，将最终的注意力输出向量返回给GPU。GPU接着进行后续的全连接层等计算。

实操心得：数据布局是关键。如何将KV矩阵切片并映射到多个AttAcc设备的bank上，直接影响负载均衡和通信开销。我们的策略是确保每个bank的计算负载大致相等，并尽量让单个请求的多次生成步骤中，与GPU的数据交换是连续和聚合的，以减少通信延迟。在实践中，我们采用了类似[TRiM]论文中提出的张量缩减映射策略的变种，取得了不错的效果。

4. 性能评估与对比分析

我们构建了一个详细的周期级模拟器来评估AttAcc的性能。基线系统是配备640GB HBM的NVIDIA DGX A100系统。对比方案有两个：1) 基线系统；2) 假设基线系统显存翻倍至1280GB（DGX-1280GB）；3) GPU + AttAcc的异构系统，其中GPU持有326GB模型权重，AttAcc提供额外的954GB容量用于KV Cache，总内存亦为1280GB。AttAcc的聚合内部带宽设定为64 TB/s（是DGX外部带宽的4倍），其DRAM访问能效基于文献数据设定为比普通HBM高3.5倍。

4.1 吞吐量提升：打破SLO约束下的批量限制

我们测试了不同输入/输出序列长度组合下，系统能达到的吞吐量（Tokens/sec）。在没有服务等级目标约束的理想情况下，随着批量增大，纯GPU系统的吞吐量增长会先后遇到两个天花板：内存容量天花板和注意力延迟天花板。

容量天花板：如图3(a)所示，对于长序列（Lin=2048），DGX-640GB在批量约35时就耗尽了显存，无法继续提升吞吐。DGX-1280GB虽然容量更大，能支持更大批量。
延迟天花板：但即使容量无限，由于注意力层延迟随批量线性增长，为了满足严格的SLO（如每个token 30ms），DGX-1280GB的批量也被限制在了一个较低的值，从而限制了其吞吐量。

AttAcc的破局点在于：它将注意力计算移到了高带宽的PIM内部，极大降低了该阶段的处理延迟。因此，在相同的SLO约束下，AttAcc系统能够支持比纯GPU系统大得多的批量。如图4所示，在50ms SLO、Lin=2048的场景下，DGX+AttAcc的吞吐量达到了DGX-1280GB的2倍以上；在更严格的30ms SLO下，优势扩大到近6倍。这意味着，在保证用户体验（低延迟）的前提下，AttAcc能让单台服务器承载更多的并发用户。

4.2 能效优势：省下的是“搬运数据”的电

能效提升同样显著，主要来自两个方面：

权重数据访存减少：由于AttAcc允许系统运行在更大的批量下，模型权重被GPU重复利用的次数更多，平均每个token的权重访存能耗下降。
KV数据访存能耗大幅降低：这是最大的贡献来源。在GPU方案中，庞大的KV矩阵需要在HBM和GPU SM之间来回搬运，功耗很高。在AttAcc方案中，KV矩阵安静地待在AttAcc内部，计算时数据仅在AttAcc内部的高能效总线上移动。外部总线只传输微小的查询向量和结果。如图5所示，在不考虑SLO的情况下，AttAcc系统相比DGX-640GB，每生成一个token的片外内存访问能耗降低了最高69%。

表格：不同配置下性能与能效对比摘要（以Lin=2048, Lout=2048为例）

系统配置	最大支持批量 (无SLO)	吞吐量 (Tokens/s) @ 50ms SLO	相对基线吞吐量	每Token内存访问能耗	相对基节能耗比
DGX (640GB)	~35	受限 (批量小)	1.0x (基线)	1.0x (基线)	1.0x
DGX (1280GB)	>100	中等 (受限于注意力延迟)	~1.5x	0.85x	~1.18x
DGX + AttAcc	>100	高(延迟瓶颈缓解)	3.24x	0.31x	3.22x

注意：这里的能效比是“性能/功耗”的提升。AttAcc通过降低执行时间（提升性能）和降低单位计算的能耗（主要是访存能耗），实现了能效的平方级改善。对于大规模部署的数据中心，这直接意味着电费和冷却成本的显著下降。

5. 实践考量与未来展望

AttAcc的设计理念虽然清晰，但要走向实际部署，还需要在工程和生态层面解决一系列问题。

5.1 软件栈与编程模型

让开发者像用CUDA一样方便地使用AttAcc是关键。我们设想了一个分层的软件栈：

编译器/运行时层：需要扩展现有的深度学习编译器（如TVM、Torch-MLIR），使其能够识别计算图中的注意力层算子，并自动将其切分、标注。一部分计算（FC层）生成GPU代码，另一部分（注意力层）生成AttAcc的指令或调用其驱动API。
驱动与内核层：需要为AttAcc设备开发驱动程序，管理设备内存、任务队列、与GPU的同步等。内核层则提供高度优化的注意力计算内核，充分利用AttAcc的硬件特性。
高层框架集成：最终目标是在PyTorch或TensorFlow中，提供一个类似torch.nn.Attention的模块，但其后端实现可以透明地分派到AttAcc上执行。用户无需修改模型代码，只需指定部署设备即可。

5.2 系统集成与挑战

互联瓶颈：即使AttAcc大幅减少了数据交换量，GPU与多个AttAcc设备之间的互联带宽和延迟仍然重要。PCIe Gen5可能成为初期瓶颈，未来需要CXL 2.0/3.0或更高速的定制互联来充分发挥性能。
负载均衡与调度：在多个AttAcc设备间动态分配不同请求的KV Cache是一个复杂的调度问题。需要考虑请求的序列长度、生命周期，以及设备间的负载均衡，避免出现“热点”设备。
容错与可靠性：将关键状态（KV Cache）放在一个可能独立于GPU的加速器中，需要设计新的容错机制。例如，如何快速恢复一个故障AttAcc设备上存储的上下文状态？

5.3 设计空间的延伸

我们目前基于HBM-PIM的设计只是PIM的一种形态。更大的设计空间有待探索：

近内存计算：将计算单元放在内存控制器附近或3D堆叠存储的底层逻辑芯片上，可能是性能与灵活性之间更好的折衷。
存内计算：利用新型非易失存储器（如ReRAM）的模拟计算特性，直接在存储单元内完成矩阵-向量乘的乘加运算，有望实现极致的能效，但面临精度、器件变异和制造挑战。
异构内存层级：将热门的、活跃的KV Cache放在AttAcc（或更快的存内计算单元）中，将较冷的或归档的Cache移至更便宜的大容量内存中，构建分层的内存系统。

5.4 不适用场景澄清

AttAcc并非万能。它的优势场景非常特定：Transformer类生成模型的大批量、长上下文推理。对于以下场景，其优势可能不明显甚至不适用：

模型训练：训练过程以计算密集的前向/后向传播为主，KV矩阵每次迭代都会变化，写入开销巨大，无法发挥PIM“一次写入、多次读取”的优势。
小批量或短序列推理：当批量很小或序列很短时，注意力层本身的计算和访存量不大，GPU足以应对，引入额外硬件带来的收益可能无法覆盖其成本和管理复杂度。
非自回归模型：一些图像生成模型（如Diffusion Transformer）虽然也用注意力，但其推理模式不同，KV矩阵的复用模式需要重新评估。

我个人在实际探索中的体会是，硬件架构的创新必须与上层应用的工作负载特征深度结合。AttAcc的成功，源于它精准地捕捉并利用了Transformer推理中一个非常具体但影响巨大的数据访问模式。它给我们带来的启示是，在“后摩尔时代”，通过软硬件协同设计，针对特定领域的关键瓶颈进行“外科手术式”的加速，可能是持续提升系统效能的最有效路径。对于广大AI应用开发者而言，理解底层计算和访存特征，不再将GPU视为黑盒，是进行高效部署和优化的第一步。未来，随着类似AttAcc的专用加速器出现，AI推理系统的架构可能会变得更加异构和多样化，而驾驭好这样的系统，将是工程师们新的核心竞争力。