PaperlayerView exact revision on arXiv

Source-linked AI summary

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

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

arXiv:2502.15470v2cs.ARcs.AIcs.DCcs.LG

TL;DR

LLM decoding is a major inference bottleneck whose kernel characteristics change with runtime parallelism, limiting static hardware mappings and fixed PIM designs. PAPI combines online kernel characterization with heterogeneous GPUs, CPUs, and hybrid PIM units, achieving 1.8× and 11.1× speedups over state-of-the-art heterogeneous and PIM-only accelerators.

  • Problem

    Changing batch size and speculation length alter kernel arithmetic intensity, making static mappings and single-capability PIM designs inadequate for varying decoding demands.

  • Method

    PAPI uses lightweight runtime scheduling based on online kernel characterization across GPUs, host CPUs, and performance-optimized and memory-capacity-optimized PIM units.

  • Results

    PAPI provides 1.8× and 11.1× performance improvements over state-of-the-art heterogeneous and PIM-only LLM accelerators, respectively.

  • Takeaways & Limitations

    Dynamic scheduling and heterogeneous PIM capabilities let PAPI adapt kernel placement to changing computation and memory-bandwidth demands during LLM decoding.

  • Takeaways & Limitations

    Prior comparisons and designs show that FC and attention kernels can require different computation and memory-bandwidth capabilities, constraining fixed-capability PIM approaches.

Abstract

from arXiv · show

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.

1. Introduction

LLM decoding dominates inference time, while batching and speculative decoding dynamically alter kernel demands between compute-bound and memory-bound behavior. PAPI addresses this variability with online kernel characterization, dynamic scheduling, and heterogeneous hardware tailored to different workloads.

  • 96% of GPT-3 175B inference time is consumed by serial decoding for input and output lengths of 32.The decoding share increases as output length grows.
  • Batching and speculative decoding generate multiple tokens per iteration, changing resource utilization and making decoding kernels compute-bound or memory-bound.Decoding parallelism directly affects memory and computation resource utilization.
  • Dynamic decoding parallelism can shift kernels between compute-bound and memory-bound states, undermining static kernel scheduling.Static mappings may send computation-intensive kernels to PIM units or memory-intensive kernels to computation-centric accelerators.
  • PAPI uses online kernel-property identification to dynamically schedule kernels to suitable units in a GPU-and-PIM heterogeneous architecture.The design targets dynamically changing parallelism properties during LLM inference.
  • PAPI combines runtime bottleneck prediction, GPU/PIM/CPU heterogeneity, and performance-optimized and memory-capacity-optimized PIM units.The two PIM types address memory-intensive kernels with different computational demands and memory footprints.
  • 1.8×, 1.9×, and 11.1× speedups are reported over A100+AttAcc, A100+HBM-PIM, and AttAcc, respectively.The comparisons include heterogeneous GPU-PIM systems and a PIM-only accelerator.

2. Background

LLM inference alternates between prefill and sequential decoding, whose repeated memory movement motivates parallel decoding. Batching and speculative decoding increase token-level parallelism through concurrent requests or draft-token verification.

  • 2.1. LLM Inference: Each transformer decoder contains QKV generation, multi-head attention, projection, and feedforward kernels, comprising FC layers and an attention layer.These kernels use general matrix-vector multiplication computations.
  • 2.1. LLM Inference: Serial decoding generates output tokens one at a time across multiple decoding iterations after prefill processes the input sequence.Each iteration uses the last output token to generate a new token until the end-of-sentence token.
  • 2.1. LLM Inference: Repeatedly loading weight and KV matrices from off-chip memory during serial decoding creates data-movement and performance overheads.KV matrices produced by QKV generation are reused in later attention iterations.
  • 2.2. Optimization Techniques in LLM Inference: Batching processes multiple input sequences concurrently, enabling request-level parallelism during inference.Mixed continuous batching can adjust batch size at runtime as resources and incoming requests change.
  • 2.2. Optimization Techniques in LLM Inference: Speculative decoding combines serial draft decoding with parallel verification of 2-10 predicted tokens by the large model.The draft model predicts tokens that the original LLM verifies simultaneously.

3. Motivation

LLM decoding parallelism varies with batch size, speculation length, and request progress, changing kernel arithmetic intensity and whether kernels are compute- or memory-bound. These dynamics make fixed hardware mappings and uniform PIM designs inefficient.

  • Kernel behavior: FC arithmetic intensity increases with speculation length and becomes compute-bound beyond speculation length 6 at batch size 32, while attention remains memory-bound.Attention changes only slightly because batching primarily improves weight reuse, which does not affect the attention kernel.
  • Dynamic parallelism: Runtime batch size and speculation length vary substantially because of SLO limits, memory capacity, request completion, continuous batching, and adaptive speculation.Static batching reduces runtime request-level parallelism as requests finish, whereas continuous batching changes it as requests are added.
  • Kernel behavior: FC and attention kernels can switch between compute-bound and memory-bound behavior on GPUs as decoding parallelism changes.This variation follows changing arithmetic intensity across runtime parallelization levels.
  • Limits of prior work: Prior heterogeneous accelerators statically map kernels to fixed devices, which can mismatch their changing computation and memory demands.Examples include always assigning FC to GPUs or PIM and attention to PIM or NPUs, despite runtime variation.
  • Limits of prior work: FC and attention can both be memory-bound yet require different PIM capabilities because their arithmetic intensities differ substantially.At batch size 4 and speculation length 8, FC has 31.7 FLOPs/Byte versus attention’s 7.0 FLOPs/Byte, implying about 4.5× higher memory bandwidth needs for attention under equal compute throughput.
  • Design motivation: PAPI addresses these dynamics with a heterogeneous platform combining PIM, GPU, and host CPU resources with parallelism-aware runtime kernel assignment.The design uses distinct computation throughput and memory bandwidth characteristics to match changing demands.

4. PAPI: Overview

PAPI monitors changing decoding parallelism and assigns kernels to suitable hardware at runtime. Its architecture combines a host CPU, high-performance processor, and heterogeneous PIM units for differing compute and memory demands.

  • Dynamic scheduling: PAPI uses a low-overhead kernel bottleneck predictor to assign kernels to suitable computing platforms at runtime.The scheduling framework is designed to monitor dynamic parallelism online while limiting latency and energy overhead.
  • Heterogeneous architecture: PAPI integrates a host CPU, high-performance processor, and physically separate PIM devices to support compute-bound and memory-bound LLM kernels.The high-performance processor handles computation-centric workloads, while PIM devices provide memory-centric execution.
  • Hybrid PIM units: PAPI provides FC-PIM units with relatively high computation capability and Attn-PIM units with larger memory capacity for different kernel demands.The two PIM types are intended to address different parallelization levels of FC and attention kernels.

5. PAPI Dynamic Scheduling

PAPI estimates FC-kernel arithmetic intensity online to identify whether kernels are memory-bound or compute-bound, then schedules them to suitable hardware at runtime. Its scheduler uses both initial and token-level runtime decisions as decoding parallelism changes.

  • Runtime Scheduling: PAPI schedules FC kernels to processing units or FC-PIM units at runtime based on whether they are compute-bound or memory-bound.The mechanism is designed for low latency and low energy consumption.
  • Memory-Boundedness Identification: Arithmetic intensity is estimated from runtime-known parallelism levels RLP and TLP, which approximate AI as RLP × TLP.The FC weight matrix has dimensions (h, h), and h denotes the model hidden dimension.
  • Estimator Accuracy: The estimator closely matches measured arithmetic intensity for GPT-3 66B FC kernels across most tested RLP and TLP configurations.At very large parallelization, such as RLP=128, the estimate is slightly higher, but the scheduling decision remains correct.
  • Initial Scheduling: Initial scheduling uses batch size and speculation length to estimate intensity, sending kernels above threshold α to PUs and others to FC-PIM units.The threshold α is selected offline by comparing execution times on PIM and PU units under varying parallelization levels.
  • Runtime Scheduling: Runtime scheduling monitors changing parallelism and can move FC kernels between PUs and FC-PIM units as estimated arithmetic intensity changes.The token-level scheme tracks finished requests through <|eos|> tokens and monitors TLP updates through a dedicated register.

6. PAPI Architecture

PAPI combines computation-centric processing units with complementary FC-PIM and Attn-PIM designs to address different computation, bandwidth, power, area, and memory-capacity demands. Its PIM configuration uses data reuse to support higher parallelism while separating FC computation needs from attention memory needs.

  • Design Constraints: PAPI evaluates PIM energy, power, and area constraints using an HBM-based PIM simulator and CACTI-3DD area estimates.The design accounts for DRAM access, data transfer, computation, bank area, FPU area, and HBM-die area limits.
  • FC-PIM Design: 33.1% of overall PIM energy comes from DRAM access when one fetched DRAM row is reused for 64 FC computations.Reusing DRAM data reduces access energy and creates additional energy budget for PIM cores.
  • FC-PIM Design: Data reuse ≥4 makes the 4P1B configuration meet the HBM power budget while using more FPUs per DRAM bank.The figure compares power across data reuse levels and FPU counts per bank.
  • FC-PIM Design: PAPI uses 96 banks per HBM memory unit with a 4P1B configuration to satisfy the single-die area constraint.The maximum number of memory banks is smaller than 97 under the stated area assumptions.
  • Attn-PIM Design: Attn-PIM is disaggregated from the high-performance processor to provide scalable memory capacity for growing KV caches and longer sequence lengths.FC-PIM instead emphasizes higher computation capability, while attention kernels require greater memory capacity.
  • Complementary PIM Units: Separately optimizing FC-PIM and Attn-PIM computation and capacity lets PAPI match FC kernels’ higher computation and bandwidth demands with attention kernels’ higher capacity needs.The two device types share bank-level computation and memory hierarchy but differ in processing units per bank.

7. Evaluation

PAPI is evaluated against heterogeneous GPU-PIM and PIM-only baselines across varied decoding configurations and datasets. It improves performance and energy efficiency by dynamically matching kernels to suitable hardware and by using heterogeneous PIM units.

  • End-to-End Performance: PAPI achieves 1.8×, 1.9×, and 11.1× speedups over A100+AttAcc, A100+HBM-PIM, and AttAcc-only, respectively.The comparison covers evaluated decoding parallelization settings with batch sizes of 4, 16, or 64 and speculation lengths of 1, 2, or 4.
  • End-to-End Performance: PAPI achieves 1.7×, 1.7×, and 8.1× speedups over A100+AttAcc, A100+HBM-PIM, and AttAcc-only on Dolly general-qa.The speedups are lower than for creative-writing because creative-writing has longer outputs and larger dynamic changes in parallelization levels.
  • Energy Efficiency: PAPI improves average energy efficiency by 3.4× and 3.1× over A100+AttAcc for creative-writing and general-qa, respectively.PAPI offloads parts of FC kernels to FC-PIM devices, reducing data movement and using low-power processing cores in memory.
  • PIM Architecture Analysis: PIM-only PAPI achieves 2.3× average speedup over AttAcc-only, increasing from 1.6× at lower parallelism to 2.7× at higher parallelism.The compared PIM-only PAPI system includes Attn-PIM and FC-PIM devices and excludes the A100.
  • PIM Architecture Analysis: FC kernels dominate decoding time, while PIM-only PAPI provides 2.9× FC-kernel speedup; communication accounts for 28.2% of decoding time.Attention runs 1.7× slower on Attn-PIM than on AttAcc-only because of the lower-area 1P2B design choice.

8. Related Work

Prior LLM accelerators use heterogeneous GPU-PIM systems to address data movement and varying kernel demands, but PAPI targets dynamically changing demands without sacrificing output quality. It combines runtime adaptation with multiple PIM capabilities.

  • PIM-Enabled LLM Accelerators: PIM integrates processing cores within memory units to provide high bandwidth and mitigate data movement in low-arithmetic-intensity LLM kernels.Prior work applies PIM particularly to memory-bound kernels such as attention.
  • PIM-Enabled LLM Accelerators: Prior heterogeneous PIM systems statically assign kernels to GPUs, NPUs, or PIM devices, whereas PAPI targets dynamically varying computation and memory needs.PAPI is presented as an architecture and runtime framework for changing parallelization levels in real-world LLM workloads.
  • Approximation-Based LLM Acceleration: PAPI does not sacrifice quality in LLM serving while providing performance and energy benefits over state-of-the-art systems.This distinguishes PAPI from approximation approaches such as pruning and quantization, which target scenarios that tolerate approximate results.

9. Conclusion

PAPI addresses dynamically changing parallelization in real-world LLM services with heterogeneous compute units and lightweight runtime scheduling. The evaluation reports 1.8× and 11.1× performance improvements over state-of-the-art LLM inference systems.

  • Conclusion: PAPI supports three computing-unit types and lightweight scheduling that monitors dynamic parallelization and assigns FC and attention kernels to suitable hardware.The design targets varying computation and memory demands caused by batching and speculative decoding.
  • Conclusion: PAPI provides 1.8× and 11.1× performance improvement over state-of-the-art LLM inference systems.The paper positions these results as evidence for heterogeneous PIM-enabled systems that handle dynamic real-world execution scenarios.
Loading 2502.15470v2…