PAPI: Exploiting Dynamic Parallelism in Large Language Model Decoding with a Processing-In-Memory-Enabled Computing System

发表时间: 2025-03 · arXiv:2502.15470 (ASPLOS 2025)

原文: https://arxiv.org/abs/2502.15470

作者/机构：Yintao He (SKLP, ICT CAS; UCAS), Haiyu Mao (King’s College London; ETH Zürich), Christina Giannoula (Univ of Toronto; Vector Institute; ETH Zürich), Mohammad Sadrosadati (ETH Zürich), Juan Gómez-Luna (NVIDIA), Huawei Li (SKLP, ICT CAS; UCAS), Xiaowei Li (SKLP, ICT CAS; UCAS), Ying Wang (SKLP, ICT CAS), Onur Mutlu (ETH Zürich)

速读

一句话结论 本文提出了一种名为 PAPI 的支持存内计算（PIM）的异构计算架构，通过在运行时动态感知大语言模型解码阶段并行度的变化，将计算密集型和访存密集型算子实时调度到最合适的硬件单元，从而在真实推理场景下实现了最高 11.1 倍的速度提升。

要解决什么问题 现有的大语言模型推理加速方案通常采用静态调度策略，将全连接（FC）算子固定分配给 GPU 等计算中心处理器，将多头注意力（Attention）算子固定分配给存内计算单元。这种做法卡在无法适应真实服务场景中动态变化的解码并行度。在实际推理中，由于用户服务级别目标限制、内存容量限制以及动态批处理机制，请求级并行度（即同时处理的请求数）会不断波动；同时，投机解码技术也会根据运行时状态动态调整词元级并行度（即单次迭代生成的词元数）。这两种并行度的变化会直接改变算子的算术强度。例如，当请求级和词元级并行度较小时，FC 算子受限于内存带宽，属于访存密集型；但当并行度增大时，FC 算子的算术强度随之升高，会动态转变为计算密集型。而 Attention 算子虽然始终是访存密集型，但其对内存带宽和容量的需求也与 FC 算子截然不同。原有的静态映射和单一规格的 PIM 硬件设计，会导致计算密集型算子被错误地塞给算力不足的 PIM，或者访存密集型算子被留在 GPU 导致显存带宽瓶颈，最终造成严重的算力闲置和延迟激增。

怎么做的 核心思路是摒弃一刀切的静态分配，设计一个包含主机 CPU、GPU 和混合 PIM 单元的异构系统，并在运行时根据算子的实时算术强度将其动态路由到最匹配的硬件上。这一机制能有效绕开静态调度导致的资源错配卡点。系统由三个关键部件构成。首先是动态并行度感知调度器，它负责在每次解码迭代后，通过统计当前批次中未完成请求的数量 $RLP$ 和系统设定的投机长度 $TLP$，以极低的硬件开销在线预测 FC 算子的算术强度。其核心判定公式为：

$$ \mathrm{AI} \approx RLP \times TLP $$

调度器将该估算值与预设的内存受限阈值 $\alpha$ 进行比较，若大于 $\alpha$ 则判定为计算密集型，将其调度给 GPU 执行；若小于 $\alpha$，则判定为访存密集型，下发给 PIM 执行。其次是针对 FC 算子设计的 FC-PIM 单元，由于 FC 算子在低并行度下仍有较高的计算需求，该单元采用了高并发配置（每个 DRAM 库配置 4 个浮点运算单元），并通过复用 DRAM 激活行的数据来降低访存功耗，使其在满足 HBM 面积和功耗约束的前提下提供充足算力。为了确保单个 HBM 裸片的总面积不超过上限 $A_{Max}$，其内存库数量 $m$ 与单库浮点单元数 $n$ 需满足面积约束公式：

$$ m ( n \times A_{FPU} + A_{bank} ) \leq A_{Max} $$
最后是针对 Attention 算子设计的 Attn-PIM 单元，由于 Attention 算子始终是访存密集型且键值缓存的内存占用随序列长度线性增长，该单元采用了大容量、低并发配置（每两个 DRAM 库共享 1 个浮点运算单元），并在物理上与 GPU 解耦，通过 PCIe 或 CXL 总线连接，从而灵活扩展内存容量。这三个部件各司其职：调度器充当流量探针与路由，FC-PIM 吸收低并行度下的 FC 访存压力，Attn-PIM 专职吞吐海量的键值缓存，GPU 则作为高并行度下 FC 算子的算力担当。

效果如何 实验评估使用了 LLaMA-65B、GPT-3 66B 和 GPT-3 175B 三个模型，在 Dolly 数据集的问答和创意写作任务上进行测试。硬件基准配置为 6 张 NVIDIA A100 GPU 搭配 90 个 HBM3 内存设备。对比基线包含三种代表性路线：第一种是 A100+AttAcc，代表最先进的异构静态调度路线（GPU 跑 FC，PIM 跑 Attention）；第二种是 A100+HBM-PIM，代表采用三星商业化 PIM 硬件的异构路线；第三种是 AttAcc-only，代表纯 PIM 计算路线（所有算子均在 PIM 上执行）。量化结果显示，在创意写作任务（具有较长输出长度和剧烈的并行度动态变化）中，PAPI 相比 A100+AttAcc、A100+HBM-PIM 和 AttAcc-only 分别实现了 1.8 倍、1.9 倍和 11.1 倍的端到端速度提升。在能效方面，PAPI 相比 A100+AttAcc 路线平均提升了 3.4 倍，因为它将部分 FC 算子卸载到了低功耗的 PIM 核心上，减少了高能耗的 GPU 数据搬运。纯 PIM 路线（AttAcc-only）在并行度极低时表现尚可，但一旦并行度拉高，其孱弱的算力会导致性能严重崩塌，这也反向证明了 PAPI 动态调度的必要性。代价与局限性方面，作者指出在纯 PIM 模式下，设备间的通信时间占到了解码阶段总执行时间的 28.2%，这意味着在极大规模部署时，互联总线带宽可能成为新的瓶颈，未来需要集成更先进的网络技术来进一步释放硬件潜力。

主要贡献

大型语言模型（LLMs）的推理过程依赖于耗时的解码阶段。当前的先进解码技术（如批处理和推测解码）会使得解码内核在计算受限（compute-bound）和内存受限（memory-bound）之间动态切换。现有研究通常静态地识别这些内核，并将其固定映射到由存内计算（PIM）单元和以计算为中心的加速器（如GPU）组成的异构架构中。这种静态映射策略存在两个核心问题：第一，由于用户需求和系统参数的动态变化，内核特性会实时改变，导致静态调度次优；第二，即使同为内存受限内核，其异构程度也很高，采用“一刀切”的统一PIM单元设计效率低下。

为了解决上述问题，本文提出了PAPI（PArallel Decoding with PIM），这是一种支持PIM的异构架构，旨在通过将计算受限或内存受限的内核动态调度到最合适的硬件单元来加速LLM解码。本文的核心创新点包括：
1. 在线内核特征表征与动态调度：提出了一种动态并行感知任务调度框架，通过低硬件开销的内核瓶颈预测器，在运行时动态地将内核调度到最合适的硬件单元。
2. 支持PIM的异构计算系统：设计了一个协调调度以计算为中心的处理单元（GPU）和混合PIM单元的架构。
3. 混合PIM架构设计：系统包含两种不同类型的PIM单元（性能优化型和内存容量优化型），以满足具有不同计算需求和内存占用的内存密集型内核。

实验结果表明，在LLaMA-65B、GPT-3 66B和GPT-3 175B模型上，PAPI相较于最先进的异构LLM加速器（GPU+PIM）和纯PIM加速器，分别实现了 $1.8\times$ 和 $11.1\times$ 的性能加速。

背景知识与研究动机

LLM结构与推理阶段。LLM结构包含多个基于Transformer的解码器。每个解码器包含四个内核：QKV生成、多头注意力、投影和前馈网络（FC）。这些内核可分为全连接（FC）层和多头注意力层，均由通用矩阵向量乘法（GEMV）计算组成。LLM推理分为预填充（Prefill）和解码（Decoding）两个阶段。预填充阶段同时处理输入序列中的多个词元以生成第一个输出词元。解码阶段则串行地逐个生成输出词元，直到遇到 <eos> 词元。相比预填充，解码阶段占据了端到端推理的绝大部分时间。在串行解码中，系统需要频繁地在片外内存和片上缓存之间搬运庞大的权重矩阵和KV缓存，导致极高的数据移动开销。
图1：(a) LLM结构。(b) 串行解码的LLM推理。(c) 带有批处理的并行解码过程。(d) 带有推测解码的并行解码过程。

LLM推理中的优化技术。为了克服串行解码的性能瓶颈，研究人员开发了批处理（Batching）【15-17, 31，Orca: A distributed serving system for transformer-based generative models等+2022/2023/2024+OSDI/arXiv/ISCA/NeurIPS】和推测解码（Speculative Decoding）【18-21，Fast inference from transformers via speculative decoding等+2023/2024+ICML/arXiv/ACL】。批处理允许单个解码步骤并发生成来自不同用户请求的多个词元，从而实现请求级并行（RLP）。混合连续批处理还可以根据资源动态调整批次大小。推测解码通过小草稿模型预测后续词元，再由大模型并行验证，从而实现词元级并行（TLP），即在单次解码迭代中并发解码多个词元。

LLM推理的算术强度分析。通过在拥有312 TFLOPS算力和 $1935 GB/s$ 内存带宽的NVIDIA A100 GPU上对OPT-30B模型进行Roofline模型分析，发现当批大小较小（如4、8、16）时，解码阶段受限于内存，FC和注意力内核均被内存带宽卡脖子。当批大小 $\geq 32$ 时，FC内核转变为计算受限，而注意力内核依然是内存受限的，因为批处理对注意力内核没有数据复用效果。进一步改变推测长度（批大小固定为32）发现，FC和注意力内核的算术强度均随推测长度增加而上升。当推测长度超过6时，FC内核变为计算受限；而注意力内核的算术强度虽有小幅提升，但仍保持内存受限状态。
图2：使用OPT-30B的Roofline模型，(a) 不同的批大小（推测长度 $\sigma = 8$）和(b) 不同的推测长度（批大小 $\Lambda = 32$）。点的颜色越深，并行度越高。

LLM推理中并行度的动态变化。在真实的LLM任务中，批大小和推测长度在运行时会发生显著变化。初始请求级并行度（Initial RLP）受限于服务级别目标（SLO，如延迟要求）、系统内存容量限制（长序列需要更多KV缓存）以及动态批处理机制的触发条件。运行时请求级并行度（Runtime RLP）在静态批处理中会随着短序列请求的完成而逐渐降低；在混合连续批处理中，也会因新请求的动态加入而不断波动。同时，词元级并行度（TLP）也会在运行时被动态调整以优化系统性能。
图3：批处理中每个请求所需的解码迭代次数，说明了随着解码迭代次数的增加，剩余并行请求的数量如何变化。

现有PIM架构的局限性。现有的PIM异构架构【22-26，Attacc! unleashing the power of pim for batched transformer-based generative model inference等+2024+ASPLOS/arXiv】存在两个主要缺陷。第一，现有工作静态地将FC和注意力内核分配给GPU或PIM。然而，在低并行度配置下（如批大小1、推测长度8），HBM-PIM和AttAcc的性能优于A100 GPU；但在高并行度配置下（如批大小 $\geq 16$），A100 GPU的性能则显著优于PIM架构。这种并行度的不可预知性要求系统必须具备动态调度能力。第二，现有工作仅支持单一计算和内存带宽能力的PIM设备。由于FC内核和注意力内核在算术强度上差异巨大（如在批大小4、推测长度8时，FC的算术强度为31.7 FLOPs/Byte，而注意力为7.0 FLOPs/Byte），单一PIM设备无法同时高效满足这两种内核截然不同的计算和内存带宽需求。
图4：具有不同并行化级别（不同批大小和推测长度）的LLM推理中FC内核的归一化延迟。

方法细节

PAPI架构核心组件。PAPI框架包含三个关键部分。首先是异构架构，该架构由主机CPU、配备PIM内存单元（FC-PIM）的高性能处理器以及物理隔离的PIM单元（Attn-PIM）组成。高性能处理器内部包含处理单元（PUs，如GPU张量核心）、FC-PIM设备和硬件调度器。其次是混合PIM单元，PAPI设计了两种不同的PIM单元：FC-PIM提供较高的计算能力以服务于FC内核，而Attn-PIM提供更大的内存容量以专门处理注意力内核。最后是动态并行感知调度，硬件调度器在运行时监控并行化水平的变化。由于注意力内核始终是内存受限的，它被固定分配给Attn-PIM；而FC内核则由调度器通过低成本的识别步骤，动态地卸载到FC-PIM（当其内存受限时）或高性能处理器的PUs（当其计算受限时）上执行。
图5：PAPI计算系统概述，及其动态并行感知调度器的示例。

FC内核的内存受限识别。调度机制通过估计FC内核的算术强度来判断其是否受限于内存。假设FC内核的权重矩阵维度为 $(h, h)$，输入维度为 $(RLP \times TLP, h)$，其中 $h$ 为隐藏层维度。其算术强度 $AI$ 的计算公式为：
$AI = \frac{RLP \times TLP \times h^2 \times 2}{(2 \times RLP \times TLP \times h + h^2) \times 2}$
在先进的LLM中，$h$ 通常非常大（如GPT-3 175B中 $h = 12288$），因此算术强度可以被精确近似为 $AI \approx RLP \times TLP$。在GPT-3 66B模型的评估中，该估计值与实际测量的算术强度高度吻合。仅在极高并行度（如 $RLP = 128$）时，估计值略大于实际值，但由于此时内核已绝对处于计算受限状态，这种微小偏差不会影响调度决策的准确性。
图6：GPT-3 66B模型中FC内核的实际测量算术强度与估计算术强度。

运行时调度的具体实现。基于估计的算术强度，调度过程在主机CPU上分为初始调度和运行时调度两个步骤。在初始调度阶段，系统在LLM服务启动前，将 $RLP$ 设为批大小，$TLP$ 设为系统预设的推测长度。系统计算 $RLP \times TLP$ 并将其与通过离线迭代评估得出的内存受限阈值 $\alpha$ 进行比较。如果估计值大于 $\alpha$，则判定FC内核为计算受限并分配给PUs；否则判定为内存受限并分配给FC-PIM。在运行时调度阶段，系统采用词元级调度方案跟踪并行度变化。先在每次解码后，将当前批次所有请求的输出词元收集到一个向量中。再统计该向量中 <eos> 词元的数量，若数量大于零，说明有请求已完成，据此更新 $RLP$。同时，通过主机CPU监控系统软件对专用寄存器中 $TLP$ 值的修改指令。接着，系统计算最新的 $RLP \times TLP$ 以预测下一次解码的算术强度。最后，将该预测值与阈值 $\alpha$ 比较，以决定是否需要将FC内核在PUs和FC-PIM之间进行重新调度。

FC-PIM的硬件设计。为了满足FC内核的计算需求，FC-PIM需要具备较高的计算并行度并满足功耗约束。在传统的1P1B（每Bank一个处理核心）PIM设计中，DRAM访问能耗占总能耗的96.7%。通过利用批处理和推测解码带来的数据复用特性，当一次DRAM访问被用于64次计算时，DRAM访问能耗占比大幅降至33.1%。这种数据复用使得在HBM功耗预算内增加FPU数量成为可能。当数据复用级别 $\geq 4$ 时，4P1B（每Bank四个FPU）配置的功耗显著降低并完全符合HBM功耗标准。为了在单颗HBM裸片严格的面积限制（$121 mm^2$）内容纳额外的FPU，系统通过减少内存容量来释放空间。设HBM裸片有 $m$ 个Bank，每个Bank有 $n$ 个FPU，总面积需满足 $m(n \times A_{FPU} + A_{bank}) \leq A_{Max}$。代入4P1B参数计算得出，最大Bank数必须小于97。因此，FC-PIM设计为每个HBM内存单元包含96个Bank（即8层堆叠中的3个Bank Group）。
图7：(a) 无DRAM数据复用时执行FC内核的PIM能耗分解。(b) 当一次DRAM访问被用于64次计算（即数据复用级别 $\delta = 64$）时，执行FC内核的PIM能耗分解。(c) 具有不同数据复用级别和每个Bank不同数量FPU的PIM架构的功耗。

Attn-PIM的硬件设计。注意力内核的计算强度较低，但其KV缓存的内存占用随序列长度线性急剧增加。为此，PAPI设计了独立且解耦的Attn-PIM单元。为了在固定面积预算内平衡计算与内存，Attn-PIM采用了1P2B（每两个Bank共享一个FPU）的低计算并行度配置，这有效保证了其在处理缺乏数据复用的注意力内核时不会超出功耗预算。同时，通过将大量Attn-PIM设备与高性能处理器解耦连接，系统能够灵活提供庞大的内存容量，以支持生成大量输出词元的长序列请求。

系统互连与数据分区。在系统集成方面，由于FC-PIM设备需要向处理单元传输海量的权重参数，因此采用NVLink等高速互连技术以防止数据传输成为瓶颈。而Attn-PIM设备主要处理字节级的Q向量等小数据传输，因此采用PCIe或CXL等标准互连即可满足带宽需求且更具成本效益。在数据分区方面，注意力内核的各个注意力头被分配到不同的Attn-PIM单元中。$K^T$ 矩阵在伪通道和Bank-Group级别按列划分，在Bank和乘法器级别按行划分；$V$ 矩阵则采取相反的划分方式。对于FC内核，庞大的权重矩阵被切分为2D块映射到各个FC-PIM设备，并在设备内部采用与 $K^T$ 矩阵相同的划分策略。

架构的实用性与对新兴模型的支持。FC-PIM和Attn-PIM共享相同的Bank级计算结构和内存层次，计算逻辑嵌入在外围电路中，避免了对DRAM核心阵列的修改，确保了与现有HBM技术的兼容性和易部署性。此外，针对混合专家（MoE）模型，FC-PIM能够通过在同一个DRAM Bank中存储不同专家的权重切片，高效执行稀疏操作，从而最小化因稀疏性导致的FPU闲置，并降低推理的能耗和延迟。

实验环境

  • 模型与数据集：评估了三种基于Transformer的大型语言模型：LLaMA-65B、GPT-3 66B和GPT-3 175B，采用FP16数据类型。数据集使用Dolly数据集中的creative-writing（长输出）和general-qa（常规问答）任务。
  • 硬件与软件配置：开发了基于Ramulator2和AttAcc的模拟器。对比基线包括：
    1. A100+AttAcc：6块NVIDIA A100 GPU + AttAcc PIM单元（1P1B）。GPU执行FC，PIM执行注意力。
    2. A100+HBM-PIM：6块A100 GPU + 三星HBM-PIM（1P2B）。
    3. AttAcc-only：纯PIM平台，所有内核均在AttAcc PIM单元上运行。
    4. PAPI配置：FC-PIM容量为12GB，Attn-PIM容量为16GB。系统共配备90个HBM设备（30个用于FC权重，60个用于注意力）。GPU通过NVLink连接。所有HBM均为HBM3，单引脚速率5.2Gbps，运行频率333MHz。

实验结果

端到端性能与能效。

  • 在Dolly creative-writing数据集上，PAPI相较于A100+AttAcc、A100+HBM-PIM和AttAcc-only，分别实现了 $1.8\times$、$1.9\times$ 和 $11.1\times$ 的端到端性能加速。这一显著提升归功于PAPI能够根据实时并行度动态调度任务，且其混合PIM架构能精准匹配不同内核的计算需求。AttAcc-only在大多数配置下表现最差，因为纯PIM架构无法满足高并行度下FC内核庞大的计算需求。
  • 在Dolly general-qa数据集上，PAPI的加速比分别为 $1.7\times$、$1.7\times$ 和 $8.1\times$。由于该数据集输出长度较短，解码阶段的瓶颈效应相对较弱，且并行度动态变化幅度较小，因此加速效果略低于creative-writing数据集。
  • 在能效方面，PAPI在两个数据集上相较于A100+AttAcc分别提升了 $3.4\times$ 和 $3.1\times$。这是因为PAPI将部分FC内核从高能耗的GPU卸载到了低功耗的FC-PIM设备上，大幅减少了数据移动。相较于AttAcc-only，PAPI也获得了 $1.15\times$ 和 $1.01\times$ 的微弱能效优势。
    图8：四个评估设计在Dolly creative-writing数据集上的端到端加速比（上）和能效（下）比较。值归一化为A100+AttAcc。
    图9：三个评估设计在Dolly general-qa数据集（GPT-3 175B）上的端到端加速比（a）和能效（b）比较。

对并行化水平的敏感性分析。
* 请求级并行度（RLP）：在固定推测长度为1，批大小从4增加到128时，PAPI在所有RLP设置下均保持最佳性能。当批大小较小时，AttAcc-only性能优于A100+AttAcc；但随着RLP增加，FC内核转为计算密集型，AttAcc-only的执行时间急剧增加，表现远逊于A100+AttAcc。
* 词元级并行度（TLP）：在固定批大小为4，推测长度从1增加到8时，PAPI平均比A100+AttAcc和AttAcc-only快 $1.5\times$ 和 $3.0\times$。随着TLP的增加，PAPI相对于A100+AttAcc的加速比逐渐收敛，因为系统智能地将越来越多转变为计算受限的FC内核调度回了GPU。
图10：LLaMA-65B在(a) 不同批大小（推测长度 $\sigma=1$）和(b) 不同推测长度（批大小 $\Lambda=4$）下的端到端加速比。

PAPI纯PIM架构性能分析。
为了单独验证混合PIM设计的有效性，构建了移除GPU、仅包含FC-PIM和Attn-PIM的“纯PIM版PAPI”，并与AttAcc-only进行对比。结果显示，纯PIM版PAPI平均实现了 $2.3\times$ 的加速。在更高并行度（如批大小64，推测长度4）下，加速比扩大至 $2.7\times$，证明了FC-PIM（4P1B）能提供更充沛的算力。对每个词元的执行时间分解显示，FC内核主导了总执行时间。纯PIM版PAPI在处理FC内核时速度提升了 $2.9\times$；尽管注意力内核在Attn-PIM（1P2B）上的运行速度比在AttAcc-only（1P1B）上慢了 $1.7\times$，但这种架构妥协成功控制了面积和功耗，且整体性能依然大幅领先。
图11：在Dolly creative-writing数据集的解码阶段，仅PIM的PAPI相对于AttAcc-only的性能加速比。
图12：LLaMA-65B模型推理（批大小=4，推测长度=4）解码阶段每个词元的执行时间分解，对比AttAcc-only与仅PIM的PAPI。

相关工作补充

在支持PIM的LLM加速器领域，现有工作如AiM【56，System architecture and software stack for gddr6-aim+2022+HCS】、AttAcc【23，Attacc! unleashing the power of pim for batched transformer-based generative model inference+2024+ASPLOS】和IANUS【25，Ianus: Integrated accelerator based on npu-pim unified memory system+2024+ASPLOS】多采用静态映射策略，无法适应大批次等计算受限场景。SpecPIM【24，Specpim: Accelerating speculative inference on pim-enabled system via architecture-dataflow co-exploration+2024+ASPLOS】提出了基于遗传算法和蒙特卡洛树搜索的调度方法，但其离线搜索的计算复杂度极高，在并行度动态变化的真实场景中反复运行会导致不可接受的性能开销。在其他LLM加速器方面，DFX【11，Dfx: A low-latency multi-fpga appliance for accelerating transformer-based text generation+2022+MICRO】和AMXGPU【92，Exploiting intel® advanced matrix extensions (amx) for large language model inference+2024+IEEE CAL】等架构未能考虑运行时的并行度动态变化。此外，SpAtten【103，Spatten: Efficient sparse attention architecture with cascade token and head pruning+2021+HPCA】等利用剪枝和量化的近似算法会牺牲模型输出质量。PAPI则是首个在不牺牲模型质量的前提下，通过轻量级运行时框架解决LLM动态计算和内存需求波动的架构。

结论

现实世界中的LLM服务广泛采用批处理和推测解码等先进并行优化技术，这导致推理过程中的并行化水平发生动态变化。因此，LLM推理中的全连接内核和注意力内核表现出不断波动的计算和内存需求。为了无缝适应这种动态需求，本文提出了PAPI。该计算系统支持三种具有不同计算和内存带宽能力的计算单元，并配备了一个轻量级的调度框架，能够以极低的开销监控LLM推理中的动态并行化水平，将全连接和注意力内核卸载到最合适的计算单元。评估结果表明，PAPI相比最先进的LLM推理系统实现了 $1.8\times$ 到 $11.1\times$ 的性能提升。本工作为利用异构PIM系统应对新兴机器学习模型在真实动态执行场景中的挑战开辟了新的研究方向。