NeuPIMs: NPU-PIM Heterogeneous Acceleration for Batched LLM Inferencing

TL;DR

The paper addresses the bottleneck in batched LLM inference arising from the coexistence of GEMM-dominated QKV/FFN computations and bandwidth-bound GEMV operations in MHA. It introduces NeuPIMs, a hardware-software co-design that couples a GEMM-centric NPU with GEMV-optimized PIM accelerators and implements dual row buffers plus a sub-batch interleaving scheduler to enable concurrent NPU+PIM execution. Key contributions include a dual-row-buffer PIM microarchitecture, an extended PIM command interface with GEMV primitives, an MHA-aware latency estimator, a greedy load-balancing scheduler, and scaling strategies for pipeline and tensor parallelism. Through cycle-accurate evaluation on GPT-3 variants with Alpaca and ShareGPT workloads, NeuPIMs delivers up to 3x throughput gains over GPU-only and 1.6x over naïve NPU+PIM baselines, with significantly improved resource utilization, demonstrating practical potential for PIM-enabled LLM inference.

Abstract

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.

Figures (15)

• Figure 1: (a) Mathematical components of decoder blocks that constitute LLMs, (b) NPU-only baseline accelerator equipped with non-PIM memory (e.g., GPU), (c) NPU+PIM integrated baseline accelerator, and (d) the proposed NeuPIMs accelerator.
• Figure 2: Model architecture and inference in LLMs.
• Figure 3: Operators in a LLM decoder block.
• Figure 4: Arithmetic intensities of LLM layers.
• Figure 5: GPU resource utilization for four different LLMs.