SoftmaxPIM: An HBM-Based PIM Architecture for Accelerating GeMV–Softmax Execution Pipeline

Abstract

LLM inference is dominated by GeMV and softmax computations during attention. These computations impose substantial compute and memory bandwidth demands. Prior works have demonstrated that HBM-based PIM architectures effectively accelerate GeMV. However, softmax remains a critical bottleneck due to frequent data movement. In this work, we make the case for computing softmax operations in memory and integrate them with in-memory GeMV operations to provide a comprehensive in-memory LLM accelerator engine. We propose SoftmaxPIM, an in-HBM compute architecture that tightly integrates GeMV with a lightweight near-bank softmax unit. The proposed softmax unit employs an approximation-based softmax formulation that removes reduction operations, enabling the entire GeMV–softmax pipeline to execute within the HBM memory stack. By eliminating host-memory transfers, the proposed design significantly reduces softmax latency. Our simulation results demonstrate up to 42.7× speedup over a CPU and 2.64× over an RTX3090 GPU, as well as up to 1.45× improvement over the prior in-memory attention accelerator AttAcc.