VEXP: A Low-Cost RISC-V ISA Extension for Accelerated Softmax Computation in Transformers

虽然变压器模型主要由浮点 (FP) 矩阵乘法主导，但通过专用硬件或众核可编程系统的强力加速，性能瓶颈已转移至如 Softmax 这样的非线性函数。加速 Softmax 是一项挑战，因其具有非逐点、非线性的特性，其中指数运算（exponentiation）是最耗资源的步骤。为了解决这一问题，我们设计了一个定制的算术模块，用于 Bfloat16 指数运算，该模块基于一种改进的 Schraudolph 方法近似算法，并通过自定义指令集架构 (ISA) 扩展将其集成到计算集群中 RISC-V 核心的浮点单元 (FPU) 中，面积开销仅为 1%。通过优化软件内核以充分利用这一扩展，与基线集群相比，我们在 Softmax 的执行上实现了延迟减少 162.7 倍，能耗降低 74.3 倍，从而使 GPT-2 配置下的 FlashAttention-2 内核性能提升了 8.2 倍，能效提高了 4.1 倍。此外，所提出的方案能够使多集群系统高效地执行预训练变压器模型（如 GPT-2、GPT-3 和 ViT）的端到端推理，分别实现高达 5.8 倍的延迟减少和 3.6 倍的能耗降低，且无需重新训练模型，同时几乎不损失精度。

论文试图解决Transformer模型中Softmax函数计算效率低的问题，尤其是在Bfloat16数据格式下，由于其非线性和指数运算复杂性，成为性能瓶颈。这是一个现有问题，但随着硬件对矩阵乘法加速的提升，Softmax的优化变得更为关键。

论文提出了一种基于Schraudolph方法的自定义Bfloat16指数运算近似算法，并将其集成到RISC-V核心的FPU中，通过ISA扩展实现。这种方法在几乎不增加硬件面积的情况下显著提升了Softmax计算的速度和能效，同时优化了软件内核以充分利用硬件加速。相比传统方法，该方案在特定任务上实现了显著的性能和能效提升。

实验设计针对FlashAttention-2内核（GPT-2配置）进行了测试，结果表明Softmax计算的延迟和能耗分别减少了162.7倍和74.3倍，整体性能提升了8.2倍，能效提高了4.1倍。此外，在多集群系统中，端到端推理GPT-2、GPT-3和ViT等预训练模型时，实现了5.8倍的延迟降低和3.6倍的能耗减少。论文未提及具体数据集，但明确表示无需重新训练模型且精度损失可忽略。目前尚无开源代码信息，但未来可以进一步研究如何将此方法扩展到其他非线性函数或更广泛的深度学习模型中。

相关研究包括：1) 专注于低精度计算优化的工作，如Google的Bfloat16设计及其在TPU中的应用；2) 针对Softmax加速的研究，例如使用Cordic算法或分段多项式近似的硬件实现；3) Transformer专用硬件加速器设计，如NVIDIA的Hopper架构中引入的FP8支持；4) 其他关于定制指令集或微架构扩展的研究，例如IBM的TrueNorth项目或MIT的Eyeriss加速器。类似论文标题可能为《Efficient Hardware Acceleration for Transformer Models》或《Custom ISA Extensions for Deep Learning Kernels》。