SoftmAP: Software-Hardware Co-design for Integer-Only Softmax on Associative Processors

📄 arXiv: 2411.17847v1 📥 PDF

作者: Mariam Rakka, Jinhao Li, Guohao Dai, Ahmed Eltawil, Mohammed E. Fouda, Fadi Kurdahi

分类: cs.AR, cs.AI

发布日期: 2024-11-26

备注: Accepted in DATE 2025

💡 一句话要点

SoftmAP：面向关联处理器的全整数Softmax软硬件协同设计

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: Softmax优化 内存计算 软硬件协同设计 低精度量化 大型语言模型 边缘计算 整数计算

📋 核心要点

  1. 大型语言模型中的Softmax等非线性算子对量化敏感，成为资源受限设备上的性能瓶颈。
  2. SoftmAP通过软硬件协同设计，利用内存计算硬件实现全整数低精度Softmax。
  3. 实验表明，SoftmAP在能量延迟积方面比A100和RTX3090 GPU提升高达三个数量级。

📝 摘要（中文）

为了使大型语言模型(LLMs)能够在资源受限的设备上运行，近期的研究致力于降低其计算和内存开销。尽管压缩技术取得了进展，但诸如Softmax和Layernorm之类的非线性算子仍然是瓶颈，因为它们对量化非常敏感。我们提出了SoftmAP，一种软硬件协同设计方法，它使用内存计算(IMC)硬件实现全整数低精度Softmax。与A100和RTX3090 GPU相比，我们的方法在能量延迟积方面实现了高达三个数量级的改进，从而在不影响性能的情况下，使LLM更易于部署。

🔬 方法详解

问题定义：论文旨在解决大型语言模型在资源受限设备上部署时，Softmax算子由于对量化敏感而导致的性能瓶颈问题。现有方法在低精度量化下，Softmax的精度损失较大，计算和内存开销仍然很高，难以满足边缘设备的需求。

核心思路：论文的核心思路是采用软硬件协同设计，利用内存计算(IMC)硬件的并行性和低功耗特性，实现全整数低精度Softmax。通过定制化的硬件加速器和优化的软件算法，在保证精度的前提下，大幅降低计算复杂度和功耗。

技术框架：SoftmAP的整体框架包括：1) 软件层面的整数Softmax算法优化，包括量化策略和数值逼近；2) 硬件层面的内存计算加速器设计，用于高效执行整数Softmax计算；3) 软硬件协同优化，根据硬件特性调整软件算法，实现最佳性能。主要模块包括：输入量化模块、指数近似计算模块、归一化模块和输出反量化模块。

关键创新：论文的关键创新在于提出了一个全整数低精度Softmax的软硬件协同设计方案。与传统的基于GPU的Softmax实现相比，SoftmAP利用内存计算的并行性和低功耗特性，显著降低了计算复杂度和功耗。此外，论文还针对内存计算硬件的特点，优化了整数Softmax算法，提高了计算效率。

关键设计：论文的关键设计包括：1) 低精度整数量化方案，在保证精度的前提下，降低计算和存储开销；2) 基于查找表(LUT)的指数近似计算方法，避免了复杂的浮点运算；3) 内存计算加速器的架构设计，包括计算单元的布局和数据流的优化；4) 软硬件协同优化策略，根据硬件特性调整量化参数和算法流程。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

SoftmAP在能量延迟积方面实现了显著的性能提升。实验结果表明，与A100和RTX3090 GPU相比，SoftmAP的能量延迟积降低了高达三个数量级。这表明SoftmAP在降低功耗和提高计算效率方面具有显著优势，使其成为在资源受限设备上部署大型语言模型的理想选择。

🎯 应用场景

SoftmAP技术可广泛应用于边缘设备上部署的大型语言模型，例如智能手机、物联网设备和嵌入式系统。通过降低Softmax算子的计算和内存开销，SoftmAP能够使这些设备在本地运行复杂的AI模型，从而提高响应速度、保护用户隐私并降低网络带宽需求。未来，SoftmAP有望推动AI技术在资源受限环境下的普及。

📄 摘要（原文）

Recent research efforts focus on reducing the computational and memory overheads of Large Language Models (LLMs) to make them feasible on resource-constrained devices. Despite advancements in compression techniques, non-linear operators like Softmax and Layernorm remain bottlenecks due to their sensitivity to quantization. We propose SoftmAP, a software-hardware co-design methodology that implements an integer-only low-precision Softmax using In-Memory Compute (IMC) hardware. Our method achieves up to three orders of magnitude improvement in the energy-delay product compared to A100 and RTX3090 GPUs, making LLMs more deployable without compromising performance.