如何实现一个高效的Softmax CUDA kernel？——OneFlow 性能优化分享

Softmax操作是深度学习模型中最常用的操作之一。在深度学习的分类任务中，网络最后的分类器往往是Softmax + CrossEntropy的组合：

尽管当Softmax和CrossEntropy联合使用时，其数学推导可以约简，但还是有很多场景会单独使用Softmax Op。如BERT的Encoder每一层的attention layer中就单独使用了Softmax求解attention的概率分布；GPT-2的attention的multi-head部分也单独使用了Softmax 等等。

深度学习框架中的所有计算的算子都会转化为GPU上的CUDA kernel function，Softmax操作也不例外。Softmax作为一个被大多数网络广泛使用的算子，其CUDA Kernel实现高效性会影响很多网络最终的训练速度。那么如何实现一个高效的Softmax CUDA Kernel？本文将会介绍OneFlow中优化的Softmax CUDA Kernel的技巧，并跟cuDNN中的Softmax操作进行实验对比，结果表明，OneFlow深度优化后的Softmax对显存带宽的利用率可以接近理论上限，远高于cuDNN的实现。

GPU基础知识与CUDA性能优化原则：

对GPU基础知识的介绍以及CUDA性能优化的原则与目标可以参考之前的文章：

https://zhuanlan.zhihu.com/p/271740706

其中简单介绍了GPU的硬件结构与执行原理：

• Kernel：即CUDA Kernel function，是GPU的基本计算任务描述单元。每个Kernel都会根据配置参数在GPU上由非常多个线程并行执行，GPU计算高效就是因为同时可以由数千个core（thread）同时执行，计算效率远超CPU。

• GPU的线程在逻辑上分为Thread、Block、Grid三级，硬件上分为core、warp两级；

• GPU内存分为Global memory、Shared memory、Local memory三级。

• GPU最主要提供的是两种资源：计算资源 和 显存带宽资源。如果我们能将这两种资源充分利用，且对资源的需求无法再降低，那么性能就优化到了极限，执行时间就会最短。在大多数情况下，深度学习训练中的GPU计算资源是被充分利用的，而一个GPU CUDA Kernel的优化目标就是尽可能充分利用显存带宽资源。

如何评估一个CUDA Kernel是否充分利用了显存带宽资源？

对于显存带宽资源来说，“充分利用“指的是Kernel的有效显存读写带宽达到了设备显存带宽的上限，其中设备显存带宽可以通过执行 cuda中的的bandwidthTest得到。 Kernel的有效显存带宽通过Kernel读写数据量和Kernel执行时间进行评估：
$$当前Kernel的有效显存带宽=读写数据量/执行时间$$

Naive的Softmax实现：

在介绍优化技巧之前，我们先看看一个未经优化的Softmax Kernel的最高理论带宽是多少。如下图所示，一个最简单的Softmax计算实现中，分别调用几个基础的CUDA Kernel function来完成整体的计算：

假设输入的数据大小为D，shape = (num_rows, num_cols), 即 D = num_rows * num_cols ，最Navie的操作会多次访问Global memory，其中：

• ReduceMax = D + num_rows （read 为 D， write 为 num_rows）
• BroadcaseSub = 2 * D + num_rows （read 为 D + num_rows，write 为 D）
• Exp = 2 * D （read 、write 均为D）
• Redu