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

📄 arXiv: 2502.15470v2 📥 PDF

作者: Yintao He, Haiyu Mao, Christina Giannoula, Mohammad Sadrosadati, Juan Gómez-Luna, Huawei Li, Xiaowei Li, Ying Wang, Onur Mutlu

分类: cs.AR, cs.AI, cs.DC, cs.LG

发布日期: 2025-02-21 (更新: 2025-02-27)

备注: To appear in ASPLOS 2025

💡 一句话要点

PAPI：利用PIM的动态并行加速大语言模型解码

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 大语言模型 LLM解码 内存计算 PIM 异构计算 动态调度 在线特征化

📋 核心要点

  1. 现有LLM解码加速方法静态映射内核到异构架构，忽略了内核计算/访存特性随参数动态变化的问题。
  2. PAPI通过在线内核特征化动态调度内核到合适的硬件单元，并设计了支持PIM的异构计算系统。
  3. 实验表明，PAPI在三个常用LLM上分别比最先进的异构和仅PIM加速器实现了1.8倍和11.1倍的加速。

📝 摘要（中文）

大型语言模型（LLMs）被广泛应用于自然语言理解和文本生成。LLM模型依赖于一个耗时的步骤，即LLM解码来生成输出tokens。一些先前的工作侧重于使用并行技术（如批处理和推测解码）来提高LLM解码的性能。目前最先进的LLM解码同时具有计算密集型和内存密集型内核。一些先前的工作静态地识别这些不同的内核，并将它们映射到由内存计算（PIM）单元和以计算为中心的加速器组成的异构架构。我们观察到，LLM解码内核的特性（例如，内核是否是内存密集型）会由于参数变化而动态变化，以满足用户和/或系统的需求，这使得（1）静态内核映射到PIM单元和以计算为中心的加速器变得次优，并且（2）由于即使在内存密集型内核中也存在很大程度的异构性，因此设计PIM单元的“一刀切”方法效率低下。在本文中，我们的目标是在考虑所涉及内核的动态变化特性的同时加速LLM解码。我们提出了PAPI（使用PIM的并行解码），这是一种支持PIM的异构架构，它利用计算密集型或内存密集型内核到合适的硬件单元的动态调度。PAPI具有两个关键机制：（1）在线内核特征化，以在运行时将内核动态调度到最合适的硬件单元，以及（2）支持PIM的异构计算系统，该系统和谐地协调以计算为中心的处理单元和具有不同计算能力的混合PIM单元。我们在三个广泛使用的LLM上的实验结果表明，PAPI分别比最先进的异构LLM加速器和最先进的仅PIM LLM加速器实现了1.8倍和11.1倍的加速。

🔬 方法详解

问题定义：现有的大语言模型解码加速方法，如静态内核映射到异构架构，无法适应LLM解码过程中计算和访存特性的动态变化。这种静态映射导致硬件资源利用率低下，无法充分发挥PIM等加速技术的潜力。此外，即使是内存密集型内核也存在异构性，采用“一刀切”的PIM设计效率不高。

核心思路：PAPI的核心思路是根据LLM解码过程中内核的动态特性，在线地将内核调度到最合适的硬件单元上执行。通过动态调度，PAPI能够充分利用异构计算资源，包括传统的计算单元和PIM单元，从而实现更高的解码效率。这种动态调度策略能够更好地适应LLM参数变化带来的影响，提高加速器的整体性能。

技术框架：PAPI的整体架构是一个支持PIM的异构计算系统，包含以下主要模块：1) 在线内核特征化模块：该模块在运行时分析LLM解码内核的计算和访存特性，确定其属于计算密集型还是内存密集型。2) 动态调度器：根据内核的特性，动态调度器将内核分配到最合适的硬件单元上执行。3) 异构计算单元：系统包含传统的以计算为中心的处理单元和具有不同计算能力的混合PIM单元。这些单元协同工作，共同完成LLM解码任务。

关键创新：PAPI的关键创新在于其动态调度策略和支持PIM的异构架构。传统的静态映射方法无法适应LLM解码过程中内核特性的动态变化，而PAPI通过在线内核特征化和动态调度，能够更好地利用异构计算资源。此外，PAPI的异构架构包含多种类型的PIM单元，能够更好地适应不同类型内核的需求。

关键设计：PAPI的关键设计包括：1) 在线内核特征化算法：该算法需要高效准确地识别内核的计算和访存特性。2) 动态调度策略：调度策略需要考虑不同硬件单元的性能特点和内核的需求，以实现最佳的资源分配。3) PIM单元的设计：PIM单元需要根据LLM解码的需求进行优化，例如，支持高效的矩阵乘法和向量加法等操作。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

PAPI在三个广泛使用的LLM（具体模型未知）上进行了实验，结果表明，PAPI分别比最先进的异构LLM加速器和最先进的仅PIM LLM加速器实现了1.8倍和11.1倍的加速。这些结果表明，PAPI的动态调度策略和支持PIM的异构架构能够显著提高LLM解码的性能。

🎯 应用场景

PAPI可应用于各种需要加速大语言模型解码的场景，例如云端推理服务、边缘设备上的自然语言处理应用等。通过提高解码效率，PAPI能够降低推理延迟，提升用户体验，并降低计算成本。未来，PAPI有望推动大语言模型在更多实际场景中的应用。

📄 摘要（原文）

Large language models (LLMs) are widely used for natural language understanding and text generation. An LLM model relies on a time-consuming step called LLM decoding to generate output tokens. Several prior works focus on improving the performance of LLM decoding using parallelism techniques, such as batching and speculative decoding. State-of-the-art LLM decoding has both compute-bound and memory-bound kernels. Some prior works statically identify and map these different kernels to a heterogeneous architecture consisting of both processing-in-memory (PIM) units and computation-centric accelerators. We observe that characteristics of LLM decoding kernels (e.g., whether or not a kernel is memory-bound) can change dynamically due to parameter changes to meet user and/or system demands, making (1) static kernel mapping to PIM units and computation-centric accelerators suboptimal, and (2) one-size-fits-all approach of designing PIM units inefficient due to a large degree of heterogeneity even in memory-bound kernels. In this paper, we aim to accelerate LLM decoding while considering the dynamically changing characteristics of the kernels involved. We propose PAPI (PArallel Decoding with PIM), a PIM-enabled heterogeneous architecture that exploits dynamic scheduling of compute-bound or memory-bound kernels to suitable hardware units. PAPI has two key mechanisms: (1) online kernel characterization to dynamically schedule kernels to the most suitable hardware units at runtime and (2) a PIM-enabled heterogeneous computing system that harmoniously orchestrates both computation-centric processing units and hybrid PIM units with different computing capabilities. Our experimental results on three broadly-used LLMs show that PAPI achieves 1.8$\times$ and 11.1$\times$ speedups over a state-of-the-art heterogeneous LLM accelerator and a state-of-the-art PIM-only LLM accelerator, respectively.