arXiv:2504.11227cs.ARcs.LG2025-04被引 8

为Transformer加速设计低成本硬件扩展,显著降低Softmax计算延迟与能耗。

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

  • 通过定制RISC-V指令集扩展,集成新型指数近似算法加速Bfloat16指数运算。
  • Softmax计算延迟降低162.7倍,能耗减少74.3倍,性能提升8.2倍。
  • 适用于需高效推理的Transformer模型部署,无需重训练且精度损失极小。

尽管Transformer主要依赖浮点矩阵乘法,但专用硬件或可编程多核系统的加速使非线性函数(如Softmax)成为性能瓶颈。由于其非逐点、非线性特性,尤其是指数运算最为耗时,加速难度大。为此,我们设计了一种基于Schraudolph方法的新型近似算法,用于Bfloat16指数运算,并将其集成至计算集群中RISC-V核心的浮点单元(FPU),通过自定义指令集架构(ISA)扩展实现,面积开销仅为1%。通过优化软件内核以利用该扩展,相比基线集群,Softmax执行延迟降低162.7倍,能耗减少74.3倍,在GPT-2配置下,FlashAttention-2内核实现8.2倍性能提升和4.1倍能效提高。此外,该方法使多集群系统能够高效完成预训练Transformer模型(如GPT-2、GPT-3、ViT)的端到端推理,延迟和能耗分别最多降低5.8倍和3.6倍,无需重新训练且精度损失可忽略。

原文摘要 · Abstract (English)

While Transformers are dominated by Floating-Point (FP) Matrix-Multiplications, their aggressive acceleration through dedicated hardware or many-core programmable systems has shifted the performance bottleneck to non-linear functions like Softmax. Accelerating Softmax is challenging due to its non-pointwise, non-linear nature, with exponentiation as the most demanding step. To address this, we design a custom arithmetic block for Bfloat16 exponentiation leveraging a novel approximation algorithm based on Schraudolph's method, and we integrate it into the Floating-Point Unit (FPU) of the RISC-V cores of a compute cluster, through custom Instruction Set Architecture (ISA) extensions, with a negligible area overhead of 1\%. By optimizing the software kernels to leverage the extension, we execute Softmax with 162.7$\times$ less latency and 74.3$\times$ less energy compared to the baseline cluster, achieving an 8.2$\times$ performance improvement and 4.1$\times$ higher energy efficiency for the FlashAttention-2 kernel in GPT-2 configuration. Moreover, the proposed approach enables a multi-cluster system to efficiently execute end-to-end inference of pre-trained Transformer models, such as GPT-2, GPT-3 and ViT, achieving up to 5.8$\times$ and 3.6$\times$ reduction in latency and energy consumption, respectively, without requiring re-training and with negligible accuracy loss.

Transformer硬件加速RISC-VSoftmax

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。