PowerInfer-2：智能手机上的大语言模型快速推理

24年6月来自上海交大的论文“PowerInfer-2: Fast Large Language Model Inference on a Smartphone”。

PowerInfer-2是一个专为在智能手机上高速推理大语言模型 (LLM) 而设计的框架，对于规模超出设备内存容量的模型特别有效。PowerInfer-2 的关键见解是将传统矩阵计算分解为细粒度神经元簇（cluster）计算来利用智能手机中的异构计算、内存和 I/O 资源。具体而言，PowerInfer-2 具有多形态神经元引擎，可调整 LLM 推理各个阶段的计算策略。此外，它还引入分段神经元缓存和细粒度神经元簇（cluster）-级流水线，可有效最小化和隐藏 I/O 操作造成的开销。PowerInfer-2 能够在两部智能手机上支持各种 LLM 模型，与最先进的框架相比，速度提高了 29.2 倍。值得注意的是，PowerInfer-2 是首个在智能手机上以每秒 11.68 个tokens的生成速率为 TurboSparse-Mixtral-47B 模型提供服务的系统。对于完全适合内存的模型，PowerInfer-2 可将内存使用量减少约 40%，同时保持与 llama.cpp 和 MLC-LLM 相当的推理速度。

研究人员探索了两种在资源受限条件下提供 LLM 推理的有前途方法。

鉴于智能手机的内存容量有限，一种策略是部署缩小版的 LLM。例如，谷歌的 Gemini Nano 3.25B [14] 使用不到 2GB 的内存，它通过减少智能功能以适应内存限制，代表了一种妥协。这是因为更大的模型具有增强的智能，这种现象被称为“规模化定律” [18]。

或者，一些技术旨在降低推理过程中 LLM 权重的计算和存储需求。PowerInfer [31] 将热激活神经元分配给 GPU 并将冷神经元分配给 CPU，将个人计算机 (PC) 上的推理速度提高了 11 倍。另一种方法，LLM in a Flash [4]，通过使用基于闪存的 NVMe 存储来存大型模型权重，从而缓解内存限制。然而，这些解决方案在智能手机上并不适用，因为智能手机的性能较弱，硬件和存储设备异构性，带宽较低，并且由于命令队列单一，不支持并发访问。这使得 I/O 做的事情成为移动设备上 LLM 推理的常见瓶颈。

首先分析通用闪存（UFS） 4.0 的随机和连续读取吞吐量。一个显著的特点是，UFS 的读取带宽随读取块大小而变化。通常，无论是连续读取还是随机读取，块越大，读取带宽越大。例如，当块大小设置为 512KB 时，连续读取和随机读取带宽分别达到最大值 4 GB/s 和 3.5 GB/s。当块大小减小到 4KB 时，带宽最小，随机读取带宽为 450 MB/s。

UFS 随机读取表现出一个有趣的现象，即随机读取性能受随机读取范围影响，随机读取范围越小，带宽越高。如图是XPU的计算性能和I/O吞吐量比较。如图（a）所示是XPU的执行时间对比，计算不同批次的矩阵-向量乘积。如图（b）所示，在 UFS 4.0 中，如果将 4KB 随机读范围设置为 128MB、256MB、512MB，128MB 范围的带宽最高，达到 1GB/s，而 512MB 范围的带宽最低，低于 850MB/s。值得注意的是，其他块大小下这种现象并不明显。因此，128MB 范围内的 4KB 随机读取带宽超过了 8KB 和 12KB 块大小。

读取带宽受发出读取命令的 CPU 影响。CPU 核频率越高，读取带宽就越大。如表所示，当用频率为 3.3GHz 的大核进行随机读取时，4KB 读取带宽达到 1 GB/s。相反，当用频率为 2.2GHz 的小核进行相同的随机读取时，带宽只有大约 760 MB/s。出现这种相关性的原因是，发起读取的 CPU 核需要运行 UFS 驱动程序，因此更高的频率可以更快地处理与 UFS 相关的 I/O 操作，包括中断和队列管理。

与 NVMe 不同，移动设备中的 UFS 存储只有一个命令队列，本质上缺乏内部并发能力。因此，与使用单核相比，使用多核发起 I/O 命令不会带来更高的 I/O 带宽。如上表所示，由于 UFS 命令队列中的竞争，使用多核进行 4KB 随机读取甚至会使 I/O 性能下降高达 40%。

总之，当需要将某些模型权重存储在移动设备的存储介质上时，高效的 LLM 系统必须充分考虑存储介质的性能特性，以最大限度地提高 I/O 带宽并最大限度地减少与 I/O 操作相关的性能开销。

LLM 推理包括两个阶段：预填和解码阶段。在预填阶段，LLM 在一次迭代中处理用户的提示，从而生成第一个token。另一方面，解码阶段涉及 LLM 以自回归的方式一次一个地顺序生成token。在预填阶段生成的token用作生成第二个token的输入。然后，第二个token用作 LLM 的输入，从而促进第三个单词的生成。此序列持续，直到输出序列完成或到达序列结束 (EOS) token。

这两个阶段表现出不同的计算模式，需要优化两个关键指标：预填阶段的第一个token时间 (TTFT) 和解码阶段的token之间时间 (TBT)。预填阶段在一次迭代中处理所有提示tokens，带来相当大的计算负担；相反，解码阶段每次迭代仅处理一个 token，因此计算需求相对较低。因此，LLM 推理系统必须利用专门为这些阶段设计的计算策略来有效优化性能指标。

传统的LLM推理通常以矩阵计算为基本推理单位，在智能手机异构硬件环境中会带来较大的计算和I/O开销，这种粗粒度的计算无法充分发挥XPU灵活的计算能力，更糟糕的是，如果矩阵权重的一段存储在存储设备上，则在开始矩阵计算之前必须延迟将这些权重加载到内存中，从而导致相当长的I/O等待时间。

专为智能手机设计的高速LLM推理框架PowerInfer-2，实现三个目标：1）低推理延迟：最小化预填充阶段（TTFT）和解码阶段（TBT）的推理延迟；2）低内存占用：减少推理过程中的内存使用量，即使模型大小超出设备内存限制，也能实现LLM的低延迟推理； 3）灵活性：确保设计可以无缝适应具有不同计算、内存和存储容量的智能手机。

提出一种称为神经元簇（neuron cluster）的计算抽象单元，它专门为异构计算场景中的 LLM 推理而设计。PowerInfer-2 以神经元簇的粒度执行计算和 I/O 操作，神经元簇可以在计算过程中由多个激活的神经元动态组成，神经元的数量由计算单元的计算能力决定。例如，在解码阶段，当由 CPU 核心执行计算时，分配给每个 CPU 核的神经元簇大小小于在预填充阶段 NPU 计算期间处理的神经元簇大小。用这种抽象，PowerInfer-2 可以充分利用具有不同计算能力的 XPU，有效地隐藏 I/O 开销。

如图是 PowerInfer-2 的整体架构，它分为在线（右侧部分）和离线（左侧部分）程序。在线部分以神经元簇粒度提供推理服务，包括四个协作组件：多形态神经元引擎、内存神经元缓存、灵活的神经元加载和神经元簇-级 I/O 流水线。

神经元引擎

多形态神经元引擎在预填充和解码阶段使用完全不同的计算模式。对于预填充阶段，神经元簇包含来自权重矩阵的所有神经元，主要依赖于 NPU，因为它可以高效地处理大型矩阵乘法。对于解码阶段，它会调用预测器来确定在启动计算之前将激活哪些神经元。然后，引擎将这些激活的神经元合并为一个小的神经元簇，并利用 CPU 核动态地计算神经元簇，从而大幅减少运行时的计算需求和内存使用量。

在开始推理计算之前，计算引擎会从神经元缓存中检索神经元权重，该缓存经过优化，可利用 LLM 推理中神经元-级的访问局部性。如果发生一个缓存丢失，PowerInfer-2 将启动 I/O 命令以从存储中获取未缓存的神经元权重。为了减少 I/O 延迟，PowerInfer-2 引入一种流水线机制，可同时处理神经元簇和 I/O 操作。此外，PowerInfer-2 通过自适应地捆绑和加载神经元来最大限度地减少 I/O 开销，这取决于模型的量化。

为了自动适应不同的模型或智能手机，在开始在线推理之前，对最初在新智能手机上使用的每个模型进行一次离线程序。此过程涉及接收三种类型的输入：模型权重、用户输入和硬件规格。它输出一个执行规划，描述在线推理中涉及的每个组件配置并指导在线程序。

具体来说，离线规划器会输出计算、内存和 I/O 的配置。对于计算，规划器会根据 CPU 和 NPU 的计算强度确定它们在不同阶段或层级的使用比。在内存配置方面，为了在内存使用率和推理性能之间取得平衡，规划器允许用户在运行 PowerInfer-2 之前设置所需的推理速度。基于此速度设置，PowerInfer-2 会计算所需的最佳缓存大小。对于 I/O 配置，规划器会触发分析器来测量模型的稀疏性以及热/冷神经元的分布。
如图是多形态神经引擎两个阶段，NPU为中心预填和CPU为中心解码的示意图。

图（a）演示 CPU 和 NPU 如何协作以 Transformer 层粒度执行预填充阶段推理。NPU 计算需要使用与 CPU 共享的有限内存。因此，在 NPU 计算开始之前，CPU 应将所需的矩阵权重预加载到此共享内存中。在特定的 LLM 层中，在 NPU 进行任何矩阵乘法之前，多个 CPU 核会从神经元缓存中读取量化矩阵权重，并提前将这些矩阵权重反量化为 fp16，最终将结果存储在 CPU 和 NPU 之间的共享内存中。同时，PowerInfer-2 使用另一个大核将下一层的所有矩阵权重异步预加载到神经元缓存中。中核的反量化、NPU 的计算和大核的 I/O 操作同时进行，以减少 I/O 开销。值得注意的是，由于预填充阶段涉及密集矩阵而非稀疏计算，通过 I/O 进行权重加载可以利用顺序读取将大块数据加载到内存中，从而最大限度地利用 UFS 的 I/O 带宽。

图 （b） 展示不同 CPU 核心进行的解码阶段推理。CPU 核首先从神经元缓存中读取注意模块的权重，并使用输入向量进行计算。然后运行预测器来确定后续权重矩阵中神经元的激活状态。第三步，CPU 核将激活的神经元分成几个簇，每个核负责使用输入向量计算其簇内的激活神经元，并最终在一个屏障处聚合结果。如果这些神经元在神经元缓存中，CPU 核将使用输入向量进行计算。在发生缓存丢失的情况下，神经元不在神经元缓存中，CPU 核上运行的 I/O 线程将神经元异步加载到缓存中，最终通知计算线程完成计算。

内存的神经缓存

高效的缓存设计可以避免昂贵的存储 I/O 活动，从而优化端到端推理性能。缓存系统的有效性取决于推理过程中局部性的存在。然而，传统的 LLM 推理需要遍历生成每个 token 的所有权重，没有局部性，导致任何缓存设计都无效。

LLM in a Flash [4] 提出利用稀疏激活在推理过程中有选择地加载权重。它还将共同激活的神经元捆绑在一起，并将它们从 Flash 一起加载，以减少 I/O 操作。然而，这种方法忽略了神经元激活的倾斜分布，其中一些热神经元激活更频繁，并且与大多数其他神经元高度相关。这给设计有效的缓存策略带来了挑战。首先，这些流行的神经元位于不同的神经元束中，并从 Flash 冗余加载，浪费了 I/O 带宽。其次，移除这些热神经元会将剩余神经元共同激活的可能性降低到 20% 以下，从而使捆绑机制在减少 I/O 操作方面无效。

为了解决这个问题，PowerInfer-2 引入一种分割神经元缓存设计，该设计针对 LLM 内的各种数据类型量身定制。它将缓存划分为多个区域，每个区域都有特定的预取和退出策略。注意块权重较小且激活较少，在整个运行时都会被预加载和保留。

相比之下，FFN 块容易频繁激活热神经元，它对这些神经元使用基于最近最少使用 (LRU) 的动态退出策略。这种方法确保热神经元更有可能留在缓存中，而冷神经元则经常被退出并根据需要从闪存中加载。重要的是，退出过程不涉及写入存储，而只是从内存中丢弃权重。

PowerInfer-2 利用经典的双队列方法来实现其 LRU 神经元缓存，该方法以单个神经元粒度来管理 LLM 权重。系统维护两个双向链表队列，标记为激活和非激活，其中队列中的神经元顺序由其最近访问的时间决定，最近访问的神经元位于队列的头部。

在运行时，所有神经元最初都加入非激活队列。重新访问时，它们被提升到激活队列的前面。激活队列中的神经元在后续访问时被移到头部。为了管理缓存容量，当激活队列填满缓存空间的 90% 时，激活队列尾部的神经元将移动到非激活队列，直到激活队列的占用率降至 90% 以下。如果缓存达到容量，则丢弃非激活队列尾部的神经元以腾出空间容纳新条目。

灵活的神经元上载

尽管配备了可以有效存储激活神经元的神经元缓存，但推理过程仍然不可避免地会产生未缓存神经元的 I/O 操作。为了优化 I/O 读取吞吐量并最大限度地减少 I/O 操作，PowerInfer-2 还捆绑相关神经元。尽管一旦移除热神经元，单个 FFN 权重矩阵内的共同激活就会变得不常见，但不同矩阵中相应位置的神经元通常会一起激活。例如，Gate、Up 和 Down 矩阵中第 i 个神经元的共同激活概率高达 80%。因此，PowerInfer-2 选择根据神经元粒度而不是矩阵结构来存储神经元权重，将 Gate、Up 和 Down 矩阵中第 i 个神经元的权重连接成一个条目。

PowerInfer-2 进一步为不同模型引入了不同的 I/O 加载策略，考虑到量化方法和 UFS I/O 的固有特性。对于未量化的模型，由于每个神经元占用的存储空间较大，PowerInfer-2 使用更大粒度的随机读取来提高 I/O 带宽。例如，Llama-7B-FP16 中的单个神经元占用 8KB，Gate、Up 和 Down 矩阵的神经元总大小为 24KB。PowerInfer-2 通过一次随机 I/O 读取即可高效地将整个 24KB 激活包传输到内存中。

对于 4 位量化模型，包大小设置为 8KB。以 Llama-7B 模型为例，其中每个神经元量化为 4 位精度并占用 2.5KB（量化的 int4 值为 2KB，量化组的 FP16 尺度为 0.5KB），组合包大小达到 7.5KB。为了与存储介质的最小读取粒度 4KB 保持一致，PowerInfer-2 补充了额外的 0.5KB，将总数四舍五入为 8KB。但是，PowerInfer-2 不是在单个 I/O 操作中加载这些 8KB 包，而是选择 4KB 粒度。

此外，考虑到这些包内 80% 的共同激活可能性，这些捆绑的神经元未被共同激活的概率仍有近 20%。因此，组合这两个 4KB 随机读取可能会导致带宽浪费。为了缓解这种情况，对于使用 4 位量化的模型，PowerInfer-2 会延迟第二个 4KB 读取，直到获得门-神经元乘法的结果。具体来说，PowerInfer-2 使用预测器来确定门矩阵中神经元的激活，并根据此信息启动包的第一部分加载。之后，如果门-神经元的输出（激活函数）非零，PowerInfer-2 将继续加载包的第二部分，从而最大限度地减少不必要的 I/O 操作。

神经元簇-级的流水线

PowerInfer-2 还旨在通过将计算与 I/O 活动重叠来隐藏 I/O 开销。一种直接的方法是矩阵级重叠，它发出 I/O 命令从存储中检索矩阵神经元，同时处理内存中已有的神经元。当从存储中加载神经元时，它们会立即被处理。虽然这种矩阵级重叠方法可以在一定程度上隐藏计算过程中 I/O 操作的成本，但它仍然需要系统等待矩阵内所有神经元（包括从存储中获取的神经元）完成，然后才能继续下一个。如图 （a） 所示，假设一个矩阵包含 8 个神经元簇，其中 4 个驻留在内存中，其余 4 个驻留在存储中。一部分 I/O 操作可以隐藏在缓存的神经元计算后面。但由于 I/O 时间较长，仍然会出现 CPU 核心必须等待 I/O 完成的情况。

为了消除 I/O 操作的等待时间，PowerInfer-2 引入了神经元簇-级流水线机制。该机制基于一个洞察：以神经元簇为粒度，可以将来自多个矩阵的神经元簇计算中的 I/O 操作重叠。具体来说，PowerInfer-2 打破矩阵计算之间的障碍；一旦一个神经元簇完成计算，它就会立即开始计算内存中的下一个矩阵中神经元簇。该机制有效地减少了等待气泡，如图（b） 所示。

PowerInfer-2将一个神经元簇的执行过程分为5个连续的阶段，分别是：确定Gate、Up、Down矩阵的行/列是否通过预测器被激活（Pred）、从存储器中读取Gate矩阵行的权重（GIO）、计算Gate矩阵行与输入向量的乘积（GC）、从存储器中读取Up和Down矩阵的行/列（UDIO）、分别计算Up和Down矩阵行/列与输入向量的乘积（UDC）。

PowerInfer-2创建多个计算线程和一个I/O线程来分别处理这5个阶段的计算和I/O操作。这些线程的具体数量以及它们将在哪些核心上执行由离线规划器决定。如图显示了计算和I/O线程如何工作以实现神经元簇流水线。在每个FFN块的开始时，所有神经元最初都在预测阶段，并被插入到计算队列中。计算线程处理这些神经元，仅将激活的神经元推进到后续阶段。然后，这些激活的神经元被合并到神经元簇中。如果神经元簇的门权重在内存中可用，则神经元簇将进入 G-C 阶段并返回到计算队列。如果没有，则将其设置为 G-IO 并移至 I/O 队列。同时，计算线程继续处理队列中的下一个神经元簇。并行地，I/O 线程从 I/O 队列中取出神经元，根据需要执行 I/O 任务。UDIO 和 UDC 的执行遵循与 GC 和 GIO 类似的模式。