Pimba用存内计算加速后变压器模型,提升长文本生成效率。
Pimba: A Processing-in-Memory Acceleration for Post-Transformer Large Language Model Serving
- 设计共享SPU的存内计算架构,支持变压器与后变压器模型统一加速。
- 相比GPU系统,生成速度最高提升4.1倍,功耗更低。
- 针对状态更新操作优化,采用微软MX量化,兼顾精度与能效。
Transformer是当前大语言模型的核心,但其计算与内存开销随序列长度增长,导致长文本推理难以扩展。为此,研究者探索状态空间模型(SSMs)、线性注意力和循环神经网络(RNNs)等后变压器架构。这一转变带来关键挑战:如何构建统一的推理系统,高效支持变压器与后变压器模型。我们分析了两类模型在批处理推理下的性能特征,发现二者均受限于内存带宽——前者因注意力机制,后者因状态更新。进一步分析表明:(1) 状态更新操作硬件成本高,单银行存内计算效率低;(2) 不同低精度算术方法精度-面积权衡各异,微软的MX为帕累托最优选择。基于此,我们设计了Pimba,一种由状态更新处理单元(SPUs)组成的阵列,每单元共享两个存储银行以实现交错访问。每个SPU包含使用MX量化算术的元素级乘加器,可高效执行状态更新与注意力操作。评估显示,相比优化过的GPU及GPU+PIM系统,Pimba在令牌生成吞吐量上分别提升最高达4.1倍和2.1倍。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。