最近拜讀了NVIDIA前陣子開源的fastertransformer，對CUDA程式設計不是很熟悉，但總算是啃下來一些，帶大家讀一下硬核源碼。
1. 簡介
英偉達公衆号推送的文章加上配圖其實已經把該要講的很清楚了，主要有以下幾方面：
- 為了減少kernel調用次數，将除了矩陣乘法的kernel都盡可能合并
- 針對大batch單獨進行了kernel優化
- 支援選擇最優的矩陣乘法
- 在使用FP16時使用half2類型，達到half兩倍的訪存帶寬和計算吞吐
- 優化gelu、softmax、layernorm的實作以及選用rsqrt等
不了解底層的同學可能不是很懂，沒事我剛看到的時候也不懂，也不敢問，強撸一下源碼就通透（fang qi）了
2. 硬核源碼解讀
首先簡略說一下第一點優化。Kernel在tensorflow裡的概念是operation的計算實作，在cuda裡是執行一個線程的函數，也是一次計算，隻不過tensorflow的更加宏觀些。每次tensorflow執行一個operation，都要調用對應的OpKernel，試想一個通過TF實作的transformer，有将近60個operation，計算一次要執行60次上述過程，進行頻繁的GPU排程和顯存讀寫。是以fastertransformer盡可能多地對kernel進行了合并。
2.1 整體結構
Fastertransformer目錄下主要有（以下簡稱FTF）：
- fastertransformer：主要源碼
- cuda：優化後的kernel以及對multi-head attention整體的封裝（沒過線性層）
- tf_op：tensorflow operation和OpKernel的注冊（op了解為聲明、Opkenerl是定義）
- trt_plugin：tensorRT的實作（可以支援multi streaming太贊了）
- bertencodertransformer.h：transformer整體的封裝
- sample：cpp、tensorflow、tensrflow_bert、tensorRT的調用FTF的示例
- tools：根據參數選擇最優的矩陣乘法（GEMM=General Matrix Multiplication）
接下來我主要想講一下1.1的細節，1.2可以參考這位大佬的文章，剩下的代碼可讀性很強，基本讀一兩遍就知道了。
除去矩陣乘法，作者把剩下的op合并成了4個（圖中藍色框）：
這四個op的cuda源碼分别在cuda_kernels.cu和open_attention.cu兩個檔案中，接下來研究下每個op。
2.2 add_QKVbias (open_attention.cu)
在FP32時，每個block負責處理一個word(num_head*head_size)的add bias運算，先找到要處理QKV中的一個，再進行運算，因為要transpose，是以把結果存入[bsz, num_head, seq_len, head_size]的矩陣裡。
在FP16是每個block同時處理QKV上的同一個word（可能是因為fp16計算的更快一些），在實際的計算中把half都轉成了half2計算。add的話直接用封裝好的__hadd2運算。使用half2計算的原因原文說的比較清楚：
針對半精度FP16，我們對各個kernel也進行了相應優化。首先，在kernel的實作中，将輸入的half指針轉成half2類型，并使用了half2相關的數學函數。這樣不僅僅可以達到2倍于half的訪存帶寬和計算吞吐，還可以極大地減少指令的發射數量。其次，在SoftMax以及Layer Normalization的操作中，為防止求和溢出，将資料以half2的形式讀入後，會轉成float2類型，來做求和計算。
-- NVIDIA BERT推了解決方案Faster Transformer開源啦
2.3 softmax_kernel (open_attention.cu)
在計算softmax之前對block線程數進行了區間處理，因為block裡的線程數最好是wrap大小(32)的倍數，提高計算效率。
調用kernel之前，會根據batch_size * head_num選擇不同的softmax kernel，主要是為了保證在大batch的情況下的計算效率，這裡為什麼使用120我也不是很清楚，希望懂的朋友助力一下
if(batch_size * head_num <= 120)
{
grid.x = batch_size * head_num * seq_len;
softmax_kernel_v2<DataType_><<<grid, block, 0, stream>>>(qk_buf_, attr_mask, batch_size, head_num, seq_len, scaler);
}
else
{
grid.x = batch_size * head_num;
softmax_kernel<DataType_><<<grid, block, 0, stream>>>(qk_buf_, attr_mask, batch_size, head_num, seq_len, scaler);
}
在算softmax時，分母有個求和操作，用到了經典的parallel reduce算法，可以仔細看看參考，講的比較清楚。
這裡注意，使用最初版源碼的同學們需要照着實作一個blockReduceMax，以防止數值溢出，softmax嚴謹的實作應該是：
def softmax(x):
"""Compute the softmax of vector x."""
exp_x = np.exp(x)
softmax_x = exp_x / np.sum(exp_x)
return softmax_x
2.4 transpose (open_attention.cu)
這裡要transpose回[bsz, seq_len, num_head, head_size]的矩陣。因為c++裡面矩陣是行優先存儲，隻要按順序乘過來就好了（最開始看的有點暈）。
2.5 add_bias_act & add_bias_input_layernorm (cuda_kernels.cu)
如果前面幾個函數啃下來了，這兩個就比較好懂，主要的優化點是：
- x^3 -> x*x*x: c語言中x*x和pow(x,2)哪個計算更快一點？
- rsqrt: Why is SSE scalar sqrt(x) slower than rsqrt(x) * x?
- 還有就是各種half2運算的使用
2.6 gemm (tools/gemm_test)
矩陣乘法根據fp16和fp32的不同在不同的cublas算法中選擇，選擇後記錄到http://gemm_config.in檔案中：
問了下作者，其實fp32也可以使用CUBLAS_GEMM_ALGO0_TENSOR_OP到CUBLAS_GEMM_ALGO15_TENSOR_OP的算法，隻不過存在一些風險（使用後速度提升2倍）。
2.7 trt_plugin
作者額外封裝了一個tensorRT的層，tensorRT主要是通過engine，在給定的context和stream下進行異步計算，提供了multi stream inference的可能。關于TensorRT的異步程式設計推薦一個英偉達的PPT：
CUDA C/C++ Streams and Concurrencydeveloper.download.nvidia.cn
這篇文章寫作周期比較長，主要是源碼比較硬核，邊看邊學cuda和c++，到現在也就懂了80%左右吧。不過fastertransformer是真的香，而且直接用tensorflow也很友善，各位需要inference的朋友們一定要用呀