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

Fuente: arXiv
Saved in:
Bibliographic Details
Main Authors: Kim, Wonung, Lee, Yubin, Kim, Yoonsung, Hwang, Jinwoo, Oh, Seongryong, Jung, Jiyong, Huseynov, Aziz, Park, Woong Gyu, Park, Chang Hyun, Mahajan, Divya, Park, Jongse
Format: Preprint
Published: 2025
Subjects:
Online Access:
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866912588349046784
author Kim, Wonung
Lee, Yubin
Kim, Yoonsung
Hwang, Jinwoo
Oh, Seongryong
Jung, Jiyong
Huseynov, Aziz
Park, Woong Gyu
Park, Chang Hyun
Mahajan, Divya
Park, Jongse
author_facet Kim, Wonung
Lee, Yubin
Kim, Yoonsung
Hwang, Jinwoo
Oh, Seongryong
Jung, Jiyong
Huseynov, Aziz
Park, Woong Gyu
Park, Chang Hyun
Mahajan, Divya
Park, Jongse
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2507_10178
institution arXiv
publishDate 2025
record_format arxiv
spellingShingle Pimba: A Processing-in-Memory Acceleration for Post-Transformer Large Language Model Serving
Kim, Wonung
Lee, Yubin
Kim, Yoonsung
Hwang, Jinwoo
Oh, Seongryong
Jung, Jiyong
Huseynov, Aziz
Park, Woong Gyu
Park, Chang Hyun
Mahajan, Divya
Park, Jongse
Hardware Architecture
Machine Learning
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.
title Pimba: A Processing-in-Memory Acceleration for Post-Transformer Large Language Model Serving
topic Hardware Architecture
Machine Learning
url https://arxiv.org/abs/2507.10178