Return
Algorithm-Hardware Implications of Softmax Approximations for In-Memory Computing-Based LLM Accelerators
DOI:10.1109/JETCAS.2026.3661247.png)
Abstract
En 中文
As large language models (LLMs) continue to scale in both size and particularly in sequence length, the softmax-attention mechanism of the underlying transformer architecture has emerged as a critical bottleneck. While in-memory computing (IMC) based architectures offer efficient acceleration for vector-matrix multiplications (VMMs), integrating softmax remains challenging due to its computational complexity. This work presents a comprehensive analysis of softmax approximations and softmax precision requirements in modern LLMs spanning model sizes from 125M to 405B parameters across several natural language processing (NLP) benchmarks. Based on these findings, we present an IMC-optimized reformulation of softmax for attention that results in negligible accuracy degradation, Including: 1) transformation of complexity from sequence-length-dependency to the much smaller attention head dimension; and 2) merging of softmax components within the adjacent IMC’s VMM operations. We present both digital and analog implementations. Our digital approach reduces area overhead to <inline-formula xmlns:mml="http://www.w3.org/1998/Math/MathML" xmlns:xlink="http://www.w3.org/1999/xlink"> <tex-math notation="LaTeX">$\mathrm {81.2~\mu \text {m}^{2}}$ </tex-math></inline-formula> and energy to 0.69 pJ per input-output pair, resulting in more than <inline-formula xmlns:mml="http://www.w3.org/1998/Math/MathML" xmlns:xlink="http://www.w3.org/1999/xlink"> <tex-math notation="LaTeX">$20 \times $ </tex-math></inline-formula> lower energy-delay-area product compared to state-of-the-art. Our analog implementation achieves <inline-formula xmlns:mml="http://www.w3.org/1998/Math/MathML" xmlns:xlink="http://www.w3.org/1999/xlink"> <tex-math notation="LaTeX">$\mathrm {37.1~\mu \text {m}^{2}}$ </tex-math></inline-formula> and 0.33 pJ by leveraging a spike-based encoding. These results highlight the importance of rethinking softmax computation for IMC to enable scalable transformer accelerators.
Keywords:
Softmax
attention
transformers
LLM
in-memory computing
Journal
IF:
3.8
Papers:
1.4K
Citations:
2.8K

