最近拜读了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的朋友们一定要用呀