通过内存内计算减少Transformer的通信开销,显著提升推理速度。
RED-PIM: Reducing Data Movement for Transformers using Processing-in-Memory

- 将注意力计算重构为本地化操作,降低跨内存库数据移动
- 将中间矩阵从N×N压缩至d×d,延迟降低最高达99.99%
- 适合长序列处理,实测长文档性能提升99.60%
Transformer模型广泛应用于自然语言处理、计算机视觉、网页搜索和DNA序列分析等领域。其高效性受制于注意力运算中处理单元与内存间频繁的数据搬运。处理-内存(PIM)技术通过在内存内部执行计算缓解此问题。然而,现有PIM方案存在高昂的跨库通信开销,且受限于内存库容量难以扩展,导致注意力数据需分片存储,削弱了PIM优势。本文提出RED-PIM,一种算法-架构协同设计,将注意力延迟中的跨库数据移动从O(N²)降至O(N),中间注意力矩阵由N×N缩小至d×d。通过重排矩阵操作、本地计算及优化数据传输策略,显著降低计算成本与互连流量。相比基线PIM实现,RED-PIM推理时间减少16.05%至99.99%(几何均值66.42%),长序列收益最大。在真实数据集上,长文档性能提升99.60%,短文本提升13.44%,同时保持或提升准确率。结果表明,RED-PIM在可扩展性和效率方面具有显著优势。
原文摘要 · Abstract (English)
Transformers are widely used across many domains, including natural language processing, computer vision, web search, and DNA sequence analysis. Given their broad applicability, improving the performance of transformer models is critical. However, the high volume of data movement between processing units and memory during attention operations significantly limits their efficiency. Processing-In-Memory (PIM) mitigates this issue by performing computations directly inside memory. While prior work has proposed PIM-based transformer implementations, they suffer from costly inter-bank communication, and struggle to scale due to the limited capacity of memory banks. As a result, attention-related data must be split across banks, diminishing the potential benefits of PIM. In this work, we propose RED-PIM, an algorithm-architecture co-design that reduces attention latency by minimizing inter-bank data movement from O(N^2) to O(N) and shrinking intermediate attention matrices from N x N to d x d. By reorganizing matrix operations, performing computations locally, and employing an optimized data transfer strategy, RED-PIM significantly reduces computation cost and interconnect traffic. Compared to baseline PIM implementation, RED-PIM achieves inference time reductions ranging from 16.05% to 99.99% (geometric mean of 66.42%), with the largest gains on longer sequences. On real-world datasets, RED-PIM improves performance by 99.60% for long documents and 13.44% for shorter ones, while maintaining or improving accuracy. These results demonstrate RED-PIM's effectiveness for scalable and efficient transformer inference.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。