Proposed by ETH Zurich and the University of Bologna, the VEXP extension optimizes RISC-V for transformer-based AI models by accelerating softmax computation, a critical operation in large language models.
Enhances RISC-V’s competitiveness in AI hardware, targeting efficient processing for machine learning workloads. This is a step toward closing the gap with proprietary ISAs like ARM and x86 in AI applications.
Positions RISC-V as a viable platform for next-generation AI accelerators, with potential for widespread adoption in edge and cloud AI systems.
In a revolutionary leap for AI hardware, researchers from ETH Zurich and the University of Bologna have unveiled VEXP, a RISC-V Instruction Set Architecture (ISA) extension that turbocharges softmax computation in transformer models, addressing a critical bottleneck that outstrips General Matrix Multiply (GEMM) operations. Detailed in a study by Run Wang and a team of reseachers, VEXP delivers unparalleled performance and energy efficiency for large language models (LLMs) like GPT-2, GPT-3, and Vision Transformers (ViT), cementing RISC-V’s role as a powerhouse in AI innovation.
Transformers, the cornerstone of modern AI, have seen GEMM operations optimized through dedicated hardware, making non-linear functions like softmax the new performance limiter. Unlike GEMM, which benefits from highly parallelizable, linear computations, softmax’s non-pointwise, non-linear nature driven by exponentiation demands complex, sequential processing, creating a significant bottleneck. VEXP tackles this by integrating a custom Bfloat16 (BF16) exponentiation block into the Floating-Point Unit (FPU) of RISC-V’s Snitch cluster, an energy-efficient architecture with eight RV32IMAFD cores. Leveraging an enhanced Schraudolph’s approximation algorithm, VEXP achieves a mean relative error of just 0.14%, preserving model accuracy without retraining or fine-tuning.
The results are transformative. VEXP slashes softmax latency by 162.7× and energy consumption by 74.3× compared to a baseline Snitch cluster, requiring only 6.39 pJ per exponential operation versus 3433 pJ for software-based methods. Integrated into the FlashAttention-2 kernel for GPT-2, it boosts throughput by 8.2× and energy efficiency by 4.1×. Scaled to a 16-cluster system, VEXP accelerates end-to-end inference of GPT-2, GPT-3, and ViT by up to 5.8× in latency and 3.6× in energy, with accuracy losses below 0.1%. Remarkably, these gains come with a mere 1% area overhead and 1.8% power increase, implemented in GlobalFoundries’ 12nm technology.
VEXP’s hybrid hardware-software approach introduces two new instructions FEXP for scalar and VFEXP for SIMD operations executing in just two cycles, compared to 319 cycles for baseline software exponentiation. Unlike fixed-point accelerators, which often require fine-tuning or sacrifice flexibility, VEXP maintains native BF16 precision, offering 1.4× better area efficiency and 7.4× lower power consumption than state-of-the-art softmax accelerators. Its integration into the open-source RISC-V ecosystem ensures versatility, from edge devices to datacenters.
As AI workloads escalate, VEXP’s ability to overcome softmax’s computational complexity where GEMM’s parallelism falls short positions RISC-V as a leader in energy-efficient AI hardware. With Qualcomm and Google adopting RISC-V for wearables and China’s 1nm Wuji processor showcasing its potential, VEXP could drive RISC-V’s dominance, heralding a new era of open-source AI innovation.
This research work is supported by NeuroSoC project, funded under the European Union’s Horizon Europe research and innovation programme.
The paper titled “VEXP: A Low-Cost RISC-V ISA Extension for Accelerated Softmax Computation in Transformers” is available on researchgate.net
Source: ResearchGate
