Accelerating ViT Inference on FPGA through Static and Dynamic Pruning

We use $\mathbf{W}_{q}$, $\mathbf{W}_{k}$, $\mathbf{W}_{v}$ $\in\mathbb{R}^{D\times HD^{\prime}}$ to denote the concatenation of weight matrices of all the heads. For example, $\mathbf{W}_{p}\in\mathbb{R}^{D\times HD^{\prime}}\text{ }\mbox{where}\text{ }p=\{q,k,v\}.$ The projection operation (Equation 5) projects the concatenated SA outputs of embedding dimension $HD^{\prime}$ back to dimension $D$ via $\mathbf{W}_{\text{proj}}\in\mathbb{R}^{HD^{\prime}\times D}$. To prune a weight matrix $\mathbf{W}\in\mathbb{R}^{M_{1}\times M_{2}}$, we define a parameterized score matrix $\mathbf{S}\in\mathbb{R}^{m\times n}$ such that, $(m,n)=(\left\lceil\frac{M_{1}}{b}\right\rceil,\left\lceil\frac{M_{2}}{b}\right\rceil)$ where $(b,b)$ is the block size. $\mathbf{S}_{ij}$ denotes the importance score of a parameter block of size $(b,b)$ in the weight matrix $\mathbf{W}$ denoted by the slice $\mathbf{W}(ib:\alpha,jb:\beta)$ where $(\alpha,\beta)=(\min(ib+b,M_{1}),\min(jb+b,M_{2}))$. $\mathbf{S}$ is used to construct a mask $\mathbf{M}\in\mathbb{R}^{M_{1}\times M_{2}}$ via the top-$k$ selection:

wher $\mathbf{M}_{ij}^{\text{block}}(.)$ is a block of size $(b,b)$ in $\mathbf{M}$ corresponding to $s_{ij}$. The masked weight is generated as, $\mathbf{W(\mathbf{M})}=\mathbf{W}\odot\mathbf{M}$ where $\odot$ is the element-wise Hadamard product. The generated masked weight $\mathbf{W(\mathbf{M})}$ is used for the forward pass during training. Note that top-$k$ is the target weight blocks of interest. To compute the gradient of $\mathbf{S}$ during the backward pass, a straight-through estimator (STE) [40, 41, 42] is used that neglects the gradients of $\mathbf{M}$ with respect to $\mathbf{S}$. Additionally, the pruning of $\mathbf{W}_{p}$ in row dimension and the pruning of $\mathbf{W}_{\text{proj}}$ in column dimension follows the same pattern (denoted as alternate pattern), as shown in Figure 2. For example, a head removed from $\mathbf{W}_{p}$ makes the corresponding head in $\mathbf{W}_{\text{proj}}$ redundant, and vice-versa.

MLP Pruning. The weight matrices in MLP are $\mathbf{W}_{\text{int}}\in\mathbb{R}^{D\times D_{\text{mlp}}}\quad\mbox{and}\quad\mathbf{W}_{\text{out}}\in\mathbb{R}^{D_{\text{mlp}}\times D}$. The pruning of $\mathbf{W}_{\text{int}}$ and $\mathbf{W}_{\text{out}}$ follows the approach for pruning MSA. A key difference, however, is in how the score parameters are defined for $\mathbf{W}_{\text{int}}$ and $\mathbf{W}_{\text{out}}$. Specifically, we define the scores as: $\mathbf{S}_{\text{linear}}\in\mathbb{R}^{D_{\text{mlp}}}\text{ }\mbox{where}\text{ }\text{linear}=\{\text{int},\text{out}\}$. The score vectors are defined to prune entire columns of $\mathbf{W}_{\text{int}}$ and entire rows for $\mathbf{W}_{\text{out}}$ (see Figure 3). Masks are generated column-wise/row-wise through top-$k$ selection. The natural parameter partitioning along the heads in MSA makes block pruning more effective in terms of removing entire heads. MLP parameters, on the other hand, lack such a partitioning. We thus focus on removing entire columns/rows for MLP parameters. For model training, we add a norm of the sigmoid of scores to the training loss [23]:

where the loss is updated to penalize the presence of a model parameter, thus driving the model to be sparse. The extent of this penalization is controlled by the hyper-parameter $\lambda$.

MSA is divided into four stages (shown in Figure 7): stage (i) computes $\mathbf{Q}$, $\mathbf{K}$ and $\mathbf{V}$ through $[\mathbf{Q},\mathbf{K},\mathbf{V}]=\mathbf{Z}[\mathbf{W}_{q},\mathbf{W}_{k},\mathbf{W}_{v}]$. $\mathbf{Z}$ is the input token matrix and $[\;]$ is matrix concatenation. Let $\mathbf{Q}_{h},\mathbf{K}_{h}$ and $\mathbf{V}_{h}$ denote the query, key, and value matrices for a head $h$, where $(0\leq h<H)$. The algorithm for executing each matrix multiplication (e.g., dense token matrix $\mathbf{Z}$ multiply by sparse weight matrices $\mathbf{W}_{q},\mathbf{W}_{k},\mathbf{W}_{v}$) in MSA using MPCA is shown in Algorithm 2 and an example is shown in Figure 8. Each CHM computes its corresponding head $[\mathbf{Q}_{h},\mathbf{K}_{h},\mathbf{V}_{h}]$. Since we perform block partitioning for each matrix, the computation is executed in a block-wise fashion. Within each CHM, the PEs within a column share the same column of weight, which is stored in the column buffer (CB). PEs of the same row share the same row of $\mathbf{Z}$. We use $\text{PE}(i,j,k)$ to denote a PE in the $k^{\text{th}}$ CHM at location $(i,j)$ within the CHM. A column of PEs: $\text{PE}(:,j,k)$ share the same column of weight with the corresponding header information (See Figure 5). Thus, each PE in $j^{\text{th}}$ column utilizes the shared header indices (Figure 5) to fetch the corresponding data block from the token matrix (in the local input buffer) to perform block-wise matrix multiplication. The partial results are accumulated in local result buffers.

The stage (ii) computes the attention scores $\mathbf{A}_{h}=\mbox{softmax}(\frac{\mathbf{Q}_{h}\mathbf{K_{h}}^{T}}{\sqrt{D^{\prime}}}),\text{ }(0\leq i<H)$. $\mathbf{Q}_{h}\mathbf{K}_{h}^{T}$ is dense block-wise matrix multiplication executed via the MPCA module (See Algorithm 2). $\mathbf{Q}$ is buffered in the GFB, and $\mathbf{K}^{T}$ is buffered in the CB. The output data blocks of $\mathbf{Q}_{h}\mathbf{K}_{h}^{T}$ are sent to EM module for element-wise scaling (by $1/\sqrt{D^{\prime}}$) and exponentiation to obtain $\text{exp}(\frac{\mathbf{Q}_{h}\mathbf{K}_{h}^{T}}{\sqrt{D^{\prime}}})$. Then, we utilize MPCA to compute the scaling factors for $\mbox{softmax}(\frac{\mathbf{Q}_{h}\mathbf{K}_{h}^{T}}{\sqrt{D^{\prime}}})$. The rows of matrix $\text{exp}(\frac{\mathbf{Q}_{h}\mathbf{K}_{h}^{T}}{\sqrt{D^{\prime}}})$ and their corresponding computed scaling factors, are streamed from MPCA to EM to perform the scaling to obtain the attention scores, $\mathbf{A}^{H}$.

The stage (iii) computes the self-attention $\mathbf{A}_{h}\mathbf{V}_{h}$. It is similar to the computation of $\mathbf{Q}_{h}\mathbf{K}_{h}^{T}$. The stage (iv) computes the projection (Equation 5). It is similar to the computation of $\mathbf{Q}$, $\mathbf{K}$ and $\mathbf{V}$ described in stage (i) as $\mathbf{W}_{\text{proj}}$ is the block-wise sparse matrix due to pruning.

Since the weights of MLP are pruned for entire columns or rows (for $\mathbf{W}_{\text{int}}$ and $\mathbf{W}_{\text{out}}$ respectively), MLP layers can be mapped into dense block-wise matrix-matrix multiplication executed by MPCA (Algorithm 2). This computation is similar to the computation of MSA (computing $\mathbf{Q}\mathbf{K}^{T}$). GELU activation is computed using the EM module.

We design a token dropping hardware module (TDHM) for on-the-fly token dropping and reorganizing the remaining tokens. The token pruning is based on importance scores of tokens $\mathbf{\mathcal{S}}$ (See Secton IV-B). The attention scores $\mathbf{A}_{h}$ for all the heads are buffered in the TDHM as soon as they are computed via MSA execution. Then, scores $\mathbf{\mathcal{S}}=\frac{1}{H}\sum_{h}\mathbf{A}_{h}$ are computed via the EM. After that, a bitonic sorting network sorts the scores $\mathbf{\mathcal{S}}$ to obtain the indices of top-$k$ tokens. Original, each token has a row index in the input token matrix $\mathbf{Z}_{\text{in}}$, denoted as old token index ($\text{id}_{\text{old}}$). After sorting by $\mathbf{\mathcal{S}}$, each token is assigned new token index ($\text{id}_{\text{new}}$) which is the row index in the output token matrix $\mathbf{Z}_{\text{out}}$. Therefore, the sorting network generates ($\text{id}_{\text{old}}$, $\text{id}_{\text{new}}$, flag) for each token, where flag indicates if the token will be pruned. To organize output token matrix $\mathbf{Z}_{\text{out}}$ (stored in Old Token Buffer), an index shuffle network routes each $(\text{id}_{\text{old}},\text{id}_{\text{new}},$flag$)$ to Old Token Buffer for fetching tokens according to $\text{id}_{\text{old}}$. Then, the fetched tokens are routed to the New Token Buffer according to $\text{id}_{\text{new}}$, which generates too-$k$ token matrix and non-top-$k$ token matrix. The non-top-$k$ tokens are then fused into a single token and merged with the top-$k$ tokens to produce the output of TDHM module.

As the weight matrices are pruned in block-wise fashion, different columns in a weight matrix can have different number of data blocks, which can potentially lead to load imbalance. The multi-level parallelism in MPCA distributes the PEs across several heads (CHMs). PEs within each CHM computes on weight column blocks using multiple iterations that each iteration executes different columns of weight matrices (See Algorithm 2). This naturally reduces the impact of load imbalance due to differing sparsity levels across the columns of weight matrices. As we restrict block-wise pruning to only the weights of MSA ($\mathbf{W}_{q}$, $\mathbf{W}_{k}$, $\mathbf{W}_{v}$, $\mathbf{W}_{\text{proj}}$), this further reduces the impact of load imbalance. Note that the gain in performance by block-wise pruning the MSA parameters (leading to the removal of entire heads) is far outweighted by the load imbalance presented by such pruning. Prior methods [44] that balance such block-wise pruning across columns are disadvantaged by the fact that they cannot remove entire heads. Moreover, we perform offline workload assignment among columns of weight marices, prior to inference, such that workloads of columns are evenly distributed across different columns of PEs within a CHM.

Dynamic token pruning, leads to varying number of tokens for different layers. Moreover, block-wise weight pruning the MSA parameters leads to removal of heads within an encoder. In general, the heads removed or retained in each encoder can vary, which can potentially lead to runtime hardware underutilization. For example, if the number of rows of token matrix $\frac{N}{b}<p_{t}$, $p_{t}-\frac{N}{b}$ rows of PEs in a CHM will be idle. As we utilize multi-level parallelism in MPCA, through selecting proper $p_{\text{t}}$ (parallelism in token dimension) and $p_{\text{h}}$ (parallelism in head dimension), we can alleviate the underutilization. We use $\frac{N_{\text{min}}}{b}$ to denote the minimum number of row blocks of all the intermediate token matrices and use $H_{\text{min}}$ to denote the minimum number of heads of all the layers. Through setting $p_{\text{t}}\ll\frac{N_{\text{min}}}{b}$, the PEs utilization in a CHM will be $>\frac{\frac{N_{\text{min}}}{p_{t}\times b}}{\lceil\frac{N_{\text{min}}}{p_{t}\times b}\rceil}$ (Suppose $6\times p_{\text{t}}<\frac{N_{\text{min}}}{b}$. The utilization will be $>85\%$ ). Similar, we can set $p_{\text{h}}\ll H_{\text{min}}$.

We perform theoretical analysis for the performance achieved by the codesign and its hardware resource utilization. We denote the total computational resources utilized by the MPCA, EM and TDHM as $R_{\text{MPCA}},R_{\text{EM}}$ and $R_{\text{TDHM}}$, respectively. $R_{\text{MPCA}}$ is proportional to the total number of computation units: $p_{\text{t}}p_{\text{h}}p_{\text{c}}p_{\text{pe}^{2}}$. Compared to $R_{\text{MPCA}}$, the resources used by $R_{\text{TDHM}}$ and $R_{\text{EM}}$ are negligible, and thus ignored for analysis. The total $R_{\text{Total}}$ (DSPs and LUTs) are $R_{\text{Total}}=(c_{1}p_{\text{t}}p_{\text{h}}p_{\text{c}}p_{\text{pe}}^{2},c_{2}p_{\text{t}}p_{\text{h}}p_{\text{c}}p_{\text{pe}}^{2})$, where $c_{1}$ and $c_{2}$ denote the amount of DSPs and LUTs utilized by a single computation unit. The size of the (global) feature buffer, column buffer and the (global) result buffer, associated with the MPCA, are $b^{2}p_{t}\gamma$, $b^{2}p_{c}\gamma$ and $b^{2}p_{\text{t}}p_{\text{h}}p_{\text{c}}$, respectively. Here, $b$ is the block size and $\gamma$ is a constant that equals the (maximum) total number of row blocks required to compute a single output block. We match the buffer sizes across compute units to improve the dataflow performance of the accelerator. This gives a total buffer size of $4\times\max({b^{2}p_{\text{t}}p_{\text{h}}p_{\text{c}},b^{2}p_{t}\gamma})$ for the EM module and $2\times\max({b^{2}p_{\text{t}}p_{\text{h}}p_{\text{c}},b^{2}p_{t}\gamma})$ for the TDHM module. The EM module requires a buffer to store the input, scaling factor, addtion factor and the output. Similarly, TDHM requires an input and an output buffer. The total size of required buffers $B_{\text{Total}}$ is given as, $B_{\text{Total}}=b^{2}p_{t}\gamma+b^{2}p_{c}\gamma+b^{2}p_{\text{t}}p_{\text{h}}p_{\text{c}}+6\times\max({b^{2}p_{\text{t}}p_{\text{h}}p_{\text{c}},b^{2}p_{t}\gamma})$. $R_{\text{Total}}$ and $B_{\text{Total}}$ are the estimation of resource utilization. The main design $p_{\text{t}}$, $p_{\text{h}}$ and $p_{\text{c}}$, with $\gamma$ are empirically set according the resource of target FPGA platform (See Section VI for details).

Based on algorithm 2, the number of cycles to perform either SBMM, DBMM or DHBMM is estimated in table III. Note that DHBMM is DBMM computed head-wise (as in stage (ii) of MSA execution). In table III, the cycles for SBMM/DBMM are the cycles required to multiply a matrix of dimension $(M_{1},M_{2})$ with a matrix of dimension $(M_{2},D)$. $D^{\prime}$ is the size per head, $b$ is the block size and $\phi$ is the ratio of retained dense blocks to total blocks within a column of the matrix. Note that for DBMM, $\phi$ is $1$ and for SBMM, $\phi$ is assumed similar in each column block for simplicity. $(M_{1},M_{2})$ and $(M_{2},D)$ are the per head left and right matrix sizes, for DHBMM, with $H$ being the total number of heads. The cycle estimates in Table III can be used to compute the total cycles for the MSA and the MLP blocks.

We compare the latency and throughput for executing the pruned model with baseline CPU and GPU (Figure 9 and 10). The latency of our accelerator is measured when batch size is $1$, and the throughput is calculated by $\frac{1}{\text{latency}}$. For comparing the latency, we set the batch size as $1$ for CPU and GPU because a larger batch size will increase the latency for CPU and GPU. For throughput comparison, we set the batch size as $8$ for CPU and GPU, which can fully exploit their thread-level parallelism. On average, our FPGA accelerator achieves a latency reduction of $12.8\times$ and $3.2\times$, compared with CPU and GPU, respectively. The lower latency of our FPGA accelerator is due to the followings: (1) our MPCA module with load balance strategy fully exploits the computation parallelism within the pruned model. The FPGA accelerator achieves a higher speedup with higher pruning ratios (smaller $r_{b}$ and $r_{t}$). In contrast, the CPU and GPU cannot efficiently handle the computational irregularity caused by weight pruning and dynamic token dropping. (2) CPU and GPU have complex cache hierarchies, leading to higher memory access latency for executing ViT inference, leading to increased latency. On average, our FPGA accelerator achieves $3.6\times$ and $0.45\times$ throughput speedup compared with CPU and GPU, respectively. Our FPGA accelerator achieves a lower throughput ($0.45\times$) than GPU, because GPU has much higher peak performance ($50\times$) and eternal memory bandwidth. When the pruning ratios become high (e.g., $r_{b}=0.5$ and $r_{t}=0.5$), our throughput gets closer to GPU, which indicates that our FPGA accelerator has higher efficiency for executing the ViT model with larger pruning ratio.

We compare the proposed codesign with the state-of-the-art ViT Accelerators [48, 37, 35] on FPGA as shown in Table VII. Prior works use at most one pruning approach. ViTAcc [48] and [37] use int4 or int8 to represent the weights and activations. In contrast, our work is the first algorithm-hardware codesign to combine two pruning approaches. In terms of latency, our accelerator achieves $6.2-18.5\times$ latency reduction compared with the prior accelerator. As different accelerators use different numbers of computation units, which directly influences their peak performance (shown in Table V), we further normalize the latency by their respective peak performance ($\text{Normalized Latency}=\text{Latency}\times\text{Peak Performance}$) to obtain a fair comparison. Our accelerator achieves a normalized speedup of $1.5-4.5\times$ compared with SPViT [35] and achieves a normalized speedup of $0.72-2.1\times$ compared with HeatViT [37]. Our accelerator achieves higher speedup by executing the model with a higher pruning ratio. The achieved speedup is attributed to (1) in addition to token pruning, we further utilize the model pruning to reduce the computational complexity compared with [35, 37], (2) our architecture design using MPCA can efficiently utilize the block-wise data sparsity in the pruned model.