Pimba: A Processing-in-Memory Acceleration for Post-Transformer Large Language Model Serving
Problem Statement
As LLMs scale to longer contexts, transformer attention costs grow prohibitively, driving the field toward post-transformer architectures (Mamba-2, linear attention, RNNs) with different computational profiles. Current serving systems and accelerators are largely specialized for one architecture family, leaving a gap for a unified, efficient hardware solution that handles both transformer attention and post-transformer state updates without duplicating design effort.
Key Novelty
- Empirical finding that both transformer attention and post-transformer state updates are memory-bandwidth bound under batched inference, unifying two seemingly disparate architecture families under one hardware bottleneck
- Discovery that state-update operations have substantially higher hardware cost than attention, making naive per-bank PIM acceleration inefficient—motivating a shared, cross-bank compute design
- Systematic accuracy-vs-area analysis of low-precision arithmetic schemes for PIM, identifying Microsoft's MX format as the Pareto-optimal choice for hardware-efficient quantized computation
- Novel Pimba architecture combining PIM and quantization: State-update Processing Units (SPUs) shared between bank pairs with interleaved access, enabling one engine to execute both attention and state-update operations efficiently
Evaluation Highlights
- Up to 4.1x higher generation throughput compared to LLM-optimized GPU systems
- Up to 2.1x higher generation throughput compared to GPU+PIM systems
- Hardware efficiency gains from sharing SPUs across bank pairs, reducing area overhead relative to naive per-bank PIM designs
Signal Assessment
Methodology
- Characterize and compare performance bottlenecks of transformer and post-transformer LLMs under batched inference to identify a shared memory-bandwidth limitation
- Analyze hardware cost differences between state-update and attention operations to determine optimal PIM compute granularity
- Evaluate accuracy-area tradeoffs across low-precision arithmetic formats to select MX as the quantization scheme for in-memory compute
- Design the Pimba architecture with SPUs shared across bank pairs and SPEs implementing MX-quantized element-wise multiply-add units
- Benchmark Pimba's generation throughput against LLM-optimized GPU and GPU+PIM baselines
System Components
A compute unit shared between two DRAM banks to enable interleaved access, amortizing the high hardware cost of state-update operations across banks
The core engine within each SPU containing element-wise multipliers and adders that use MX-based quantized arithmetic to execute both state updates and attention
Microsoft's MX (microscaling) low-precision format, selected as the Pareto-optimal choice balancing accuracy and hardware area cost for quantized arithmetic
Processing-in-Memory paradigm applied at the bank level to reduce data movement for memory-bandwidth-bound attention and state-update computations
Results
| Metric/Benchmark | Baseline | Pimba | Delta |
|---|---|---|---|
| Generation throughput vs LLM-optimized GPU | 1x | up to 4.1x | +310% |
| Generation throughput vs GPU+PIM system | 1x | up to 2.1x | +110% |
Key Takeaways
- Under batched inference, memory bandwidth—not compute—is the dominant bottleneck for both transformer attention and post-transformer state updates, motivating bandwidth-centric accelerator design
- A unified accelerator can serve both transformer and post-transformer LLMs efficiently, reducing the need for architecture-specific hardware as model families diversify
- Naive per-bank PIM designs are inefficient for state-heavy operations; sharing compute engines across bank pairs is a practical way to balance hardware cost and performance
- MX (microscaling) quantization offers a favorable accuracy-area tradeoff for in-memory low-precision arithmetic, making it a strong candidate for future PIM and quantized-inference hardware designs
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) (e.g., Mamba-2), 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 not only emerging post-transformer LLMs but also existing transformer models 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 largely bounded by memory bandwidth under batched inference—due to attention in transformers and state updates in post-transformers. Inspired by this finding, we propose Pimba, an accelerator solution that aims to address the memory bottleneck by jointly leveraging (1) Processing-in-Memory (PIM) paradigm and (2) LLM quantization. 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 a Pareto-optimal choice. Building on these insights, we design the architecture of Pimba as an array of State-update Processing Units (SPUs), each shared between two banks to enable interleaved access. 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.1 × and 2.1 × higher generation throughput, respectively.