紧急发布：利用内存计算加速变压器推理的硬件-软件协同设计

《IEEE Transactions on Circuits and Systems for Artificial Intelligence》：HASTILY: Hardware-Software Co-Design for Accelerating Transformer Inference Leveraging Compute-in-Memory

本文提出基于计算内存（CIM）架构的HASTILY加速器，通过统一计算与查找模块（UCLM）实现softmax加速和内存优化，结合细粒度流水线和并行计算策略，将序列长度相关度从二次降低到线性，在65nm工艺下对BERT模型实现4.4-9.8倍吞吐量提升，并显著提高能效比。

　　

摘要：

变换器已成为大多数机器学习应用中神经网络架构的支柱。由于它们的广泛应用，人们提出了多种加速注意力机制（注意力机制是变换器的基本构建模块）的方法。本文通过软硬件协同设计的方法来解决加速注意力机制所面临的挑战，同时利用了内存计算（Compute-in-Memory, CIM）架构。我们开发的名为HASTILY的加速器在能耗和面积方面表现出色，旨在加速softmax运算（注意力机制中的关键操作），并大幅降低随着输入序列长度增加而呈二次方增长的芯片内存需求。该加速器采用了名为“统一计算与查找模块”（Unified Compute and Lookup Modules, UCLMs）的新型CIM单元，这些单元在同一SRAM阵列中集成了查找和乘累加功能，从而比标准CIM阵列占用更少的面积。UCLMs采用TSMC 65nm工艺制造，能够同时执行指数运算和矩阵向量乘法操作。为了进一步提高性能，HASTILY还采用了细粒度的流水线策略来调度注意力层和前馈层，将序列长度的二次方依赖性降低为线性依赖性。此外，对于涉及计算指数值最大值和总和的快速softmax运算，我们通过“归约-收集”（reduce-and-gather）策略在多个核心上实现并行处理。我们使用专为注意力计算定制的编译器和标准的CIM模拟器对所提出的架构进行了评估。结果表明，在INT-8精度下，与Nvidia A40 GPU和基准CIM硬件相比，BERT模型的端到端吞吐量（TOPS）分别提升了4.4倍、9.8倍和1.7倍。同时，在能效方面（TOPS/W），HASTILY比A40 GPU提高了16倍至36倍，并且与基准CIM硬件相当。我们的评估代码库将在github上开源。

显示更多