# Hardware Accelerator for Multi-Head Attention and Position-Wise Feed-Forward in the Transformer
<body>
<div style="text-align: right;">
<p>CASLAB Group Meeting</p>
<p>2024.02.20</p>
<p>IEEE SOCC</p>
<a href="https://arxiv.org/abs/2009.08605">Paper link</a>
</div>
---
[TOC]
---
## Background
### Transformer
#### 1. Transformer Structure
#####     以中英文翻譯為例來說明
#####     由 Encoder和Decoder組成，分別包含n, m個block。
#####     Step1
#####      input向量X，X=單字的Embedding(從原始data萃取出的特徵)+位置的Embedding
#####     Step2
#####      將X傳到Encoder，經過n個Encoder之後得到矩陣C，C = X_n * d(通常d=512)。
#####      其中每個Encoder block output dimension = input dimension
#####     Step3
#####      Encoder output C傳到Decoder，Decoder依序輸出透過Mask來遮住後面的單字。
#####    
#### 2. Position Embedding(PE)
#####    
#####     d = PE dimension = 單字dimension = input dimension
#### 3. Self-Attention
#####    
#####     Transformer中的Multi-Head Attention就是由多個Self-Attention組成
#####     在計算的時候會用到矩陣Q, K, V，三者皆由X轉換而來
#####     X * WQ = Q; X * WK = K; X * WV = V;
#####     得到Q, k, V之後即可求Self-Attention output
#####    
#####    
#####     其中d_k是Q, K的列數
#####     接著Q * K transpose之後，做softmax
#####     其中Softmax過後，每一行的和=1
#####    
#### 4. Multi-head Attention(MHA)
#####     MHA包含多個Self-Attention
#####    
#####     各自生成h個輸出矩陣 Z_1 ~ Z_h 之後
#####    
#####     將 Z_1 ~ Z_h 連接後做線性變換，最終output Z與input X的dimansion相同
#### 5. Add and Norm
#####     X表示MHA或FFN的input
#####    
#####     這邊的Add是指residual connection
#####    
#####     這邊的Norm是指Layer Normalization
#####     
#####     由上圖可以看出BatchNorm 與 LayerNorm最大的差別
#####     
#####     
#####     
#### 6. Feed Forward
#####     兩層的全連接層，第一層的Activation Function是Relu，第二層不用Activation Function
#####     其中input是X，輸出與X的dimension相同
#####    
#### 7. Decoder 與 Encoder差異
#####     包含兩個MHA
#####     第一個MHA有Masked
#####     第二個MHA的K, V使用Encoder的輸出，Q是Masked的輸出
#####     最後有一個Softmax
#### 8. Decoder Masked(第一個MHA)
#####     
#####     正常做完Q*K transpose之後 * Masked Matrix
#####     其中Masked Matrix的Masked部分即一極小負數，如此一來接著做softmax後，該部分趨近0
#### 9. Decoder Masked(第二個MHA)
#####     K, V使用Encoder的輸出，Q是Masked的輸出
### Systolic array
#####     
#####     上圖為SA的經典示意圖，因為往往access memory都比實際計算時間還要長，所以下面那個就是一個1D SA的範例，一次將同一個data隨時間傳遞給一整排的PE
#####     SA的優點是結構簡單且成本低，但缺點是靈活性很差，而且資料需要適當安排，下面以一個2D 3*3 SA為例
#####     
#####     本篇文章就是用2D SA的概念，來加快GEMMs運算
## Paper : Hardware Accelerator for Multi-Head Attention and
Position-Wise Feed-Forward in the Transformer
### Introduction
#####     本篇是第一個為Transformer提出特定硬體加速器的文章。
#####     MHA和FFN是Transformer的bottleneck，因此本篇文章也只著重在這2個部分，即下圖紅色框框，傳入硬體即為Q, K, V matrix，不管前面的Embedding處理。
#####     
#####     貢獻 :
#####     1.提供一個有效的方法劃分Transformer的巨大陣列，使得MHA和FFN可以共享硬體資源。
#####     2.提供了第一個硬體架構設計，可以完成這兩個Res-block的計算，其中為了確保SA的硬體Reuse rate足夠高，計算流程需要特別設計。
#####     3.兩個最複雜的非線性函數scaled masked-softmax和LayerNormalization，都經過高度優化，變得對硬體更加友善。其中LayerNormalization的delay被盡可能地減少。
#####     在Quantization使用int8的Transformer後，在Xilinx xcvu13pfhga2104-3-e FPGA上評估，當max sequence length=64且batch size=1的情況下，相比NVIDIA V100 GPU，MHA的部分有14.6X加速比，FFN則有3.4X
### Partitioning Matrices in the MHA and FFN
#####      Q, K and V通常形狀都是相同的，皆為s*d_model
#####      
#####     
#####     從論文"Attension is all you need"中的這張表，可以看出來基本上transformer 的d_model參數都是64的倍數。
#####     所以本篇論文以s * 64為一塊，去拆分每一個大矩陣。
### MHA ResBlock
#####     
#####     以上面這張圖來看的話，Q, K and V的形狀都是s * d_model，每一個head的weight都是d_model * 64的形狀，如此一來乘s * d_model形狀的矩陣之後，會得到s * 64的Q_h, K_h and V_h，接著K_h做transpose後與Q_h乘，然後做softmax後 * V_h，這樣就做完一個head的一次self-attention。接著將多個head的output接在一起之後， * W_G，即完成了一次MHA。
#####     註: 上面提到的Q, K and V是加速器的input，即一般transformer的input X(推測是為了保有彈性，沒有硬性規定這邊的Q=K=V)
#####     MHA的output + Q，即residual connection(殘差連接)，得到Matrix G。
#####     Matrix G接著做LayerNorm之後就是一次MHA
#####     註 : 這張圖上的 general matrix-matrix multiplications (GEMMs)幾乎都能用s * 64的SA完成，且使用率很高。
#####     除了Q * K_transpose的時候，如果s比64小的話就會需要再額外切成像下圖這樣，有可能會造成浪費，且需要分好幾次做。
#####     
#####     
#####     但如果是s比64大，的話就會需要再額外切成像下圖這樣，有可能會造成浪費，且需要分好幾次做。
#####     
#####     
#####     只有當s = 64的時候，才不會有浪費。
#####     但這部分只佔所有GEMMs運算的一小部分
#####     
#####     因為s通常不大於128的關係，所以可以化減成這條式子。
### FFN ResBlock
#####     
#####     上圖是FFN的部分，這邊就很單純，X * W_1 + b_1，接著重點是ReLU，因為這邊要取Max(0, XW_1 + b_1)，因此取ReLU的話>0的就會被保留，要是都沒有>0的部分就都會是0。接著 * W_2後 + b_2即FFN的輸出。接著 + input X 即residual connection(殘差連接)，之後做LayerNorm，如此一來即Encoder的output。
### Partition Weight Matrix
#####     
#####     本篇論文是以s * 64為單位，所以以上圖為例的話，實際上會拆成許多s * 64的矩陣，且以常見的transformer model來看，s通常都是能整除d_model的。
### Pseudo Code
#####     
#####     由上到下
#####     第一個for loop :
#####     跑h次self-attension，這邊的h是指head數量
#####     第二個for loop :
#####     將self-attension的輸出 * W_G，且 + Q (即Add and Norm的Add)，這邊的h是指weight被切成h塊
#####     第三個for loop :
#####     用ReLU取max(0, XW_1 + b_1)
#####     第四個for loop :
#####     * W_2 + b_2之後 + X (即Add and Norm的Add)
### System Architecture
#####     
#####     -------------------------------------------------------------------------------- MHA start --------------------------------------------------------------------------------
#####     Step1.
#####     左邊的mux選Q，上面的mux選weight memory(W_Q傳進來)，傳入SA module，做完QW得到Q之後，傳到旁邊的adder加bias
#####     Step2.
#####     左邊的mux選K，上面的mux選weight memory(W_K傳進來)，傳入SA module，做完KW得到K之後，傳到旁邊的adder加bias
#####     Step3.
#####     左邊的mux選Q，上面的mux選Temp2(從上面adder的output拉過來)，傳入SA module，做完Q * K_transpose之後傳入Softmax Module(這裡面會做/(d_k^0.5))
#####     Step4.
#####     Softmax Module在做事的時候，SA是空閒的，所以可以平行運算。因此左邊的mux選V，上面的mux選weight memory(W_V傳進來)，傳入SA module，做完VW得到V之後，傳到旁邊的adder加bias
#####     Step5.
#####     左邊的mux選Temp1(softmax output傳進來)，上面的mux選Temp2(從旁邊adder的output拉過來)，傳入SA module，做完 * V之後，即完成一次self-attension
#####     Step6.
#####     左邊的mux選P，上面的mux選weight memory(W_G傳進來)，傳入SA module，做完PW得到P之後，傳到旁邊的adder加bias
#####     Step7.
#####     Step6的output傳到下面的adder，input Q 拉進來adder(綠色那條線)，加完之後傳到LayerNorm module
#####     Step8.
#####     從LayerNorm Module出來之後即完成MHA
#####     註 : step1 ~ step5要做h次，才會做step6
#####     -------------------------------------------------------------------------------- MHA end --------------------------------------------------------------------------------
#####     -------------------------------------------------------------------------------- FFN start --------------------------------------------------------------------------------
#####     Step1.
#####     左邊的mux選X，上面的mux選weight memory(W_1傳進來)，傳入SA module，做完XW得到X之後，傳到旁邊的adder加bias
#####     Step2.
#####     將adder的output傳入ReLU，即得到P
#####     Step3.
#####     左邊的mux選P，上面的mux選weight memory(W_2傳進來)，傳入SA module，做完PW得到P之後，傳到旁邊的adder加bias
#####     Step4.
#####     Step3的output傳到下面的adder，input X 拉進來adder(綠色那條線)，加完之後傳到LayerNorm module
#####     -------------------------------------------------------------------------------- FFN end --------------------------------------------------------------------------------
### Softmax module
#####     
#####     公式4中的D = Q * K_transpose, Y = softmax output, i = head number, j = D's row number
#####     由公式4可知，當M=0，即沒有Mask的時候，正常做softmax，當M=1時，輸出直接給0。
#####     因為softmax和V matrix的SA運算是平行的，所以只要確保softmax的運算不要變成critical path即可。觀察式子可以發現又做除法又做exp對於硬體負擔還是蠻大的，因此將公式4化簡為公式5。
#####     
#####     其中X_max是為了避免overflow的常見方法，即分子分母同除exp(X_max)。
#####     接著使用log sum-exp trick，即可推出最終式，詳細推導如下。
#####     
#####     如此一來，就可以避免掉除法的計算。
#####     
#####     上圖中，向右shift 3 bits即除8的意思
#####     Stage 1 : 藍箭頭
#####     當M = 0時，找D最大值，後傳入EXP Unit
#####     當M = 1時，mux選擇負無限大
#####     Stage 2 : 紫箭頭
#####     D和D_max傳入EXP Unit之後，相減做exp，接著傳出去做累加
#####     Stage 3 : 綠箭頭
#####     將累加完的exp(D - D_max)做ln
#####     Stage 4 : 橘箭頭
#####     接著做完最後的相減之後做exp，輸出即為所求Y
#####     註 : 推測EXP Unit裡面有減法器
### Layer Normalization module
#####     
#####     註 : 這裡講的64h cycle指G matrix有64h rows
#####     原本在做Layer Normalization的流程是最上面的框框，計算完Matrix G之後依序算出E(G)和var(G)再求出output
#####     其中E(G)和var(G)是2個最花時間的步驟，因此此篇論文的方法就是想辦法讓這2個步驟能和Calculate Matrix G平行運算，如此一來就能大幅度降低總體運算時間。
#####     
#####     式子6是常見的LayerNorm公式，其中ε = 10^-8，為了確保分母不要為0，E(G)為G每一列的平均，var(G)為G每一列的變異數，γ是縮放參數，β是平移參數
#####     
#####     
#####     本論文中，處理E(G)的方法是在計算Matrix G的時候就先累加，不用等全部的G算完才累加。
#####     處理var(G)的方法是將式子8改寫完下式，推導也是常見的變異數公式
#####     
#####     
#####     接著就是硬體實現的部分
#####     
#####     藍色框框是將計算完的G傳進來做累加，算出E(G)之後求出E(G)^2
#####     橘紅色框框是做(G - E(G)) * γ
#####     綠色框框是計算變異數，將G^2 之後做累加，然後與E(G)^2從藍色框框傳進來，與ε一同加減之後，開根號
#####     最終傳回橘紅色框框做相除，然後+ β即為LayerNorm的output
#####     註 : 開根號是使用LUT查表
### Result : Quantization
#####     Training Data :
#####     IWSLT 2016 German-English parallel corpus，德英平行語料庫是指由國際語音和語言技術協會（International Workshop on Spoken Language Translation，簡稱IWSLT）在2016年發布的一個德語到英語的平行語料庫
#####     “tst2014”
#####     註 : BLEU score為Decoder output與答案做比較得到的分數
| 配置 | BLEU score |
| -------- | -------- |
| All use FP32 | 23.88 |
| GEMMs use INT8 | 23.48 |
| GEMMs use INT8 and quantize softmax module | 23.57 |
### Result : Hardware Implementation
#####     
| FPGA | |
| -------- | -------- |
| Parameter | batch size = 1, s = 64 |
| Board | Xilinx xcvu13p-fhga2104-3-e FPGA |
| Environment | Vivado 2018.2|
| Operating Freq. | 200MHz|
| Power | 16.7W (13.3W dynamic and 3.4W static)|
#####     
#####     
#####     上表是在s=64最理想的情況下，跑transformer base model，只看MHA與NVIDIA V100相比，加速比為14.6X，只看FFN，加速比為3.4X
### Evaluation
#####     s * 64 MACs
#####     Input SRAM：2 * (s * dmodel) B (GEMMs INT8)
#####     Weight SRAM：(dk * dff) B
#####     Bias SRAM：(4 * dmodel) B
#####     Output SRAM：(s * dmodel) B
#####     Others SRAM：(s * dff) B and (s * dk) and (s * dmodel)
#####     
| Suppose transformer base model | h=8, dk=64, s=64 |
| -------- | -------- |
| input | 64KB |
| weight | 128KB |
| bias | 2KB |
| output | 32KB |
| temp1 | 32KB |
| temp2 | 4KB |
| P | 128KB |
| total | 390KB |
---
## Reference
[【機器學習2021】Transformer (下)](https://youtu.be/N6aRv06iv2g?si=DEdk0J5Pbfqzzydg)
0:11:30 Transformer architecture
0:23:40 NAT
0:26:00 Cross attention
0:56:00 GLEU score
[Transformer模型詳解](https://zhuanlan.zhihu.com/p/338817680)
[Attention is all you need 導讀](https://www.youtube.com/watch?v=nzqlFIcCSWQ)
[LayerNormalization詳解](https://blog.csdn.net/sinat_34072381/article/details/106173365)
[Systolic Array](https://zhuanlan.zhihu.com/p/26522315)
</body>
Sign in with Wallet
Connect another wallet