Source-linked AI summary
Hardware Accelerator for Multi-Head Attention and Position-Wise Feed-Forward in the Transformer
Siyuan Lu, Meiqi Wang, Shuang Liang, Jun Lin, Zhongfeng Wang
TL;DR
Transformer accelerators face intensive matrix computation, complicated data flow, and complex nonlinear functions, while existing DNN accelerators largely target CNNs or RNNs. This paper proposes an FPGA accelerator for Transformer MHA and FFN ResBlocks using shared matrix partitioning, systolic-array-aware scheduling, and optimized nonlinear functions, achieving 14.6× and 3.4× speed-ups over an NVIDIA V100 GPU.
Problem
Transformer hardware acceleration is difficult because of intensive matrix computations, complicated data flow, and complex nonlinear functions, while prior DNN accelerators mainly target CNNs or RNNs.
Method
The paper designs a reconfigurable FPGA accelerator for MHA and FFN ResBlocks using shared matrix partitioning, systolic-array computation flow, and optimized nonlinear functions.
Results
14.6× and 3.4× speed-ups were measured for the MHA and FFN ResBlocks, respectively, against an NVIDIA V100 GPU implementation.
Takeaways & Limitations
The FPGA design lays a foundation for efficient hardware accelerators targeting multiple Transformer networks and supports operation in mobile or embedded systems.
Abstract
from arXiv · showhide
Designing hardware accelerators for deep neural networks (DNNs) has been much desired. Nonetheless, most of these existing accelerators are built for either convolutional neural networks (CNNs) or recurrent neural networks (RNNs). Recently, the Transformer model is replacing the RNN in the natural language processing (NLP) area. However, because of intensive matrix computations and complicated data flow being involved, the hardware design for the Transformer model has never been reported. In this paper, we propose the first hardware accelerator for two key components, i.e., the multi-head attention (MHA) ResBlock and the position-wise feed-forward network (FFN) ResBlock, which are the two most complex layers in the Transformer. Firstly, an efficient method is introduced to partition the huge matrices in the Transformer, allowing the two ResBlocks to share most of the hardware resources. Secondly, the computation flow is well designed to ensure the high hardware utilization of the systolic array, which is the biggest module in our design. Thirdly, complicated nonlinear functions are highly optimized to further reduce the hardware complexity and also the latency of the entire system. Our design is coded using hardware description language (HDL) and evaluated on a Xilinx FPGA. Compared with the implementation on GPU with the same setting, the proposed design demonstrates a speed-up of 14.6x in the MHA ResBlock, and 3.4x in the FFN ResBlock, respectively. Therefore, this work lays a good foundation for building efficient hardware accelerators for multiple Transformer networks.
I. INTRODUCTION
The paper addresses the difficulty of accelerating Transformer inference by proposing a reconfigurable FPGA architecture for the MHA and FFN ResBlocks, the model’s two most complex layers. It combines shared matrix partitioning, systolic-array-aware computation flow, and optimized nonlinear functions, achieving substantial speedups over a GPU implementation.
- I. INTRODUCTION: Transformer accelerators are motivated by high computation complexity, large memory requirements, intensive matrix operations, complicated data flow, and complex nonlinear functions.These constraints make Transformer networks difficult to operate on mobile and embedded systems and difficult to implement efficiently in hardware.
- I. INTRODUCTION: The proposed reconfigurable systolic-array architecture targets the MHA and FFN ResBlocks, which contain most Transformer parameters and computations.These are identified as the Transformer’s two most complex layers.
- I. INTRODUCTION: Efficient matrix partitioning lets the MHA and FFN ResBlocks share most hardware resources.The approach is intended to support both ResBlocks within one accelerator design.
- I. INTRODUCTION: A carefully designed computation flow aims to maintain high utilization of the systolic array, the design’s largest module.The paper presents this as part of its first hardware architecture supporting both ResBlocks.
- I. INTRODUCTION: Scaled masked-softmax and layer normalization are optimized for hardware, with layer-normalization latency reduced as an architectural bottleneck.The optimizations reduce hardware complexity and latency of the system.
- I. INTRODUCTION: 14.6× and 3.4× speed-ups were achieved for the MHA and FFN ResBlocks, respectively, versus an NVIDIA V100 GPU implementation under INT8, sequence-length-64, batch-size-1 settings.The evaluation used a Xilinx xcvu13pfhga2104-3-e FPGA and the Transformer base model.
II. BACKGROUND AND MOTIVATION
Transformer encoder and decoder layers are composed of MHA and FFN ResBlocks, which contain the principal computational structures targeted by the accelerator. MHA uses multiple attention heads with shared input tensors, while FFN uses two linear sublayers separated by ReLU.
- A. The Model Architecture of the Transformer: Each Transformer encoder and decoder layer contains an MHA ResBlock and an FFN ResBlock.The embedding and softmax output layers are outside the scope of this work.
- A. The Model Architecture of the Transformer: An MHA ResBlock has h attention heads, each receiving the same V, K, and Q input tensors.The passage introduces the attention-head structure and its shared inputs.
- A. The Model Architecture of the Transformer: The mask operation removes softmax inputs corresponding to illegal connections, while d_k equals 64 and h equals 8 in the base model or 16 in the big model.These parameter values define the stated Transformer configurations.
- A. The Model Architecture of the Transformer: The FFN ResBlock consists of two fully connected linear sublayers with a ReLU activation between them.This defines the FFN computation structure used in the Transformer layers.
B. Transformer-Based Pre-Trained Models
The paper targets the MHA and FFN ResBlocks because they dominate Transformer computation and storage, and partitions their matrices so one s × 64 systolic array can support most GEMMs.
- B. Transformer-Based Pre-Trained Models: MHA and FFN ResBlocks occupy most Transformer storage and FLOPs, motivating dedicated acceleration for these widely used structures.
- B. Transformer-Based Pre-Trained Models: The design partitions WG, W1, and W2 so most GEMMs in both ResBlocks can use the same s × 64 systolic array.The partitioning follows dmodel = 64h and dff = 4dmodel = 256h.
- B. Transformer-Based Pre-Trained Models: The accelerator must handle larger products P × WG, X × W1, and ReLU(XW1 + b1) × W2 beyond the linear sublayers within individual attention heads.
- B. Transformer-Based Pre-Trained Models: The resulting matrix strategy supports the remaining larger multiplication with little systolic-array utilization impact through zero padding or partitioning Qi when needed.Zero padding applies when s is below 64; otherwise Qi is partitioned.
IV. HARDWARE ARCHITECTURE DESIGN FOR THE PROPOSED ACCELERATOR
The complete accelerator uses the partitioned matrices in a shared architecture centered on an s × 64 systolic array, with computation scheduled to keep that array active.
- IV. HARDWARE ARCHITECTURE DESIGN FOR THE PROPOSED ACCELERATOR: The complete accelerator is built from the proposed weight-matrix partitioning and organized in the top-level architecture shown in Fig. 5.
- IV. HARDWARE ARCHITECTURE DESIGN FOR THE PROPOSED ACCELERATOR: The s × 64 systolic array uses a two-dimensional processing-element grid with s rows and 64 columns, outputting product columns sequentially.
- IV. HARDWARE ARCHITECTURE DESIGN FOR THE PROPOSED ACCELERATOR: The computation flow overlaps Softmax execution with V × Wvi + BiasV so the systolic-array module can run with minimal interruption.The schedule is designed to keep the systolic array running until layer normalization begins.
A. Scaled Masked-Softmax
The Softmax module computes scaled masked softmax from the s × s matrix Qi × KT, using hardware-oriented transformations to reduce the cost of exponentiation and division.
- A. Scaled Masked-Softmax: The module names Qi × KT as D, an s × s input matrix, and produces output matrix Y using mask matrix M.
- A. Scaled Masked-Softmax: The design uses log-sum-exp and reductions for exponential and logarithmic functions, avoiding regular multipliers and lookup tables in the Softmax module.
B. Layer Normalization
Layer normalization lies on the critical path of both ResBlocks, so the design reduces its latency by starting computation early and reorganizing variance calculation.
- LayerNorm is always on the critical path because both ResBlocks must compute it before producing output.
- The straightforward LayerNorm computation adds at least 128h cycles for calculating the row mean and variance.
- The proposed method uses s accumulators connected to the module input to calculate row sums in advance.
- It also chooses an alternative variance calculation so output can begin after very few cycles once matrix G is complete.
- The LayerNorm architecture implements the x^(-0.5) unit with a lookup table.
A. Quantization of Transformer Base Model
The Transformer base model is quantized to INT8 for FPGA evaluation, with Softmax handled separately because its exponential and logarithmic functions use linear approximations.
- The Transformer base model is quantized to INT8 for a machine translation evaluation using the IWSLT 2016 German-English corpus.
- After quantizing trainable and activation matrices while retaining FP32 Softmax internals, BLEU decreases from 23.88 to 23.48.
- Quantization proceeds in two steps because Softmax uses linear approximations for its exponential and logarithmic functions.
B. Hardware Implementation Results
The FPGA implementation is evaluated at batch size 1 and sequence length 64, then compared with a GPU implementation under the same settings to assess layer latency.
- The FPGA completes the MHA and FFN ResBlocks in 21,344 and 42,099 cycles, respectively, at batch size 1 and sequence length 64.
- The design runs at up to 200MHz with total on-chip power of 16.7W, comprising 13.3W dynamic and 3.4W static power.
VI. CONCLUSION AND FUTURE WORK
The paper presents the first hardware accelerator for the Transformer’s MHA and FFN ResBlocks, with FPGA results showing speed and power benefits for embedded deployment.
- The work presents the first hardware accelerator for the Transformer’s MHA and FFN ResBlocks.
- The FPGA implementation shows promising speed and power results that can contribute to operating Transformer networks in mobile or embedded systems.