ME-ViT: A Single-Load Memory-Efficient FPGA Accelerator for Vision Transformers

Vision Transformer (ViT) is a state-of-the-art deep learning architecture that utilizes the Transformer model for computer vision tasks. By leveraging self-attention mechanisms, ViTs have achieved significant advancements in image classification and object detection, revolutionizing the field of visual recognition. The ViT architecture is shown in Figure 2. The input image is broken up into patches and fed into the Transformer Encoder as a sequence. The Transformer Encoder is mainly composed of a multi-headed self-attention block (MSA), a multi-layer perceptron block (MLP), and layer normalization blocks (LN).

where $\mathrm{W^{H}}$ is the hidden layer weights, $\mathrm{B^{H}}$ is the hidden layer bias, $\mathrm{W^{O}}$ is the output layer weights, and $\mathrm{B^{O}}$ is the output layer bias, $\operatorname{GeLU}$ is an activation function[26].

In this work, we design a ViT accelerator by implementing the above functions on FPGA and optimizing the memory accesses in model inference.

Field Programmable Gate Arrays (FPGAs) have been extensively used for accelerating machine learning tasks [27][28]. FPGAs consist of a programmable interconnect of logic gates and on-chip memories (BRAMs, URAMs, LUTRAMs), allowing custom hardware architectures to be designed. This flexibility enables hardware optimizations to meet the specific requirements of various machine learning tasks.

There have been several proposed architectures for accelerating ViT and, more generally, Transformer inference on FPGAs [29, 23, 25, 30]. These architectures typically quantize weights and activations to 8 bits to reduce model size and computation requirements [22, 23]. Some architectures compute Softmax and LayerNorm on the FPGA [30, 25], while others perform this on the host CPU [23]. Computing these functions on the CPU reduces the complexity of the FPGA design, but increases the memory overhead. Wang et al. [25] propose ViA, a ViT accelerator which performs the full calculation per layer on an FPGA but requires write-backs between layers and does not implement the full-size ViT models. Nag et al. [30] target an edge FPGA device and also incur frequent memory reads and write-backs between layers. Sun et al. [23] propose an accelerator design that requires the CPU to perform the Softmax and LayerNorm functions and also necessitates loading and unloading for each matrix multiplication. Lu et al. [29] propose an accelerator for the MSA and MLP layers of a Transformer with on-chip computation of Softmax and LayerNorm, but require memory access between matrix multiplications.

All of the above model architectures write back intermediate layer calculations and frequently load weights for matrix multiplication, which incur a large memory overhead [29, 25, 30, 23]. As these architectures scale to larger sizes, they become limited by memory bandwidth constraints. ME-ViT aims to address these memory constraints by eliminating intermediate write backs between layers and only loading weights once into the design, ensuring the minimal amount of data is transferred. This policy enables our design to easily scale without being constrained by memory limitations.

The Memory-Efficient Processing Element (ME-PE) is designed such that distinct input parameters such as weights and biases are only loaded once into the design from off-chip memory (e.g., DRAM). In addition, there is no unloading of any intermediate values. These two objectives of the ME-PE ensure that the minimal amount of memory traffic is incurred. A final objective of this design is to minimize the number of computation stalls for internal data movement within the PE. The hardware architecture of the ME-PE is shown in Figure 5, centered around a systolic array with multi-purpose buffers to manage data around it.

There are 3 modes of operation for the ME-PE shown in Figure 6: 1) Linear Projection (LP) Mode; 2) Multi-headed Self-Attention (MSA) Mode; and 3) Multi-Layer Perceptron (MLP) Mode. Each ME-PE handles all modes on the same architecture with only one extra buffer created to handle MLP Mode (see Section III-D). Since each ME-PE performs all calculations needed for ViT inference, individual layer calculations do not need to be written back to memory. In addition, the residual layer connections present in the Transformer architecture are retained within the PE and similarly do not require additional memory access.

The ME-PE consists of 3 BRAM buffers and 6 LUTRAM buffers. The Weight Buffer stores a maximum matrix size of $\mathrm{D\times D}$ bytes, and the Feature and Layer Buffers store a maximum matrix of size of $\mathrm{(N+1)\times D}$ bytes. For the base ViT model, $\mathrm{D=768}$ and $\mathrm{N=256}$. The Weight Buffer, Layer Buffer, and Feature Buffer are implemented in BRAM and use $\mathrm{160}$, $\mathrm{64}$, and $\mathrm{64}$ 36k BRAMs respectively, for a total of $\mathrm{288}$ BRAMs. While smaller BRAM allocations could store the required data, a multiple of $\mathrm{P_{SYS}}$ must be used to meet the parallel access needs of the systolic array.

The Q, K, V, Result, and two S Buffers are implemented in LUTRAM due to their small size and parallel data access. The sizes of these buffers are $\mathrm{P_{SYS}\times D_{h}}$, $\mathrm{D\times D_{h}}$, $\mathrm{D\times D_{h}}$, $\mathrm{P_{SYS}\times(2\cdot P_{SYS})}$, and $\mathrm{P_{SYS}\times D}$ respectively. Most buffers are multi-purpose, and the names suggest the general use case.

The Linear Projection (LP) Mode performs BMM to compute Feature Buffer $\mathrm{\times}$ Weight Buffer as shown in Fig 8. This operation is needed for the linear projection of input features and to compute the output linear layer in MSA. In addition to performing matrix multiplication, LP Mode performs the residual connection addition and LayerNorm. The LayerNorm result is stored in the Layer Buffer to be used as an input for the next MSA or MLP block. Since the residual sum is also used as an input to the next block, this sum is stored in the Feature Buffer.

LP Mode places a minimum requirement for BRAM usage to store both the layer matrix (L) and the weight matrix (W) since these weights are repeatedly accessed during linear projection. A matrix multiplication of the same size occurs in the output linear projection of the MSA block, so this mode is reused at that step.

Figure 7 illustrates the task scheduling for this mode. At timestamp 1, the block matrix multiplications $\mathrm{L_{1}\times W_{1}}$ and $\mathrm{L_{1}\times W_{2}}$ are concurrently performed and stored in the Result Buffer. At timestamp 2, the Result Buffer values are added to the residual connection values from the Layer Buffer and stored in the S Buffer. Meanwhile, the next pair of block matrix multiplications are performed. Timestamp 3 calculates the sum and squared sum from the new values in the S Buffer to be later used for LayerNorm. At timestamp 4, the cycle repeats with the next row block from the Feature Buffer. At timestamp 5, the mean and variance are calculated from the sums stored in LN Sum. At timestamp 6, LayerNorm is calculated (see Section III-E1) and the results are stored in the Layer Buffer. The un-normalized values are moved to the Feature Buffer and will be used later for residual connection.

The Multi-Headed Self-Attention (MSA) Mode performs the MSA operation for a single head at a time on the ME-PE architecture shown in Figure 10. At the start of operation, the Layer Buffer contains the LayerNorm result calculated from either the previous LP or MLP operation. The Feature Buffer contains the previous layer result before layer normalization to be used later for the residual connection. Task scheduling of this mode is shown in Figure 9. At the start of operation, the $\mathrm{V_{i}}$ and $\mathrm{K_{i}}$ matrices are calculated and stored in the respective buffers. The weight matrices $\mathrm{W^{V}_{i}}$, $\mathrm{W^{K}_{i}}$, and $\mathrm{W^{Q}_{i}}$ are loaded during the first few block matrix multiplication iterations during the $V_{i}$ calculation. At timestamp 1, the residual data stored in the Feature Buffer is moved to empty space in the Weight Buffer so the head outputs can be stored in the Feature Buffer. Simultaneously, the first block row of $\mathrm{Q}$ for head 1 is calculated by $\mathrm{L_{1}\times W^{Q}_{1}}$ and stored in the Q Buffer. At timestamp 2, the Q Buffer is multiplied by two block columns from the K Buffer until the full row of the S Buffer is calculated. At timestamp 3, the next row block of Q is calculated, the next residual row is stored, and the FP reciprocal is calculated. At timestamp 4, Softmax is calculated (see Section III-E2). The cycle repeats until the last head, where timestamp 5 shows how the residual data is moved to the Layer Buffer as the previously stored data there is no longer needed. This leaves the ME-PE with the output attention matrices stored in the Feature Buffer and residual connection data in the Layer Buffer at the end of the operation.

The Multi-Layer Perceptron (MLP) Mode performs the MLP operation on the ME-PE shown in Figure 11. Like with MSA Mode, the Layer Buffer contains the previously calculated LayerNorm, and the Feature Buffer contains the non-normalized values. The operation scheduling in MLP Mode is shown in Figure 12. At timestamp 1, column blocks $\mathrm{W^{H}_{1}}$ and $\mathrm{W^{H}_{2}}$ are loaded into the Weight Buffer and row blocks $\mathrm{B^{H}_{1}}$ and $\mathrm{B^{H}_{2}}$ are loaded into the Feature Buffer. At timestamp 2, the block matrix multiplications $\mathrm{L_{1}\times W^{H}_{1}}$ and $\mathrm{L_{2}\times W^{H}_{1}}$ are concurrently performed and stored in the Result Buffer. The column blocks $\mathrm{B^{H}_{1}}$ and $\mathrm{B^{H}_{2}}$ are loaded to the V Buffer and row blocks $\mathrm{B^{H}_{1}}$ and $\mathrm{B^{H}_{2}}$ are loaded into the K buffer. The row blocks $\mathrm{B^{H}_{1}}$ and $\mathrm{B^{H}_{2}}$ are staged by being moved to the S Buffer for future addition. At timestamp 3, the row blocks $\mathrm{B^{H}_{1}}$ and $\mathrm{B^{H}_{2}}$ are added to the two blocks stored in the Result Buffer. Instead of GeLU, the ReLU activation function is used for hardware simplicity. Various hardware approximations of GeLU exist and would not affect the scheduling or timing of this design. The activation function results are stored in the Q buffer. This timestamp results in a stall of the systolic array, however it is unavoidable since the result is needed immediately for the next multiply. This only incurs a delay of $\mathrm{P_{SYS}}$ clocks which is insignificant. The next pair of output biases are also loaded. At timestamp 4, the first two output layer block multiplications are computed. At timestamp 5, the results computed in timestamp 4 are added to the staged values in the $\mathrm{S_{1}}$ buffer, and the next two output layer blocks are calculated. This process repeats until all sub-blocks involving $\mathrm{M_{1,1}}$ are calculated. At timestamp 6, the $\mathrm{S_{1}}$ buffer containing the staged values is stored back into the Feature Buffer, and multiplication repeats with $\mathrm{M_{2,1}}$. At timestamp 7, the next pair of Layer Buffer row blocks is processed. This continues until all row blocks in the Layer Buffer have been multiplied by $\mathrm{W^{H}_{1}}$. After all layer row blocks have been processed, the cycle repeats with the next $\mathrm{W^{H}}$ column block. At timestamp 8, $\mathrm{S_{2}}$ is stored and the next row block from the Feature Buffer is loaded to $\mathrm{S_{2}}$. On the last cycle in timestamp 9, the residual data stored in the Weight Buffer is transferred back to the Layer Buffer, replacing the existing values as they are no longer needed.

Two separate S Buffers are needed to reduce stalls for matrix multiplication. After one S buffer is filled it gets written back to Feature Buffer BRAM. While this takes place, the second S buffer gets written to. Otherwise, matrix multiplication would need to be stalled while the S buffer gets unloaded and reloaded with new values. This adds an additional resource overhead to MLP Mode that is not present for other modes. However, since the S buffer is relatively small and implemented in LUTRAM, this incurs minimal penalty.

A Multiple ME-PE architecture (Multi-PE) is proposed that contains parallel instantiations of the ME-PE along with a scheduler to coordinate data traffic between them. The Alveo U200 contains 3 Super Logic Regions (SLRs), with SLRs 0 and 2 having 2275 DSPs each and SLR 1 only having 1317. With a $\mathrm{P_{SYS}=32}$, up to 5 PEs can fit as shown in Figure 14. A smaller $\mathrm{P_{SYS}}$ can fit more SAs in the FPGA, however the BRAM needed for such a design is the same as $\mathrm{P_{SYS}=32}$ since buffering requirements are unchanged. Therefore, only the $\mathrm{P_{SYS}=32}$ design is implemented to maximally utilize the available DSPs. Resource utilization is discussed in more detail in Section IV-B. The Multi-PE architecture can achieve remarkably higher throughput, but an increase in data traffic causes a bottleneck due to the limited 77 GB/s bandwidth to the FPGA DRAM. These results are discussed in Section IV-G.

We implement ME-ViT on the Xilinx Alveo U200 platform, consisting of 5867 DSPs, 1766 36k BRAMs, 892K LUTs, and 1831K FFs. It has 4 channels of DDR memory, and a total bandwidth of 77 GB/s. Experimental results are evaluated independently for each ME-ViT mode, and theoretical values are presented which remove extra latencies added from Vitis synthesis and Place and Route (P&R). All implementations are designed for 300 MHz, and 150 MHz figures are provided to compare with other designs. The ME-PE is designed and evaluated using Vitis HLS 2023.1. Power estimates are calculated using AMD Power Design Manager 2023.1.1.

A single ME-PE is analyzed on four common ViT model sizes shown in Table I. ME-PEs with systolic array size $\mathrm{P_{SYS}=32}$ and $\mathrm{P_{SYS}=16}$ are analyzed to provide insight into the performance scalability across FPGAs of varying DSP resources. As the size of the systolic array decreases, there is a corresponding reduction in total FPS (frames per second). Memory bandwidth also reduces despite smaller systolic arrays requiring more frequent data transfers. This relationship between scale and memory bandwidth is explored in Section IV-E. Finally, Multi-PE results are calculated based on single $\mathrm{P_{SYS}=32}$ ME-PE performance to maximally utilize the available DSPs.

ViT variants shown in Table I are evaluated on the ME-ViT architecture. ViT-B refers to the base ViT model presented in [8] but with 256 input image resolution. DeiT-B, DeiT-S, and DeiT-T refer to model sizes presented in [10], all on 224 input image resolution. The key difference between DeiT models lies in their respective model dimensions: 768, 384, and 192.

Results for hardware utilization are shown in Table II. The three ME-ViT modes are implemented separately, and throughput values in Table III are derived from the latency per mode. Since each mode largely utilizes the same resources but with different control logic, the unified design’s resources would marginally exceed that of the largest mode (MLP Mode). Resources are shown for the $\mathrm{P_{SYS}=32}$ ME-PE for the ViT-B model. Resource consumption is unchanged for DeiT-B, and BRAM usage drops to 176 and 144 for DeiT-S and DeiT-T respectively. For $\mathrm{P_{SYS}=16}$, 256 DSPs are used, with other values remaining the same as buffering requirements do not change.

Performance comparison values for ME-ViT are shown in Table III. Theoretical values are presented which are calculated by leveraging extra parallelism which cannot be achieved through Vitis HLS. These figures are achievable with Verilog synthesis. All synthesized designs achieve an operating frequency of 300 MHz, however, competing solutions (Auto Vit Acc [22] and ViTA [30]) are implemented at lower frequencies. To accurately compare design metrics, 150 MHz figures are provided.

We compare ME-ViT performance against four baseline platforms: (1) CPU only platform (Intel i7-9800X), (2) high-power GPU accelerated platform (Nvidia Titan RTX), (3) low-power GPU accelerated platform (Nvidia Jetson TX2), (4) FPGA ViT accelerator presented in Auto Vit Acc [22], and (5) Edge FPGA ViT accelerator presented in ViTA [30].

A single ME-PE achieves a theoretical throughput of 26.4 FPS, a 1.04$\mathrm{\times}$ improvement over the Auto Vit Acc implementation. Notably, Auto Vit Acc is implemented at 150 MHz which outperforms a single ME-PE at the same clock frequency due to higher DSP usage. The ME-PE has a similar FPS/DSP efficiency as Auto ViT Acc at 150 MHz, and sees a 2.16$\mathrm{\times}$ improvement at 300 MHz. The ME-PE has a high power efficiency of 2.83 FPS/Watt, outperforming all other platforms except for ViTA. In comparison to a GPU with a similar power usage (TX2), ME-ViT has a 4.42$\mathrm{\times}$ improvement in FPS/Watt, demonstrating the high efficiency of the custom architecture approach.

Overall throughput measured in frames per second (FPS) for the 4 models is shown in Table IV. FPS improves as model sizes get smaller, and $\mathrm{P_{SYS}=16}$ experiences roughly 0.25$\times$ the throughput of corresponding $\mathrm{P_{SYS}=32}$ designs. Latencies per ME-ViT mode are presented in Figure 15. MLP Mode exhibits the longest duration out of the three modes, taking approximately 60 percent of execution time across all models. A reduction of input image size from 256 to 224 results in an average 1.17$\times$ improvement for $\mathrm{P_{SYS}=32}$ and an average 1.08$\times$ improvement for $\mathrm{P_{SYS}=16}$. For both systolic array sizes, a reduction of model dimension in half results in an average improvement of 3.7$\times$.

As there exist no published results on memory bandwidth for similar FPGA architectures, a non optimized approach is calculated with the following characteristics that are common in various designs [23, 21, 30]: Each BMM loads two input matrices. If an input block matrix was used for the previous multiply, it remains loaded. 2) All calculated matrix blocks are written back to DRAM. 3) Softmax and LayerNorm[33] are calculated on the CPU and are implicitly included in intermediate write-backs.

Figure 16 shows memory bandwidth figures for ME-ViT on all four models for both $\mathrm{P_{SYS}=32}$ and $\mathrm{P_{SYS}=16}$. Total and peak bandwidth improvements are shown in Table V. Peak improvement occurs in the MSA Mode as this has the most back-and-forth traffic in the unoptimized case. Despite fewer model parameters to transfer, a higher improvement is seen on smaller models since less time is spent on computation and therefore a larger proportion of data movement needs to occur in the same time. There is an approximate 3.8$\times$ reduction in latency between $\mathrm{P_{SYS}=32}$ and $\mathrm{P_{SYS}=16}$, but an approximate 2$\times$ reduction in data transferred. This results in a larger reduction in memory bandwidth for $\mathrm{P_{SYS}=16}$.

The systolic array has a large impact on the computational efficiency, defined as the total computation performed divided by the minimum computation required. Depending on how $\mathrm{P_{SYS}}$ divides both the model dimension and the layer height, computational efficiency can greatly vary. When $\mathrm{P_{SYS}}$ poorly matches these dimensions, the BMMs at the right and bottom boundaries fill only a small portion of the systolic array, leading to wasted computation.

As shown in Figure 17, there are periodic cycles in efficiency ranging from 0.95 to 0.5, with notable peaks occurring at 11, 17, 33, 50, and 66. In addition, simply increasing $\mathrm{P_{SYS}}$ leads to an overall downward trend in efficiency. These results underscore the importance of tuning hardware to fit the model dimensions.

We analyze the performance of the Multi-PE to determine the effectiveness of ME-ViT when scaled to larger FPGAs. Multi-PE performance results are shown in Figure 18. The unoptimized design for ViT-B and DeiT-B can only support 3 ME-PEs before performance is limited by memory bandwidth. For DeiT-S and DeiT-T, the unoptimized design can only support 2 ME-PEs. ME-ViT allows 5 ME-PEs to be supported for all models, resulting in a 1.66$\times$ improvement in both FPS and GOPS (Giga Operations per Second) for ViT-B and DeiT-B, and a 2.5$\times$ improvement for DeiT-S and DeiT-T. The theoretical maximum GOPS for 5 ME-PEs is 3072, yet a maximum of 2682 is achieved due to inefficiencies with irregularly-sized matrix multiplication. Matrix blocks that do not completely fill the systolic array result in unused DSPs, leading to an overall reduction in GOPS.

For larger ViT models such as ViT-Large and ViT-Huge ($\mathrm{D=1024}$ and $\mathrm{1280}$), total BRAM usage will increase to 384 and 608 respectively with $\mathrm{P_{sys}=32}$. ME-PEs of this size can still fit inside a large FPGA like the Alveo U200, but fewer total will fit in the Multi-PE design. This results in a large number of unused DSPs, reducing the total GOPs from the theoretical maximum.