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

TL;DR

This work tackles end-to-end Transformer acceleration on processing-in-memory (PIM) hardware by addressing three core bottlenecks: the compute-write-compute (CWC) dependency in attention, the KV cache memory bottleneck, and the quadratic attention complexity. It introduces (i) a matrix-decomposition approach that removes CWC and a sub-matrix pipeline to keep ReRAM crossbars busy, (ii) a software-hardware co-design Cascade KV Cache Pruning-Quantization (CPQ) with Hierarchical Quantization Extension (HQE) to compress and transfer only non-zero KV data, and (iii) a reformulation of attention as nearest-neighbor retrieval using content-addressable memory (CAM) to reduce similarity computations. The proposed architecture combines digital CIM blocks, a computing engine, and dedicated pruning/quantization units to support end-to-end LLM inference. Empirical results show substantial gains in energy efficiency and latency compared to Nvidia A100 and FlightLLM baselines, with reported improvements such as $159.9\times$ energy efficiency and $49.6\times$ throughput versus A100, and $34.8\times$ energy efficiency and $29.2\times$ throughput versus FlightLLM for end-to-end OPT-6.7B inference.

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.

Figures (6)

• Figure 1: Comparison between FC/CONV layers and MatMul layer in ReRAM implementation yang2020retransformer: (a) MatMul layers, (b) FC/CONV layer.
• Figure 2: Illustration of eliminating the CWC dependency in scaled dot-product attention (SDA) yang2020retransformer.
• Figure 3: (a) Layer-level pipeline and (b) Sub-matrix level pipeline yang2020retransformer.
• Figure 4: Overview of the CPQ compression strategy chen2025titanus.
• Figure 5: Hierarchical quantization extension strategy chen2025titanus.