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

要約

変圧器はフローティングポイント（FP）マトリックス総拡大によって支配されていますが、専用のハードウェアまたは多くのコアプログラマブルシステムを介した積極的な加速により、パフォーマンスボトルネックがSoftMaxのような非線形関数にシフトしました。
SoftMaxの加速は、その非点向で非線形の性質のために挑戦的であり、指数は最も要求の厳しいステップです。
これに対処するために、Schraudolphの方法に基づいて新しい近似アルゴリズムを活用するBFLOAT16指数のカスタム算術ブロックを設計し、カスタム指導セットアーキテクチャ（ISA）の拡張を介して、コンピュータクラスターのRISC-Vコアのフローティングポイントユニット（FPU）に統合します。
ソフトウェアカーネルを最適化して拡張機能を活用することにより、ベースラインクラスターと比較して162.7 $ \ Times $ LaTENCYおよび74.3 $ \ TIMES $ ENERGYでSoftMaxを実行し、GPT-2構成のフラッシュ2 Kernelのために8.2 $ \ Times $のパフォーマンス改善と4.1 $ $ \ Times $のエネルギー効率を達成します。
さらに、提案されたアプローチにより、マルチクラスターシステムは、GPT-2、GPT-3、VITなどの事前に訓練された変圧器モデルのエンドツーエンドの推論を効率的に実行し、それぞれ最大5.8 $ \ Times $と3.6 $ \ Times $の削減を、それぞれ再訓練を必要とせずに、それぞれ3.6 $ \ Times $削減を達成します。

要約(オリジナル)

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.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google