Pimba: A Processing-in-Memory Acceleration for Post-Transformer Large Language Model Serving

TL;DR

Pimba introduces a processing-in-memory system optimized for both transformer and post-transformer LLMs, significantly improving inference throughput by addressing memory bandwidth and hardware cost challenges.

Contribution

The paper presents a novel PIM architecture with shared processing units for efficient state updates and attention, enabling unified serving of diverse LLM architectures.

Findings

Achieves up to 4.1x higher token throughput compared to GPU-based systems.

Effectively supports both transformer and post-transformer LLMs within a unified framework.

Reduces hardware costs by optimizing low-precision arithmetic methods.

Abstract

Transformers are the driving force behind today's Large Language Models (LLMs), serving as the foundation for their performance and versatility. Yet, their compute and memory costs grow with sequence length, posing scalability challenges for long-context inferencing. In response, the algorithm community is exploring alternative architectures, such as state space models (SSMs), linear attention, and recurrent neural networks (RNNs), which we refer to as post-transformers. This shift presents a key challenge: building a serving system that efficiently supports both transformer and post-transformer LLMs within a unified framework. To address this challenge, we analyze the performance characteristics of transformer and post-transformer LLMs. Despite their algorithmic differences, both are fundamentally limited by memory bandwidth under batched inference due to attention in transformers and state updates in post-transformers. Further analyses suggest two additional insights: (1) state update operations, unlike attention, incur high hardware cost, making per-bank PIM acceleration inefficient, and (2) different low-precision arithmetic methods offer varying accuracy-area tradeoffs, while we identify Microsoft's MX as the Pareto-optimal choice. Building on these insights, we design Pimba as an array of State-update Processing Units (SPUs), each shared between two banks to enable interleaved access to PIM. Each SPU includes a State-update Processing Engine (SPE) that comprises element-wise multipliers and adders using MX-based quantized arithmetic, enabling efficient execution of state update and attention operations. Our evaluation shows that, compared to LLM-optimized GPU and GPU+PIM systems, Pimba achieves up to 4.1x and 2.1x higher token generation throughput, respectively.