NeuPIMs-提升LLM推理的NPU-PIM异构加速

大语言模型(Large Language Models, llm)被广泛应用于各个领域，如自然语言理解、内容生成和决策支持。然而，这些模型的一个关键挑战在于内存和计算。本文专门解决了当代llm中的推理挑战。

LLM推理涉及大量大规模GEMV和GEMM的计算。为了解决这种计算需求，一种常见的做法是利用高性能机器学习(ML)加速器，如gpu和tpu。npu通常针对计算密集型任务进行优化，特别是针对GEMV的有效执行。然而，它们对于gemv的效用并不理想，因为后者的计算强度较低，这导致NPU的计算资源利用率不足。另一方面，内存中处理(Processing- in-Memory, PIM)技术虽然对GEMV不那么有效，但对带宽密集型GEMV操作显示出了希望。

为此，本工作提出了NeuPIMs，一种用于llm批量推理的新型异构加速系统。我们构建NeuPIMs，使其有效地平衡内存带宽和系统计算资源的利用，以提高整体推理吞吐量。NeuPIMs是一种新型的异构加速系统，专为批量推理大型语言模型（LLMs）而设计。该系统结合了传统的以GEMM（通用矩阵乘法）为中心的神经处理单元（NPU）和针对GEMV（通用矩阵-向量乘法）优化的内存中处理（PIM）设备。

图1.拟议的NeuPIMs加速器

LLM推理的计算特征

LLM模型架构和执行。图2展示了所有最先进的大型语言模型共享的模型体系结构。这张图在整个论文中都是一个反复出现的例子。对于输入提示(例如，“我认为这个”)，模型经历了一个总结阶段，对输入进行编码，为随后的生成阶段建立上下文。在生成阶段，模型以自回归的方式在每次迭代中产生一个令牌，使用生成的键值投影用于下一次迭代。这两个阶段都构成一个解码块序列，每个解码块由三个主要层组成: (1)QKV生成，(2)多头注意(MHA)，以及(3)一组前馈网络(ffn)。

图2.LLM模型体系结构和推理

为了更好地理解LLM推理的计算特征，我们使用两个GPT3变体GPT3- 13b和GPT3- 175b进行了屋顶线分析。图3显示了算术强度(FLOPS/byte)与性能(TFLOPS)之间的关系。我们观察到，对于这两种模型，生成阶段是严重的内存绑定，而汇总阶段是计算绑定。由于这两个阶段具有算法依赖性，并且以顺序的方式交替发生，因此使用同构计算平台可实现高资源利用率。并促使我们设计了一个异构系统，该系统结合了以计算为中心的基于收缩阵列的GEMM NPU和以内存为中心的gemv内存中处理(PIM)加速器。

图3. LLM层的算力强度

基于GPU的LLM推理服务

图4.四个不同llm的GPU资源利用率

GPU系统利用率不足。我们分析了一个配备gpu的基线系统，以了解LLM推理的计算、内存和带宽的利用率。我们将系统与NVIDIA GeForce RTX 3090 24GB和NVIDIA A100 40GB进行比较，运行四种不同的LLM模型:GPT- NeoX, LLaMa2, OPT和MPT。图4给出了利用率结果以及以误差条表示的分层变化。该图表明，尽管并行化方案存在固有的缺陷，但容量利用率接近100%。这个观察结果是直观的，因为使用的gpu数量是根据容量限制确定的。然而，计算资源的利用率一直低于40%，这表明基于gpu的LLM推理系统的成本效率低下。这种利用率不足的原因是带宽不足，尽管a100配备了HBM，提供了1555 GB/s的聚合。不幸的是，只要gemm和gemv之间的串行依赖关系持续存在，这种不平衡是不可避免的。

Naive NPU-PIM方法

解决此带宽瓶颈的一种直接方法是利用内存中处理(Processing-in-Memory, PIM)技术，该技术允许将带宽受限的计算卸载到内存中的加速器上。因此，我们设计了一个naively集成NPU-PIM加速器，利用NPU的收缩阵列架构，并结合了最先进的基于pim的gemv加速器Newton。

NPU和PIM并行执行的必要性。观察到利用率不足主要是由于PIM微体系结构的基本限制，该限制不允许主机(NPU)和PIM单元的并发执行，从而序列化不相关的资源使用。因此，在NPU加速器中实现PIM实际应用的最关键挑战是使它们能够并行执行。这一研究问题构成了本工作的重点。

 NeuPIMs调度

双行缓冲器的集成使NeuPIMs能够同时处理NPU的内存访问和PIM命令。本节描述了多头注意(MHA)层执行的计算重叠机会，以及一种用于批处理LLM推理的新型请求调度技术，称为子批交错。

1. MHA层的重叠机会

1. 

   图5.多头注意层的重叠机会

图5显示了NeuPIMs中(1)pim端logit和attend操作的重叠，(2)NPU端softmax操作的重叠。由于多头注意力的操作可以分解为头部粒度，因此即使是naïve NPU-PIM集成架构也应该有机会利用两种资源都要重叠运营。然而，它不能利用这个机会，因为操作结果不能通过PIM通道在PIM单元和矢量单元之间传递。相比之下，NeuPIMs通过采用双行缓冲器，可以同时利用NPU和PIM。这种配置允许NeuPIMs中的矢量单元存储部分logit (softmax)值，而不必等待PIM上完成整个logit操作(GEMV)，从而减少了集成单元的利用率不足。

值得注意的是，这种重叠之所以可能，只是因为存在头级并行性，而这种并行性仅在多头注意层中可用。此外，重叠的机会只存在于PIM和NPU的矢量单元之间，导致NPU的收缩阵列在执行MHA层期间大部分未使用。

2.子批交错

序列化执行的限制。图6(a)显示了解码器块操作的执行时间轴示例，运行在naïve NPU-PIM集成设备上。特别是，该图描述了批次内QKV生成、多头注意和投影& ffn之间的算法依赖关系。由于依赖关系，它们必须串行执行，这不可避免地导致NPU和PIM的利用率都很低。

为了应对这一挑战，我们提出了子批交错，将一个大批划分为两个子批，并交替执行以提高资源利用率。图6(b)描述了子批交错技术，该技术允许在PIM和NPU上的子批内分别同时执行PIM友好和NPU友好的操作，从而显著提高NPU和PIM的利用率。

图6.解码器块执行时间表示例:(a)序列化执行，(b)子批交错。

执行时间线的对比分析。让表示要在单个NeuPIMs设备上执行的解码器块的数量。图6显示，在不使用子批交错的情况下，每个解码器块中的运算符将按顺序执行，导致总执行时间等于times每个解码器块执行时间(即。， NPU-S上的QKV生成+ PIM上的MHA; NPU-S上的NPU-V +投影和ffn)。

然而，如图6所示，子批交错在NPU-S执行时间中有效地隐藏了MHA层的执行时间，导致总执行时间等于(二进制操作- 1)乘以每个子批部分执行时间(即。(在NPU-S上生成投影、ffn和QKV)加上在开始和结束时划分的单个解码器块执行时间。当交错发生时，NPU和PIM的利用率得到了提高，因为它们的执行有效地重叠了。我们的实证研究表明，交错期的NeuPIMs执行时间主要受到运行GEMM操作的NPU执行时间的限制，从而隐藏了MHA层执行的PIM执行时间。

结论

在这项工作中，我们提出了一个新的系统，NeuPIMs，它将NPU(通用ML加速器)与PIM技术集成在一起，以减轻与变压器层中不同操作及其数据流相关的限制。我们为所提出的系统引入了一种新的调度和执行策略，该策略可以更好地利用HBM内存、计算密集型NPU和PIM加速器进行LLM推理服务。结果表明，与naïvely集成NPU和PIM加速器的基线系统相比，本工作中开发的系统提供了1.6倍的吞吐量改进。