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

TL;DR

PAPI tackles the problem of dynamically changing parallelism in LLM decoding by combining online kernel characterization with a heterogeneous PIM-enabled system that includes memory-centric Attn-PIM and compute-centric FC-PIM units. The core idea is to monitor runtime parallelism (RLP and TLP) and assign FC kernels to either PIM or GPU/PUs based on estimated arithmetic intensity, while always running attention on Attn-PIM, thereby balancing compute throughput and memory bandwidth. The approach yields significant end-to-end speedups (up to 11.1× over AttAcc-only baselines) and substantial energy efficiency gains (up to 3.4×) across multiple models and datasets, demonstrating the practicality of dynamic, fine-grained kernel offloading in real-world LLM serving. By addressing MoEs and scalability through a hybrid, disaggregated PIM design, PAPI offers a pathway to scalable, energy-efficient LLM inference in heterogeneous architectures.

Abstract

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.

Figures (12)

• Figure 1: (a) LLM structure. (b) LLM inference with serial decoding. (c) Parallel decoding process with batching. (d) Parallel decoding process with speculative decoding.
• Figure 2: Roofline model using OPT-30B with (a) different batch sizes (speculation length = 8) and (b) different speculation lengths (batch size = 32). The darker the color of the dots, the higher the degree of parallelism.
• Figure 3: Decoding iterations required for each request in a batch, illustrating how the number of remaining parallel requests changes as decoding iterations increase.
• Figure 4: The normalized latency of the FC kernel in LLM inference with different parallelization levels (different batch sizes and speculation lengths).
• Figure 5: Overview of the PAPI computing system, and an example of its dynamic parallelism-aware scheduler.