通过存内计算优化注意力机制,实现Transformer模型高效加速
End-to-End Transformer Acceleration Through Processing-in-Memory Architectures
- 将注意力与前馈计算重构为存内处理,减少芯片外数据传输
- 动态压缩剪枝键值缓存,缓解长序列下的内存膨胀问题
- 将注意力重构成关联记忆操作,降低复杂度与硬件开销
Transformer已成为自然语言处理和大模型的核心,但大规模部署面临三大挑战:一是注意力机制需大量矩阵乘法及频繁的中间结果跨内存-计算单元传输,导致高延迟和高能耗;二是长序列推理中键值缓存(KV cache)可能无序增长,甚至超过模型权重大小,造成严重内存与带宽瓶颈;三是注意力随序列长度呈二次复杂度,放大数据移动与计算开销,使大规模推理效率低下。本文提出存内计算方案,重构注意力与前馈计算以最小化片外数据传输,动态压缩与剪枝KV缓存以应对内存增长,将注意力重新解释为关联记忆操作以降低复杂度与硬件占用。我们还针对现有先进加速器与通用GPU评估该设计,证明在能效与延迟方面均有显著提升。这些方法共同解决了计算开销、内存可扩展性与注意力复杂度问题,进一步实现了Transformer模型的端到端高效加速。
原文摘要 · Abstract (English)
Transformers have become central to natural language processing and large language models, but their deployment at scale faces three major challenges. First, the attention mechanism requires massive matrix multiplications and frequent movement of intermediate results between memory and compute units, leading to high latency and energy costs. Second, in long-context inference, the key-value cache (KV cache) can grow unpredictably and even surpass the model's weight size, creating severe memory and bandwidth bottlenecks. Third, the quadratic complexity of attention with respect to sequence length amplifies both data movement and compute overhead, making large-scale inference inefficient. To address these issues, this work introduces processing-in-memory solutions that restructure attention and feed-forward computation to minimize off-chip data transfers, dynamically compress and prune the KV cache to manage memory growth, and reinterpret attention as an associative memory operation to reduce complexity and hardware footprint. Moreover, we evaluate our processing-in-memory design against state-of-the-art accelerators and general-purpose GPUs, demonstrating significant improvements in energy efficiency and latency. Together, these approaches address computation overhead, memory scalability, and attention complexity, further enabling efficient, end-to-end acceleration of Transformer models.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。