COMET-3D: Compute-in-Memory Based Transformer Accelerator with Optimized Pipeline and 3D Heterogeneous Integration

Transformers have become the backbone of large language decoder and encoder models, but their compute- and memory-intensive nature makes them inefficient on traditional von Neumann architectures. Compute-in-memory (CIM) architectures offer a promising path forward by enabling in situ matrix operations and reducing memory access overhead. However, existing CIM-based accelerators suffer from: (1) unrealistic assumption of single-cycle activation of entire crossbar; (2) inefficient pipelines that are not optimized for operational unit (OU)-based execution; and (3) high analog-to-digital converter (ADC) cost. To address these limitations, we propose a latency-optimized pipeline tailored specifically for OU-based CIM execution and introduce COMET-3D—a 3D heterogeneous architecture that integrates SRAM and ReRAM-based CIM arrays with a logic die containing digital MAC units and softmax modules. The proposed architecture and dataflow maximize hardware utilization and enable efficient acceleration of the multi-head self-attention (MHSA) layer. Experimental results across BERT, GPT2, and LLAMA demonstrate that COMET-3D outperforms baseline architectures with similar compute resources by up to 34× in energy-delay product (EDP) for LLAMA, with gains of 7.4× for GPT2 and 4.3× for BERT-Large.