NeuPIMs: NPU-PIM Heterogeneous Acceleration for Batched LLM Inferencing

Guseul Heo; Sangyeop Lee; Jaehong Cho; Hyunmin Choi; Sanghyeon Lee; Hyungkyu Ham; Gwangsun Kim; Divya Mahajan; Jongse Park

Source-linked AI summary

NeuPIMs: NPU-PIM Heterogeneous Acceleration for Batched LLM Inferencing

Guseul Heo, Sangyeop Lee, Jaehong Cho, Hyunmin Choi, Sanghyeon Lee, Hyungkyu Ham, Gwangsun Kim, Divya Mahajan, Jongse Park

arXiv:2403.00579v3Mar 1, 2024cs.AR

TL;DR

LLM inference combines compute-intensive GEMMs, which NPUs handle well, with bandwidth-heavy GEMVs, for which NPUs are less efficient and PIM is better suited. NeuPIMs integrates a GEMM-centric NPU with GEMV-friendly PIM using dual row buffers and sub-batch interleaving, achieving 1.6× higher throughput than a naïve NPU+PIM baseline.

Problem
LLM inference imposes substantial memory and compute requirements, while NPUs are under-utilized on GEMVs and PIM is less effective for GEMMs.
Method
NeuPIMs combines a conventional GEMM-centric NPU with PIM accelerators and uses dual row buffers plus sub-batch interleaving to enable concurrent execution.
Results
1.6× throughput improvement over a naïve NPU+PIM integration, with gains ranging from 13% to 3× across models and datasets.
Takeaways & Limitations
NeuPIMs takes an effective initial step toward practical PIM deployment for LLM inference by improving utilization of NPU and PIM resources.
Takeaways & Limitations
Pipeline parallelism reduces the decoder blocks and batch size per device, while sub-batch interleaving can under-utilize NPU systolic arrays.

Abstract

from arXiv · show

Modern transformer-based Large Language Models (LLMs) are constructed with a series of decoder blocks. Each block comprises three key components: (1) QKV generation, (2) multi-head attention, and (3) feed-forward networks. In batched processing, QKV generation and feed-forward networks involve compute-intensive matrix-matrix multiplications (GEMM), while multi-head attention requires bandwidth-heavy matrix-vector multiplications (GEMV). Machine learning accelerators like TPUs or NPUs are proficient in handling GEMM but are less efficient for GEMV computations. Conversely, Processing-in-Memory (PIM) technology is tailored for efficient GEMV computation, while it lacks the computational power to handle GEMM effectively. Inspired by this insight, we propose NeuPIMs, a heterogeneous acceleration system that jointly exploits a conventional GEMM-focused NPU and GEMV-optimized PIM devices. The main challenge in efficiently integrating NPU and PIM lies in enabling concurrent operations on both platforms, each addressing a specific kernel type. First, existing PIMs typically operate in a "blocked" mode, allowing only either NPU or PIM to be active at any given time. Second, the inherent dependencies between GEMM and GEMV in LLMs restrict their parallel processing. To tackle these challenges, NeuPIMs is equipped with dual row buffers in each bank, facilitating the simultaneous management of memory read/write operations and PIM commands. Further, NeuPIMs employs a runtime sub-batch interleaving technique to maximize concurrent execution, leveraging batch parallelism to allow two independent sub-batches to be pipelined within a single NeuPIMs device. Our evaluation demonstrates that compared to GPU-only, NPU-only, and a naïve NPU+PIM integrated acceleration approaches, NeuPIMs achieves 3$\times$, 2.4$\times$ and 1.6$\times$ throughput improvement, respectively.

1 Introduction

LLM inference combines compute-intensive GEMMs with bandwidth-intensive GEMVs, exposing complementary strengths and utilization gaps between NPUs and PIM. NeuPIMs co-designs hardware and scheduling to execute both efficiently, improving throughput over conventional baselines.

Workload characteristics: LLM decoder blocks contain QKV generation, multi-head attention, and feed-forward networks, with batching producing GEMMs for QKV and feed-forward computation.Multi-head attention instead requires GEMVs because activation vectors provide no data-reuse opportunity.
Workload characteristics: NPUs efficiently execute GEMMs but underutilize resources on lower-arithmetic-intensity GEMVs, whereas PIM is promising for bandwidth-intensive GEMVs but less effective for GEMMs.This complementary specialization motivates a heterogeneous accelerator.
NeuPIMs design: NeuPIMs jointly exploits a GEMM-centric NPU and multiple GEMV-friendly PIM accelerators for batched LLM inference.The system is designed to balance memory bandwidth and computational-resource utilization.
NeuPIMs design: Dual row buffers enable regular memory accesses to occur concurrently with PIM GEMV operations, addressing blocked execution in existing PIMs.NeuPIMs also intersperses memory and PIM commands while preserving DRAM timing constraints.
NeuPIMs design: Sub-batch interleaving parallelizes GEMM operations from one independent sub-batch with GEMV operations from another.The scheduler partitions a batch into two sub-batches to create concurrent NPU-PIM execution.

2 Background

LLM inference uses decoder blocks containing QKV generation, multi-head attention, and feed-forward networks, whose computational characteristics differ under batching. Serving systems selectively batch compatible layers, creating a need to handle GEMM and GEMV operations together.

Each decoder block contains QKV generation, multi-head attention, and feed-forward networks.
Batching transforms QKV and FFN computations into GEMMs, while MHA remains dominated by activation-activation operations with different characteristics.
MHA layers challenge batching, whereas QKV generation and FFN layers can be batched to reuse weights across requests.
Selective batching computes attention individually while batching QKV generation and FFN layers, requiring simultaneous GEMM and GEMV computation.
NeuPIMs builds on iteration-level scheduling, allowing requests to join or leave batches at iteration boundaries.

3 Motivation

LLM inference workloads leave conventional GPU and naïve NPU-PIM systems under-utilized because bandwidth-bound and compute-bound operations use resources unevenly. These observations motivate a heterogeneous design that can execute NPU and PIM work concurrently.

GPU-based LLM serving can approach full capacity utilization while still experiencing layer-wise resource imbalance across compute, memory, and bandwidth.
Naïve NPU-PIM integration replaces GPUs with a systolic-array NPU and Newton PIM accelerator to address bandwidth-bound computations.
The proposed system targets low utilization with an NPU-PIM architecture and a scheduler that interleaves two sub-batches.
When PIM runs MHA layers, NPU utilization falls to zero and combined NPU-PIM utilization remains below 40%.
Existing PIM microarchitectures serialize host-NPU and PIM execution, preventing concurrent use of otherwise disjoint resources.

4 Overview of NeuPIMs

NeuPIMs combines NPU and PIM devices with scheduling and memory-system support for batched LLM generation. Its design separates GEMM-oriented computation from PIM-oriented work while enabling interleaved execution.

A scheduler partitions each inference batch into two sub-batches and interleaves their executions to improve NPU-PIM parallelization.
NeuPIMs assigns summarization-phase GEMMs to standalone NPUs and focuses NeuPIMs devices on generation-phase computation.
Each NeuPIMs bank has dual row buffers, separating PIM execution from regular NPU memory accesses.
Memory controllers interleave memory and PIM commands while preserving inter-command timing constraints across 32 HBM-based PIM channels.
The compiler translates LLM specifications into intermediate representations and produces NPU and NeuPIMs instruction binaries.

5 NeuPIMs Architecture

NeuPIMs extends PIM hardware and command scheduling to support concurrent memory access and GEMV execution. Its architecture coordinates variable-dimensionality PIM operations, command interleaving, and execution across channels and sub-batches.

Memory architecture: Single-row-buffer PIM accelerators load matrix rows into bank-local buffers and perform parallel partial dot products, but block simultaneous NPU execution.
Design scope: NeuPIMs is prototyped with Newton but is intended to apply to GEMV accelerators using standard DRAM microarchitecture and command interfaces.
Command interface: The command interface includes PIM_GWRITE, PIM_ACTIVATION, PIM_DOTPRODUCT, and PIM_RDRESULT for moving operands, activating rows, computing dot products, and returning results.
Memory architecture: NeuPIMs uses dual row buffers so regular memory accesses and PIM operations can proceed on separate bank rows.
Command interface: PIM_HEADER communicates variable GEMV dimensionality so memory controllers can estimate execution latency and avoid conflicts with DRAM refresh.
Scheduling: Requests are distributed across PIM channels, whose controllers broadcast commands to banks and interleave memory read/write commands with PIM commands.
Scheduling: Execution timelines contrast serialized decoder-block processing with sub-batch interleaving across NPU systolic arrays and vector units.

6 NeuPIMs Scheduling

NeuPIMs scheduling overlaps NPU and PIM execution at both the MHA-operation and sub-batch levels. It combines dual-row-buffer support with latency-aware channel balancing and sub-batch partitioning.

Overlapping Opportunities in MHA Layer: Dual row buffers let NeuPIMs handle NPU memory accesses and PIM commands simultaneously during MHA execution.This enables overlap between PIM-side logit and attend operations and NPU-side softmax operations.
Overlapping Opportunities in MHA Layer: MHA overlap is limited to head-level parallelism and leaves the NPU systolic arrays largely unused.The overlap applies between PIM and NPU vector units rather than the systolic arrays.
Sub-batch Interleaving: Without interleaving, decoder-block operators execute sequentially, producing a total time of N times the per-decoder-block execution time.The serialized sequence includes QKV generation, MHA, projection, and FFNs.
Sub-batch Interleaving: Sub-batch interleaving hides MHA execution within NPU execution, improving NPU and PIM utilization.During interleaving, execution is mostly bounded by NPU GEMM time, hiding PIM MHA time.
Channel Load Balancing: NeuPIMs balances channel loads by estimating MHA latency, sorting requests by decreasing token length, and assigning each request to the minimally loaded channel.The greedy min-load bin-packing algorithm updates estimated latency after each assignment.
Sub-batch Partitioning: Requests in each channel are divided into halves and appended to two sub-batches for interleaving.The partitioning maintains similar sub-batch sizes because NPU latency depends strongly on inference batch size.

7 Scaling NeuPIMs System

NeuPIMs can use tensor and pipeline parallelism to scale across devices, but pipeline parallelism reduces per-device decoder blocks and batch size. The system therefore prioritizes tensor parallelism and uses pipeline parallelism when model size requires it.

Pipeline Parallelism: Pipeline parallelism divides decoder layers across NeuPIMs devices and processes micro-batches in a pipeline.The batch is divided according to pipeline depth.
Pipeline Parallelism: Pipeline parallelism reduces the decoder blocks per device and the batch size, which can lower performance benefits and under-utilize NPU systolic arrays.Sub-batch interleaving further reduces batch size.
Tensor Parallelism: Tensor parallelism splits model tensors into shards executed across devices, with results aggregated after each step.The aggregation requires inter-device communication.
Tensor Parallelism: Sub-batch interleaving doubles communication frequency but leaves total communication traffic unchanged relative to a non-partitioned batch.One completed sub-batch can communicate while the other computes, reducing communication latency.
Parallelism Strategy: NeuPIMs prioritizes tensor parallelism and uses pipeline parallelism when model size is too large for tensor parallelism alone.This choice follows the differing communication and utilization implications of the two strategies.

8 Evaluation

NeuPIMs is evaluated against GPU-only, NPU-only, and NPU+PIM baselines across GPT-3 variants, datasets, batch sizes, and parallelization schemes. It improves throughput through concurrent NPU/PIM execution, with gains increasing for larger batches and longer sequences, while incurring power overhead.

Evaluation methodology: NeuPIMs is compared with GPU-only, NPU-only, and NPU+PIM systems using GPT-3 variants, ShareGPT and Alpaca datasets, and synthesized inference-serving workloads.The evaluation uses cycle-accurate simulation and varies model type, batch size, sequence lengths, and tensor/pipeline parallelism.
Throughput: 1.5× throughput improvement is achieved by NPU+PIM over NPU-only, while NeuPIMs adds 13% to 3× improvement over NPU+PIM across models and datasets.NeuPIMs gains are larger for ShareGPT and grow as batch size increases from 64 to 512.
Utilization: 64.9% NPU utilization and 26.4% PIM utilization demonstrate concurrent NPU+PIM execution in NeuPIMs.The reported utilization trends support the role of concurrent execution in improving resource use.
Batch-size sensitivity: NeuPIMs achieves the highest throughput at batch sizes of 256 or larger, indicating that larger-batch inference is preferable for this system.The system shifts the bottleneck toward NPU computation as batch size grows.
Power and energy: NeuPIMs has 1.8× higher power consumption than NPU-only while delivering 2.4× speedup and 25% energy reduction.The comparison uses an NPU-only system equipped with vanilla HBM.
Comparison with TransPIM: NeuPIMs achieves 228× higher average throughput than TransPIM, with speedups ranging from 79× to 431×.The paper attributes the gap to executing GEMMs on the NPU and TransPIM’s focus on single-batch transformer inference.

9 Discussion

NeuPIMs is efficient for batched LLM inference but is less suitable for training because training consists entirely of GEMMs. Integrating it with production software requires a translator from existing model representations to the NeuPIMs specification.

Model training: NeuPIMs has limited training efficiency because fixed-length training sequences entail GEMMs, whereas PIM targets bandwidth-bound GEMV operations.The system can be used for training, but its efficiency is limited.
Production software integration: Production integration requires a translator that converts ONNX-, PyTorch-, and JAX-defined models into the NeuPIMs LLM specification.The rest of the system stack already includes scheduling, compilation, and inference runtime components.

10 Related Work

Related work improves LLM serving through memory, kernel, partitioning, and execution optimizations, while PIM research targets GEMV acceleration. NeuPIMs focuses on simultaneous NPU and PIM execution for LLM inference.

LLM inference serving: Prior LLM serving systems optimize memory footprint, kernel execution, operator partitioning, or combinations of these strategies.NeuPIMs targets hardware utilization across compute- and I/O-suitable NPU and PIM platforms.
PIM for language model support: Existing PIM studies accelerate GEMV operations, but the cited works do not enable simultaneous execution of PIM and NPU operations for LLM inference.NeuPIMs adds a scheduling policy and heterogeneous execution approach to address this gap.

11 Conclusion

NeuPIMs integrates an NPU with PIM to address the differing compute, memory, and bandwidth demands of transformer layers. Its scheduling and execution strategy improves batched LLM inference throughput over naïve NPU+PIM integration.

Conclusion: NeuPIMs integrates a general ML accelerator with PIM to address high memory capacity, compute intensity, and bandwidth constraints in LLM inference.The system combines NPU, HBM memory, and PIM resources for inference serving.
Conclusion: 1.6× throughput improvement is achieved over a baseline that naïvely integrates an NPU with a PIM accelerator.The result is attributed to NeuPIMs’ scheduling and execution strategy.