End-to-End Transformer Acceleration Through Processing-in-Memory Architectures

TL;DR

This paper proposes processing-in-memory architectures to accelerate Transformer models by reducing data movement, managing memory growth, and lowering computational complexity, resulting in improved energy efficiency and latency.

Contribution

It introduces novel processing-in-memory techniques for Transformers, restructuring attention and feed-forward operations, and optimizing memory and complexity management.

Findings

Significant energy efficiency improvements over GPUs and accelerators

Reduced latency in Transformer inference tasks

Effective management of key-value cache growth

Abstract

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.