CUDA算力计算编程实例

CUDA算力计算编程实例

算力芯片的“CUDA难题”

“生成式AI时代和AI的iPhone时刻已经到来”。北京时间8月8日晚间，英伟达创始人黄仁勋在计算机图形学顶会SIGGRAPH上发布了GH200 Grace Hopper 超级芯片、AI Workbench等成果时这样讲道，并透露首批GH200预计于2024年第二季度出货。在全球范围内白热化算力争霸的当下，英伟达已经赢得了竞争先机，而其并行计算和编程平台CUDA可能是最重要的“幕后英雄”。凭借强大而丰富的计算库，CUDA得到了算法工程师的认可，甚至被认为是“巩固英伟达硬件市场地位的护城河”。因此，CUDA成为摆在其他算力芯片企业眼前的两难问题：兼容还是不兼容？

兼容CUDA是因为“好用”

CUDA之所以会成为算力芯片硬件厂商必须要认真考虑的一个选择，最直接的原因，是其已经实现了与算法客户的强绑定。众多算法工程师已经习惯了CUDA提供的工具库及其编程语言，向外迁移总是会存在不习惯的问题。

因此，很多算力芯片硬件厂商选择了兼容CUDA的路线——使硬件能够直接用CUDA调动起来，这样可以降低用户更换硬件平台的不适感。

换手机有“一键迁移”功能。而如果兼容CUDA——记者在采访中了解到——换芯片硬件也可以像换手机一样容易。如果芯片客户原本用的是CUDA软件和英伟达的卡，在向兼容CUDA的算力硬件迁移过程中，原有的代码一个都不改，直接就能用起来。这样一来，客户尝试选用新的算力芯片的意愿将大大提升。

算法工程师对CUDA的使用习惯并不是一天养成的。

英伟达GPU及其并行计算和编程平台CUDA进入行业时间早，这为英伟达积累了非常明显的先发优势。从当前火热的大模型的发展历程来看，许多大模型发展之初都是基于开源机器学习库PyTorch训练，最早期使用的也是英伟达的GPU，与之相配套，算法工程师便会大量地用到CUDA库代码。从机器学习出现到大模型盛行，经过十几年的发展，算法工程师们不可避免地会用到CUDA库支持的算子。这样一来，如果使用CUDA驱动硬件，大模型写好之后，编程语言就能直接运行起来，实现无缝对接。

CUDA主页

包括大模型企业在内，使用计算芯片的企业所采用的硬件基础也基本上是由英伟达的产品搭建的。在一次计算训练中，统一硬件品牌，也能更好保证其训练的稳定性。因为英伟达产品入局早，许多工程师对CUDA的使用习惯甚至是从学校受教育阶段便培养起来的，在毕业参加工作后，各企业采用的软硬件工具也是英伟达厂牌。这样一来，算法工程师在软件工具使用上存在非常大的惯性。

这对于英伟达之外的算力芯片来说意味着，如果不兼容CUDA，市场推广过程将存在一定的困难。

某算力芯片企业告诉记者，要想打开市场，短期内必须兼容CUDA生态。如果不兼容CUDA，就会出现很多问题，包括要在代码上做微调，跑模型时要实现收敛等问题，这背后的工作量相当大。但如果整个软硬件能够支持CUDA，客户就无需再做二次开发或者修改算法。这被认为是降低客户使用成本最简单的方式。

兼容CUDA也是在帮助客户节约时间。对于以大模型为主营业务的企业而言，时间就是生命。模型推演快一天、比竞争对手的产品早一天上线都是非常重要的。因此，若是硬件迁移需要工程师花大量的时间适应软件工具，将极大地降低客户对新产品的接受意愿。

 

来源：CUDA官网，《中国电子报》整理

“要站在用户的角度思考问题”，这是企业选择兼容CUDA的立足基点。客户开发大模型，是以效率为第一要义的，很多软件工程，也都是拿到现成的代码再做调优。现在很火的大模型的开发也是同样的，很多厂家先拿到小模型和机器学习的代码，再在此基础上做累加和调优，最终实现规模化。

要让客户觉得“好用”，首先要解决的是“能用”的问题。也就是首先得将客户积累了十几甚至二十几年的软件基础跑起来，解决“从0到1”的问题。只有这个问题解决了，才能考虑“好用”的问题、 “从1到100”的问题——能帮助客户通过编程支持新功能。

从企业战略的角度来看，英伟达的生态已经形成比较成熟的市场、标准和护城河。基于这样的事实，在既有标准的基础上切入，对于算力芯片企业而言，将是更加便利且高效的路径。

兼容CUDA也是一把“双刃剑”

不过，也有业界专家提醒记者，许多宣称“兼容”CUDA的产品，并不是真的“兼容”，而是通过架构的相似性，使产品可以比较容易地运行CUDA的代码。因为CUDA不是开源代码，是“黑盒子”，因此100%兼容CUDA在技术上是无法实现的。同时在安全、知识产权方面存在风险。

“兼容”CUDA就像一把双刃剑，一方面可以降低算力芯片企业的获客成本，另一方面，也会在一定程度上给产品的创新潜能带来限制。

许多算力芯片企业选择在软件栈上另辟蹊径，主要是因为担心“兼容”CUDA会成为自家产品发展“天花板”。当前“兼容”CUDA的方式，主要是试图将编程模型做到与CUDA一致，但实际底层的硬件架构难以实现与英伟达的GPU完全相同。由此，兼容CUDA很有可能只能是短期行为，长期来看将不利于芯片架构创新与性能的提升。而自研软件栈，则相当于将创新的“天花板”掌握在自己手中。

从全球来看，许多国外的做算力芯片企业，例如 Graphcore、SambaNova、tenstorrent，没有一家兼容CUDA。而之所以不去兼容CUDA，归根结底，是因为各家希望探索出更适合做AI计算的路。英伟达的GPU早期只做GPU图形渲染，只是因为其产品能够借助CUDA在通用计算领域，也就是在包括AI在内的领域用起来，才使其逐渐成为AI算力芯片领域的领头羊。但从英伟达产品的计算架构和性价比来看，都不是最合适的AI计算的产品。

换句话说，如果英伟达是最合适的，可能也就没有这么多AI芯片公司存在了。

算力芯片的后来者，要做的是从前人的发展中吸取经验，但依然要摸索自己的路。CUDA的出现，最初只是为了使GPU能够满足除视觉处理之外的其他通用计算功能。而之所以CUDA会富有市场竞争力，核心在于它实现了从硬件层到软件层的全跑通，且基于其百万级的客户体量，实现了对性能的优化。

芯片设计、软件开发，这是一件亟需创新能力的事情。如果还想走CUDA的老路，只是尽可能模仿英伟达的产品，其实就没有往把AI芯片做得更好的路径上走，兼容CUDA某种程度也是在壮大英伟达的生态，增强对CUDA的依赖性。

而如果不兼容CUDA，走自主创新的道路，产品和企业发展的路径都将更宽。因此有的企业选择了走自定义编程模型的路线，提供从硬件平台到开发者工具包、计算库和框架的一整套方案。

而对于这条路可能存在的获客困难问题，有的算力芯片企业向《中国电子报》记者表示，如果客户下定决心向英伟达生态之外的其他产品迁移，其实平台间迁移的困难并没有想象中那么大，客户还是能获得丰厚商业价值回报的。

底层逻辑是构建开发者生态

软硬件协同完善的生态是客户选择CUDA的根本原因，也是英伟达领先于其他算力芯片企业最核心的竞争力。

尽管在是否兼容CUDA的问题上，算力芯片设计企业具有不同的观点，但在搭建企业生态的问题上，记者接触的算力芯片企业都给出了一致的回答：要建设企业自己的生态体系。

至于何为生态、如何搭建生态，业界的观点存在些许差异。

有的企业选择从指令集开始，到计算库和编译器等层次进行自研，构建软硬件相结合的生态。构建自己的软件栈首先是打好基础，对标CUDA及以下的抽象层次，充分发挥自己芯片的特色，开发出一套用户可用、易用的编程模型。积硅步以至千里，最终帮助目标客户完成从英伟达到自家芯片的平滑迁移，同时又能使客户在使用自己的软件栈时充分感受到新特性带来的更优体验。

CUDA生态合作伙伴

所谓的“生态”底层逻辑是开发者生态。AI芯片作为技术属性很强的产品，其核心价值是帮助开发者在这个硬件的基础上加速其算法开发与业务部署。生态建设的成功与否取决于这款产品能否给客户带来价值。例如学生学习了这款产品的知识能够帮助其找到工作，企业使用该产品后能实现其业务目标，并且市场上也有相应熟悉该产品的人才储备等等；生态里的每一个角色都能获得利益才是生态健康成长的关键。

因此，有的算力芯片厂商给出了这样的发展建议：国内厂商协同定义自研编程模型，以此联合拓展开发者，集聚企业的力量，让更多的高校、商业伙伴使用。

大模型的出现，为生态搭建从芯片厂各自为战走向产业联合提供了契机。

如果没有大模型，各家都会选择适合自己的通用方式，缺少将供应链上下游集合到一起解决问题的利益取向。用简单的话来说，如果是三五张卡一台服务器就能解决的问题，一家企业单点调优就可以实现了。而大模型是需要大算力和大互联的东西，动辄要调动上万张卡、上千台服务器，还要考虑供电等各种问题，最后呈现的是关乎生产基础甚至国计民生的东西，它所需要的资源就不是一家企业能够做到的，因此更需要产业链的协作，更有机会调动产业链更多的资源。

不仅如此，大模型的出现，也为算力芯片市场突破提供了机会。大模型的发展使人工智能的技术路线出现了收敛的趋势。在大模型出现之前，对应一个需求的解决方案非常繁杂，每个企业都会提出自己的方案。而大模型的发展，使得算法技术路线逐渐统一于Transformer模型，从而也为算力芯片日后需要瞄准的技术方向树立了“靶子”。在这种情况下，如果算力芯片企业能够争取到大的算力供应商（例如大的互联网公司）作为合作伙伴，再以此为基础推出结合其他行业的特点做微调，进而产出适合多行业的解决方案，就有很大的机会赢得更广阔的市场前景，也就有了打破英伟达商业壁垒的机会。

关于如何实现产业链协作，实现算力芯片的破局，记者采访的企业表示，可以联通可控的供应链，在所有的供应链成员中了解自己有哪些独特的技术，评估有竞争力的地方，结合自己的技术特色解决人工智能发展中存在的问题。、

要做好自己的产业生态，当前最缺的是大量工程师的调优工作。英伟达的GPU加速库中有数学库、并行算法库、图像和视频库、通信库、深度学习库等多个类型，有100个加速库，每个计算库又积累了开源代码。这是英伟达几万工程师耕耘了20年，通过解决客户的问题所积累起来的结果。

对于算力芯片企业而言，要做好自己的软件生态，当前不存在很多科研问题，很多问题都可以通过开源代码找到答案。当前业内存在的是工程问题，需要安下心来一点一点做，需要依靠众多工程师的力量，集中力量办大事。现在算力芯片从0到1的突破已经实现了，剩下的就要看时间和积累。

 

【CUDA编程】Faster Transformer v1.0 源码详解

写在前面：本文将对 Nvidia BERT 推理解决方案 Faster Transformer 源码进行深度剖析，详细分析作者的优化意图，并对源码中的加速技巧进行介绍，希望对读者有所帮助。本文源码解读的内容仅限 Faster Transformer v1.0 版本，更高版本的源码将在后续文章中继续解读。

1 Faster Transformer

Transformer 模型最早在 2017 年由谷歌在论文中提出，抛弃了以往深度学习任务里面使用到的 CNN 和 RNN，取而代之的是一种 self-Attention 的结构，将 Attention 思想发挥到了极致，一定程度上解决了 RNN 的长序列信息丢失问题，基本取代了 RNN 的地位。
Transformer 一经面世便在 NLP 领域大杀四方，近两年一些文章开创性地将 Transformer 模型跨领域地引用到了 CV 任务中，并取得了不错地成果。这也被许多学者认为是开创了 CV 领域的新时代，甚至可能完全取代传统的卷积操作。
虽然 Transformer在多种场景下都有优秀的表现，但是在推理部署阶段，其计算性能却受到了巨大的挑战：以 BERT 为原型的多层 Transformer 模型，其性能常常难以满足在线业务对于低延迟和高吞吐的要求。以 BERT-BASE 为例，超过 90% 的计算时间消耗在 12 层 Transformer 的前向计算上。因此，一个高效的 Transformer 前向计算方案，既可以为在线业务带来降本增效的作用，也有利于以 Transformer 结构为核心的各类网络在更多实际工业场景中落地。
基于上述背景，NVIDIA GPU 计算专家团队针对 Transformer 推理提出了的性能优化方案：Faster Transformer。
Faster Transformer 是一个 BERT Transformer 单层前向计算的高效实现，其代码简洁明了，后续可以通过简单修改支持多种 Transformer 结构。目前优化集中在编码器（encoder）的前向计算。底层由 CUDA 和 cuBLAS 实现，支持 FP16 和 FP32 两种计算模式，其中 FP16 可以充分利用 Volta 和 Turing 架构 GPU 上的 Tensor Core 计算单元。

2 优化原理

在深入了解 Faster Transformer 的优化原理之前，我们先来了解一下主流深度学习框架 Tensorflow 中 Transformer 的实现情况，仅仅以一个基本的激活函数 gelu 为例，这个函数在框架中是通过 8 个类似 Pow、Add、和 Tanh 等基本 OP 来实现的。也就是说每进行一次 gelu 运算要调用 8 次基本 OP，同时底层也对应 8 次 GPU kernel 的调用，每一次调用也意味着一次显存读写，先不说 kernel 计算耗时，光是显存读写就已经是大量的开销。如何降低这部分开销？最直观的方法就是减少调用，让数据一直留在显存甚至寄存器里被访问，即 OP 融合，一次调用就实现整个计算逻辑。
出于性能最大化的考虑，在 Faster Transformer 内部，Nividia 将除矩阵乘法以外的所有 kernel 都进行了尽可能的融合，单层 Transformer 的计算流程如下图所示:

 

如图所示，基于 OP 融合的思想，Faster Transformer 只用了 14 个 kernel 就完成了原来将近 60 个 kernel 的计算逻辑。这其中，8 个 kernel 是通过调用 cuBLAS 接口计算矩阵乘法（黄色框），其余 6 个是自定义 kernel（蓝色框）。
接下来笔者将沿调用链逐步介绍每个 kernel 的优化逻辑。

3 调用链

Faster Transformer v1.0 版本源码地址如下，有兴趣的读者可以前往阅读。

https://github.com/NVIDIA/FasterTransformer/tree/v1.0/fastertransformer

通读源码后笔者对调用关系梳理如下。

BertEncoderTransformer->forward()
    ->OpenMultiHeadAttention->forward()
        ->cublasGemmEx
        ->cublasGemmEx
        ->cublasGemmEx
        ->multiHeadAttr_nofuse_kernelLauncher
            ->add_QKV_bias  (kernel)
            ->cublasGemmStridedBatchedEx
            ->softmax_kernel    (kernel)
            ->cublasGemmStridedBatchedEx
            ->transpose (kernel)
    ->cublasGemmEx
    ->add_bias_input_layernorm_kernelLauncher   (kernel)
    ->cublasGemmEx
    ->add_bias_act_kernelLauncher   (kernel)
    ->cublasGemmEx
    ->add_bias_input_layernorm_kernelLauncher   (kernel)

从调用链也可以看出，总共 14 个步骤，与示意图一一对应。核心逻辑都在两个类中实现：BertEncoderTransformer 和 OpenMultiHeadAttention。

4 OpenMultiHeadAttention

OpenMultiHeadAttention 类中有两个重要的成员方法：构造函数、forward 方法。其中构造函数内主要进行一些参数初始化功能，设备内存的申请和初始化也在该函数内进行。forward 方法内主要是多头注意力机制核心逻辑的具体实现。

4.1 cublasGemmEx for Q、K、V

forward 方法中首先就是对输入的 3 个 tensor 进行线性变换，其实就是对 3 个 tensor 分别进行 Dense 层变换，我们知道 Dense 是包含一个矩阵乘法和一个 add_bias 操作，这里只进行矩阵乘法，add_bias 操作放在后面的 kernel 进行。这里使用了 cuBLAS 接口计算矩阵乘法，具体代码如下：

check_cuda_error(cublasGemmEx(param_.cublas_handle,
    CUBLAS_OP_N, CUBLAS_OP_N,
    n, m, k,
    &alpha,
    param_.attr_kernel_Q, AType_, n,
    param_.from_tensor, BType_, k,
    &beta,
    query_buf_, CType_, n,
    computeType_,
    static_cast<cublasGemmAlgo_t>(cublasAlgo_[0])));

check_cuda_error(cublasGemmEx(param_.cublas_handle,
    CUBLAS_OP_N, CUBLAS_OP_N,
    n, m, k,
    &alpha,
    param_.attr_kernel_K, AType_, n,
    param_.to_tensor, BType_, k,
    &beta,
    key_buf_, CType_, n,
    computeType_,
    static_cast<cublasGemmAlgo_t>(cublasAlgo_[0])));

check_cuda_error(cublasGemmEx(param_.cublas_handle,
    CUBLAS_OP_N, CUBLAS_OP_N,
    n, m, k,
    &alpha,
    param_.attr_kernel_V, AType_, n,
    param_.to_tensor, BType_, k,
    &beta,
    value_buf_, CType_, n,
    computeType_,
    static_cast<cublasGemmAlgo_t>(cublasAlgo_[0])));

这里仅仅是矩阵乘法 API 的调用，按文档传参即可，这里不展开介绍，笔者计划另开一篇文章专门介绍这个 API 的调用方法。

4.2 multiHeadAttr_nofuse_kernelLauncher

4.2.1 add_QKV_bias

上面说过 Dense 层包含矩阵乘法和 add_bias 操作，其中 add_bias 操作在核函数 add_QKV_bias 中完成，源码针对两种数据类型 fp16 和 fp32 分别提供了一个 kernel，只是网络结构有所差异。
针对 fp32，每个 block 处理一个 word，总共有 batch_size * seq_len * 3 个 block，对于 Q、K、V 3 个 tensor 而言，前 batch_size * seq_len 个 block 处理 Q，中间 batch_size * seq_len 个 block 处理 K，后 batch_size * seq_len 个 block 处理 V。示意图如下：

Add_QKV_bias grid structure (FP32)

/**
 * @brief
 *
 * @tparam T                          OperationType
 * @param Q                           [batch_size, seq_len, head_num, size_per_head], query
 * @param bias_Q                      [head_num * size_per_head, ] length is the same as word's embedding dim
 * @param K                           [batch_size, seq_len, head_num, size_per_head], key
 * @param bias_K                      [head_num * size_per_head, ] length is the same as word's embedding dim
 * @param V                           [batch_size, seq_len, head_num, size_per_head], value
 * @param bias_V                      [head_num * size_per_head, ] length is the same as word's embedding dim
 * @param q_buf_                      [batch_size, head_num, seq_len, size_per_head], transpose query & add bias
 * @param k_buf_                      [batch_size, head_num, seq_len, size_per_head], transpose key & add bias
 * @param v_buf_                      [batch_size, head_num, seq_len, size_per_head], transpose value & add bias
 * @param batch_size
 * @param seq_len                     
 * @param head_num
 * @param size_per_head
 * @param word_per_block              1
 * @return __global__
 */
template<typename T>
__global__
void add_QKV_bias(T* Q, const T* bias_Q, T* K, const T* bias_K, T* V, const T* bias_V, T* q_buf_, T* k_buf_, T* v_buf_,
  const int batch_size, const int seq_len, const int head_num, const int size_per_head, const int word_per_block)
{

  T* data_ptr;
  T* buf_ptr;
  const T* bias_ptr;
  // word counts per batch
  int m = batch_size * seq_len;
  // word embedding dim
  int n = head_num * size_per_head;
  // 总共有3m个block，第一部分处理q，第二部分处理k，第三部分处理v，这里使用qkv_id区分处理哪个矩阵
  int qkv_id = blockIdx.x * word_per_block / m;
  // 矩阵偏移量
  int row_offset = (blockIdx.x * word_per_block % m) * n;

  if(qkv_id == 0)
  {
    data_ptr = Q + row_offset;
    buf_ptr = q_buf_;
    bias_ptr = bias_Q;
  }
  else if(qkv_id == 1)
  {
    data_ptr = K + row_offset;
    buf_ptr = k_buf_;
    bias_ptr = bias_K;
  }
  else
  {
    data_ptr = V + row_offset;
    buf_ptr = v_buf_;
    bias_ptr = bias_V;
  }

  int batch_id = (blockIdx.x * word_per_block % m) / seq_len;
  int head_id = threadIdx.x / size_per_head;
  int id_in_head = threadIdx.x % size_per_head;
  // word_id in seq, not data_index
  int word_start_id = (blockIdx.x * word_per_block) % seq_len;

  T bias = __ldg(&bias_ptr[threadIdx.x]);

  for(int i = word_start_id; i < word_start_id + word_per_block; ++i)
  {
    // add bias
    T tmp = data_ptr[threadIdx.x] + bias;
    // buf's shape: [bacth_size, head_num, seq_len, size_per_head]
    int target_id = batch_id * (seq_len * head_num * size_per_head) + head_id * seq_len * size_per_head +
      i * size_per_head + id_in_head;

    buf_ptr[target_id] = tmp;
    data_ptr += n;
  }
}

核函数第一部分先根据 block_id 确定当前处理的 tensor 具体是 Q、K、V 中的哪一个，从而拿到输入输出变量的内存地址。
第二部分求出 tensor 中对应元素的索引，首先我们知道输入输出 tensor 是一个四维的 array，所以应该有四个索引，按维度顺序依次是 batch_id、word_start_id、head_id、id_in_head，有读者看到这里可能会有疑问：为什么要计算这些索引，前面计算了矩阵偏移量 row_offset，完全可以在 block 内按 thread_id 索引就可以拿到对应元素。原因是在 add_QKV_bias 核函数中计算逻辑不仅仅是 add，还有 transpose，熟悉 multiHeadAttention 的读者都知道对 Q、K、V 线性映射之后，紧接着就是一个 transpose 操作，目的是把 embedding_dim 这个维度划分成多个独立的 head，每个 head 后面单独进行 attention，所以要把 head 维度移到 seq_len 维度前面。换句话说这里的 transpose 解决的是“多头”的问题，和 attention 无关。
理解了前面的逻辑，第三部分就比较简单了，先进行 add 操作，然后将结果按照 [bacth_size, head_num, seq_len, size_per_head] 的维度顺序写在输出 tensor 中，这里隐含了一个 transpose，需要注意的是这个 transpose 操作是输出的 tensor 中元素存储顺序相对于输入 tensor 而言的，并不是对输入 tensor 做了变换。

针对 fp16，每个 block 同时处理 Q、K、V 上的同一个 word，同一个线程先后处理 3 个 word 上对应元素的计算逻辑，实际计算中把 half 都转成了 half2，使用标准库中的函数 __hadd2 运算。网络结构如下：

从图中可以看出，block_size 为 embedding_dim 的一半，这是因为用了 half2 这个数据结构，实际上每个线程处理了 2 个元素，所以线程数量缩减一半。核函数内部逻辑分为 2 个部分：1、求出 tensor 中对应元素的索引。2、一次对 Q、K、V 进行 add 和 transpose 操作。

template <>
__global__
void add_QKV_bias(__half* Q, const __half* bias_Q, __half* K, const __half* bias_K, __half* V, const __half* bias_V,
  __half* q_buf_, __half* k_buf_, __half* v_buf_,
  const int batch_size, const int seq_len, const int head_num, const int size_per_head, const int word_per_block)
{
  int tid = blockIdx.x * blockDim.x + threadIdx.x;
  int batch_id = tid / (head_num * seq_len * size_per_head);
  int seq_id = (tid % (head_num * seq_len * size_per_head)) / (head_num * size_per_head);
  int head_id = (tid % (head_num * size_per_head)) / size_per_head;
  // id in head
  int id = tid % size_per_head;
  // from [batch_size, seq_len, head_num, size_per_head] tanspose to [bacth_size, head_num, seq_len, size_per_head]
  int target_id = target_index(batch_id, seq_id, head_id, id, batch_size, seq_len, head_num, size_per_head);

  int bias_id = threadIdx.x;

  half2* src_ptr = (half2*)Q;
  half2* dst_ptr = (half2*)q_buf_;
  const half2* bias_ptr = (const half2*)bias_Q;
  dst_ptr[target_id] = __hadd2(src_ptr[tid],  __ldg(&bias_ptr[bias_id]));

  src_ptr = (half2*)K;
  dst_ptr = (half2*)k_buf_;
  bias_ptr = (const half2*)bias_K;
  dst_ptr[target_id] = __hadd2(src_ptr[tid],  __ldg(&bias_ptr[bias_id]));

  src_ptr = (half2*)V;
  dst_ptr = (half2*)v_buf_;
  bias_ptr = (const half2*)bias_V;
  dst_ptr[target_id] = __hadd2(src_ptr[tid],  __ldg(&bias_ptr[bias_id]));
}

4.2.2 计算 attention scores

先来看一下 attention 的计算公式，定义如下：

其中，，也就是说这一步要解决的是一个矩阵计算，用 tensorflow 代码表示如下：

scores = tf.matmul(query, key, transpose_b=True)

针对矩阵运算，源码中直接调用了 cuBLAS API，具体代码如下：

DataType_ alpha = (DataType_)1.0f, beta = (DataType_)0.0f;
// 计算 q * k^T
check_cuda_error(cublasGemmStridedBatchedEx(cublas_handle,
    CUBLAS_OP_T, CUBLAS_OP_N,
    seq_len, seq_len, size_per_head,
    &alpha,
    k_buf_, AType_, size_per_head, seq_len * size_per_head,
    q_buf_, BType_, size_per_head, seq_len * size_per_head,
    &beta,
    qk_buf_, CType_, seq_len, seq_len * seq_len,
    batch_size * head_num,
    computeType_,
    static_cast<cublasGemmAlgo_t>(cublasAlgo_[1])));

不熟悉 attention 的读者可能会问 attention scores 的具体含义是什么，笔者在早期的文章中有过介绍，其实就是两个矩阵的词向量两两相乘，向量相乘有什么含义？相似度，这个分数就代表 Q、K 的相似度。有兴趣的读者可以移步笔者之前的文章详细了解。（https://mp.weixin.qq.com/s/0zen3ItKmDLt5rTUbF37Mg）

4.2.3 softmax_kernel

拿到 Q、K 的相似度之后，直观上只要右乘一个 V 就可以得到 attention out，其含义就是一个加权平均的概念，既然要加权平均，必然要对权值进行归一化处理，这里的 softmax 就是这个作用。关于 softmax 核函数的实现方法笔者在前两篇文章也有介绍，OneFlow 官方给出了更为高效的实现方式，其高效的原因主要在访存带宽处理上，有兴趣的读者可以移步。【CUDA编程】OneFlow Softmax 算子源码解读之WarpSoftmax，【CUDA编程】OneFlow Softmax算子源码解读之BlockSoftmax

// 计算softmax(qk)
if(seq_len <= 32)
    block.x = 32;
else if(seq_len > 32 && seq_len <= 64)
    block.x = 64;
else if(seq_len > 64 && seq_len <= 128)
    block.x = 128;
else if(seq_len > 128 && seq_len <= 256)
    block.x = 256;
else if(seq_len > 256 && seq_len <= 512)
    block.x = 512;
else
    block.x = 1024;

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);
}

源码中核函数的 block_size 是根据 seq_len 确定的，取大于 seq_len 且为 32 的  的最小值。
另外在调用 softmax kernel 之前，会根据 batch_size * head_num 选择不同的 softmax kernel，主要是为了保证在大 batch 的情况下的计算效率，这里以 120 为阈值，应该是作者的经验数值。这里作者给出了 2 个 softmax kernel 的实现。
当 batch_size * head_num > 120 时，此时 batch 内元素较多，grid_size 取 batch_size * head_num，这时一个线程内处理一个 seq_len 的数据。

Software Kernel Grid Structure

/**
 * @brief
 *
 * @tparam T
 * @param qk_buf_                 [batch_size, head_num, seq_len, seq_len]
 * @param attr_mask               [batch_size, seq_len, seq_len]
 * @param batch_size
 * @param head_num
 * @param seq_len
 * @param scaler                  缩放因子
 * @return __global__
 */
template <typename T>
__global__
void softmax_kernel(T* qk_buf_, const T* attr_mask, const int batch_size, const int head_num, const int seq_len,
  const T scaler)
{
    // grid_size = batch_size * head_num
    int batch_id = blockIdx.x / head_num;
    // batch偏移量
    int qk_offset = blockIdx.x * seq_len * seq_len;
    int mask_offset = batch_id * seq_len * seq_len;

    __shared__ float s_sum, s_max;

    // 每次处理一个seq_len的数据
    for(int i = 0; i < seq_len; ++i)
    {
      float qk = threadIdx.x < seq_len ? (float)qk_buf_[threadIdx.x + qk_offset] : 0.0f;
      float mask_val = threadIdx.x < seq_len ? (float)attr_mask[threadIdx.x + mask_offset] : 0.0f;
      // 对于某些padded的word，给一个很小的值使其近似达到不参与运算的目的
      mask_val = (1.0f - mask_val) * -10000.0f;

      float tmp = threadIdx.x < seq_len ? (float)(qk * (float)scaler + mask_val): -1e20f;

      float max_val = blockReduceMax<float>(tmp);

      if(threadIdx.x == 0)
        s_max = max_val;
      __syncthreads();

      qk = threadIdx.x < seq_len ? __expf(tmp - s_max) : 0.0f;

      float sum_val = blockReduceSum<float>(qk);

      if(threadIdx.x == 0)
      {
        s_sum = sum_val + 1e-6f;
      }
      __syncthreads();

      if(threadIdx.x < seq_len)
        qk_buf_[threadIdx.x + qk_offset] = (T)(qk / s_sum);

      qk_offset += seq_len;
      mask_offset += seq_len;
    }
}

核函数内首先计算了每个元素偏移量，对于输入 tensor 而言，每个 block 处理 seq_len * seq_len 个数据，所以 block 内元素偏移量为 blockIdx.x * seq_len * seq_len，而对于 mask 矩阵而言，其维度为 [batch_size, seq_len, seq_len]，跟 head_num 无关，所以其偏移量为 batch_id * seq_len * seq_len。
接下来是一层循环，对于 seq_len * seq_len 矩阵而言，每个线程处理当前 thread_id 列的元素，每轮循环结束，处理该列下一行的元素。在每一轮循环中，所有的线程一起处理一行数据，首先拿到数据 qk 以及 mask_val。如果 mask_val 为 0，则给 mask_val 赋一个很小的值最后加在 qk 上使 qk 值很小，以致最终这个 softmax 分量趋于 0；如果 mask_val 为 1，则 mask 不干预后续计算。每个线程拿到处理后的 qk 值即 tmp 后，进行一次块内规约，即可求出当前行的最大值 max_val，然后为了避免指数运算导致数值溢出，让 tmp 减去 max_val 并求其指数值赋给 qk ，然后对 qk 再一次块内规约求出当前行的和 s_sum，最后让 qk 除以 s_sum 即可得到 softmax 值。核函数内要注意在两次块内规约后一定要进行一次块内同步，否则可能计算错误。

当 batch_size * head_num <= 120 时，此时 batch 较小，grid_size 取 batch_size * head_num * seq_len，这时一个线程块内处理一行数据，每个线程内只处理一个的数据。

template <typename T>
__global__
void softmax_kernel_v2(T* qk_buf_, const T* attr_mask, const int batch_size, const int head_num,
  const int seq_len, const float scaler)
{
    int batch_id = blockIdx.x / head_num / seq_len;
    int seq_id = blockIdx.x % seq_len;
    int qk_offset = blockIdx.x * seq_len;
    int mask_offset = batch_id * seq_len * seq_len + seq_id * seq_len;

    __shared__ float s_sum, s_max;

    float qk = threadIdx.x < seq_len ? (float)qk_buf_[threadIdx.x + qk_offset] : 0.0f;
    float mask_val = threadIdx.x < seq_len ? (float)attr_mask[threadIdx.x + mask_offset] : 0.0f;
      
    mask_val = (1.0f - mask_val) * -10000.0f;

    float tmp = threadIdx.x < seq_len ? (float)(qk * (float)scaler + mask_val) : -1e20f;
    float max_val = blockReduceMax<float>(tmp);
    if(threadIdx.x == 0)
      s_max = max_val;
    __syncthreads();

    float qk_tmp = threadIdx.x < seq_len ? __expf((float)(tmp - s_max)) : 0.0f;
    float sum_val = blockReduceSum<float>(qk_tmp);

    if(threadIdx.x == 0)
    {
      s_sum = sum_val + 1e-6f;
    }
    __syncthreads();

    if(threadIdx.x < seq_len)
      qk_buf_[threadIdx.x + qk_offset] = (T)(qk_tmp / s_sum);
}

这种情况下不涉及循环处理，计算逻辑与前面 softmax_kernel 循环体内部计算逻辑相同，不再赘述。

4.2.4 计算多头 attention out

这一步的意思就是使用 softmax 后的相似度矩阵右乘一个 V，得到多头注意力输出，注意这时候输出 tensor 的维度为 [batch_size, head_num, seq_len, size_per_head]。源码中直接调用了 cuBLAS API，具体代码如下：

// 计算qk * v
check_cuda_error(cublasGemmStridedBatchedEx(cublas_handle,
    CUBLAS_OP_N, CUBLAS_OP_N,
    size_per_head, seq_len, seq_len,
    &alpha,
    v_buf_, AType_, size_per_head, seq_len * size_per_head,
    qk_buf_, BType_, seq_len, seq_len * seq_len,
    &beta,
    transpose_dst_, CType_, size_per_head, seq_len * size_per_head,
    batch_size * head_num,
    computeType_,
    static_cast<cublasGemmAlgo_t>(cublasAlgo_[2])));

4.2.5 transpose

前面说过，多头 attention out 的维度为 [batch_size, head_num, seq_len, size_per_head]，此时这些 head 已经完成使命了，通过独立的 head_num 组 attention 参数计算出了 attention out，最后需要做的就是把这 head_num 组 attention out 拼接起来，体现在 tensor 上就是做一次 transpose，将维度变为 [batch_size, seq_len, head_num, size_per_head]。源码针对 fp16 和 fp32 分别提供了一个核函数 transpose，计算逻辑和 add_QKV_bias 中 transpose 计算逻辑相同，索引按顺序乘即可。具体代码如下：

template<typename T>
__global__
void transpose(T* src, T* dst, const int batch_size, const int seq_len, const int head_num, const int size_per_head)
{
  int batch_id = blockIdx.x / (head_num * seq_len);
  int seq_id = blockIdx.x % seq_len;
  int head_id = (blockIdx.x % (head_num * seq_len))/ seq_len;
  dst[batch_id * (head_num * seq_len * size_per_head) + seq_id * head_num * size_per_head
    + head_id * size_per_head + threadIdx.x] = src[blockIdx.x * size_per_head + threadIdx.x];
}

template<>
  __global__
void transpose(__half* src, __half* dst,
    const int batch_size, const int seq_len, const int head_num, const int size_per_head)
{
  int tid = blockIdx.x * blockDim.x + threadIdx.x;

  int batch_id = tid / (head_num * seq_len * size_per_head);
  int head_id = (tid % (head_num * seq_len * size_per_head)) / (seq_len * size_per_head);
  int seq_id = (tid % (seq_len * size_per_head)) / size_per_head;
  int id = tid % size_per_head;

  int target_id = target_index(batch_id, head_id, seq_id, id, batch_size, head_num, seq_len, size_per_head);
  half2* src_ptr = (half2*)src;
  half2* dst_ptr = (half2*)dst;

  dst_ptr[target_id] = src_ptr[tid];
}

5 BertEncoderTransformer

BertEncoderTransformer 类中有两个重要的成员方法：构造函数、forward 方法。其中构造函数内主要进行一些参数初始化功能，设备内存的申请和初始化也在该函数内进行。forward 方法内主要是核心逻辑的实现。

5.1 attention forward

根据调用链可知，BertEncoderTransformer->forward() 中第一步就是 attention_->forward()，其中 attention_ 对象在构造函数中被定义，attention_->forward() 执行的就是第 4 节的内容。

5.2 对 attention out 做线性变换

根据流程图和调用链可知，这一步是对多头注意力的输出 tensor 做一层线性变换，右乘一个参数矩阵，其实就是一个不加激活函数的 Dense 层，分为矩阵乘法和 add bias 两个操作步骤，这里调用了 cuBLAS API 实现矩阵乘法。

DataType_ alpha = (DataType_)1.0f;
DataType_ beta = (DataType_)0.0f;
int m = batch_size_ * from_seq_len_;
int k = head_num_ * size_per_head_;
int n = k;

check_cuda_error(cublasGemmEx(param_.cublas_handle,
    CUBLAS_OP_N, CUBLAS_OP_N,
    n, m, k,
    &alpha,
    param_.attr_output_kernel, AType_, n,
    attr_out_buf_, BType_, k,
    &beta,
    attr_matmul_buf_, CType_, n,
    computeType_,
    static_cast<cublasGemmAlgo_t>(cublasAlgo_[0])));

5.3 add_bias_input_layernorm_kernel

从核函数名字可以看出，这个核函数实现了 3 个操作：add bias、add input、layernomalization。其中 add bias 是完成上一步线性变换未完成的加偏置工作，add input 是 transformer 模型中的残差结构，layernomalization 则是层归一化操作。综合起来这个核函数的作用是：对线性变换后的 attention out 加偏置，然后加上原始输入 tensor 组成一个残差结构，最后进行一次层归一化变换。源码中针对 fp16 和 fp32 分别提供了一个核函数实现，计算逻辑都一样，这里只以 fp32 为例介绍。

Add_bias_input_layernorm kernel grid structure

/**
 * @brief                       grid_size = m, block_size = n
 *
 * @tparam T
 * @param out                   [batch_size, sql_len, latent_dim]
 * @param input                 [batch_size, sql_len, latent_dim]
 * @param bias                  [latent_dim,]
 * @param gamma
 * @param beta
 * @param m                     batch_size * seq_len
 * @param n                     latent_dim
 * @return __global__
 */
template <typename T>
__global__
void add_bias_input_layernorm(T* out, const T* input, const T* bias, const T* gamma, const T* beta, int m, int n)
{
  int tid = threadIdx.x;

  __shared__ float s_mean;
  __shared__ float s_variance;
  float mean =  0.0f;
  float variance = 0.0f;

  float local_out = 0.0f;
  // add，一个block处理一行
  for(int i = tid; i < n; i += blockDim.x)
    local_out += (float)(out[blockIdx.x * n + i] + input[blockIdx.x * n + i] + __ldg(&bias[i]));
  // mean_i = sum(x_i[j] for j in range(k)) / k
  mean = blockReduceSum<float>(local_out);
  if(threadIdx.x == 0)
    s_mean = mean / n;
  __syncthreads();
  // var_i = sum((x_i[j] - mean_i) ** 2 for j in range(k)) / k + epsilon
  variance = blockReduceSum<float>((local_out - s_mean) * (local_out - s_mean));
  if(threadIdx.x == 0)
    s_variance = variance / n + 1e-6f;
  __syncthreads();
  // x_i_normalized = (x_i - mean_i) / sqrt(var_i)
  // output_i = x_i_normalized * gamma + beta
  for(int i = tid; i < n; i += blockDim.x)
    out[blockIdx.x * n + i] =
                   (T)(((local_out - s_mean) * rsqrtf(s_variance)) * (float)(__ldg(&gamma[i])) + (float)(__ldg(&beta[i])));
}

如示意图所示，核函数中每个 block 处理一行数据，共 latent_dim = head_num * size_per_head 个元素，核函数中首先计算了 add bias、add input 两个操作，并将计算结果存储在寄存器变量 local_out 中。
接下来就是标准的 layerNormalization 操作，我们先来看一下 layerNormalization 的操作步骤，参照一下 tensorflow 框架 API 文档。

For each sample x_i in inputs with k features, we compute the mean and variance of the sample:
mean_i = sum(x_i[j] for j in range(k)) / k
var_i = sum((x_i[j] - mean_i) ** 2 for j in range(k)) / k

and then compute a normalized x_i_normalized, including a small factor epsilon for numerical stability.
x_i_normalized = (x_i - mean_i) / sqrt(var_i + epsilon)

And finally x_i_normalized is linearly transformed by gamma and beta, which are learned parameters:
output_i = x_i_normalized * gamma + beta
gamma and beta will span the axes of inputs specified in axis, and this part of the inputs' shape must be fully defined.

具体地，第一步计算均值和方差，核函数中使用块内规约计算出均值 s_mean 存储在共享内存中，所有块内线程都可以访问。然后根据 s_mean 和线程内的 local_out 以及 epsilon 系数再进行一次块内规约计算出方差 s_variance 存储在共享内存中。
第二步进行归一化和线性变换，对应 tensorflow API 的二、三步，直接计算即可，没有其他技巧，公式如下：

5.4 FeedForward 结构

根据 Transformer 模型结构，多头注意力之后为了增强表达能力，加了一个 FeedForward 层，该结构内部就是两个 Dense 层，第一层 Dense 中使用了激活函数，第二层没有激活函数。所以 FeedForward 层中包含了 5 个操作：矩阵乘法、add bias、activation、矩阵乘法、add bias。

5.4.1 attention out * inter kernel

FeedForward 层第一次线性变换会扩展 tensor 的最后一个维度的长度，源码中将 latent_dim（也就是 n）扩展为原来的 4 倍，所以这里的 inter kernel 的形状为 [latent_dim, 4 * latent_dim]，矩阵运算后的输出 tensor 形状为 [batch_size, seq_len, 4 * latent_dim]。

n *= 4;
check_cuda_error(cublasGemmEx(param_.cublas_handle,
    CUBLAS_OP_N, CUBLAS_OP_N,
    n, m, k,
    &alpha,
    param_.inter_kernel, AType_, n,
    attr_matmul_buf_, BType_, k,
    &beta,
    inter_matmul_buf_, CType_, n,
    computeType_,
    static_cast<cublasGemmAlgo_t>(cublasAlgo_[1])));

5.4.2 add_bias_act_kernel

顾名思义，add_bias_act_kernel 核函数包含了 add bias 和 activation 两个操作。源码中 block_size = n / 4 实际就是 latent_dim，为什么不直接取上一次运算后的矩阵宽度 n = 4 * latent_dim 呢？这里是希望一行元素（4 * latent_dim）能在一个 block 内处理，如果 block_size 直接取 n = 4 * latent_dim，可能超过 1024，因此还取 latent_dim，线程内循环 4 次处理即可。同样，源码中针对 grid_size 也取了 m / 4，在线程中通过循环每次跨 m / 4 步长处理 4 行数据。

template <typename T>
__global__
void add_bias_act(T* out, const T* bias, int m, int n)
{
  T val, reg_bias;

  int row_id = blockIdx.x;
  int ite = n / blockDim.x;
  int tid = threadIdx.x;

  for(int i = 0; i < ite; ++i)
  {
    reg_bias = __ldg(&bias[i * blockDim.x + tid]);
    row_id = blockIdx.x;

    while(row_id < m){
      val = out[tid + i * blockDim.x + row_id * n]+ reg_bias;
      out[tid + i * blockDim.x + row_id * n] = gelu<T>(val);
      row_id += gridDim.x;
    }
  }
}

核函数中先对列进行循环，ite = 4，从全局内存读出当前列的 bias，然后针对行进行循环，步长为 m / 4，循环体内部对当前行当前列的元素进行 add bias 和 gelu 操作，这里gelu 操作是一个简单的 element-wise 操作，比较简单不再介绍。
笔者点评：这里笔者私以为没有必要 grid_size 也取 m / 4，cuda 本身对线程块的数量没有限制，完全可以直接取 m，每次每个线程只处理一行数据，一方面可以增加并行程度，另一方面代码可阅读性也更好。笔者给出代码如下，亲测可用。

dim3 grid(m);
dim3 block(n / 4);
assert(block.x <= 1024);
add_bias_act_v2<T><<<grid, block, 0, stream>>>(out, bias, m, n);

template <typename T>
__global__
void add_bias_act_v2(T* out, const T* bias, int m, int n) {
  T val, reg_bias;

  int row_id = blockIdx.x;
  int ite = n / blockDim.x;
  int tid = threadIdx.x;

  for(int i = 0; i < ite; ++i) {
    reg_bias = __ldg(&bias[i * blockDim.x + tid]);
    val = out[tid + i * blockDim.x + row_id * n]+ reg_bias;
    out[tid + i * blockDim.x + row_id * n] = gelu<T>(val);
  }
}

5.4.3 inter out * out kernel

FeedForward 层第二次线性变换将 tensor 的最后一个维度的长度转换为原始大小，源码中将 n 重新赋值为 latent_dim，所以这里的 out kernel 的形状为 [4 * latent_dim, latent_dim]，矩阵运算后的输出 tensor 形状为 [batch_size, seq_len, latent_dim]。

n = k;
k *= 4;
check_cuda_error(cublasGemmEx(param_.cublas_handle,
    CUBLAS_OP_N, CUBLAS_OP_N,
    n, m, k,
    &alpha,
    param_.output_kernel, AType_, n,
    inter_matmul_buf_, BType_, k,
    &beta,
    param_.transformer_out, CType_, n,
    computeType_,
    static_cast<cublasGemmAlgo_t>(cublasAlgo_[2])));

5.5 add_bias_input_layernorm_kernel

这个核函数的计算逻辑在 5.3 中已经介绍过了，包含加偏置项、残差结构、层归一化三个操作，不再赘述。

6 小结

至此，Transformer encoder 前向计算的 14 个操作优化技巧已介绍完毕。总结如下：

针对半精度 fp16 的优化方面。首先，在 kernel 的实现中，将输入的 half 指针转成 half2 类型，并使用了 half2 相关的数学函数。这样不仅仅可以达到 2 倍于 half 的访存带宽和计算吞吐，还可以极大地减少指令的发射数量。其次，在 softmax 以及 layerNormalization 的操作中，为防止求和溢出，将数据以 half2 的形式读入后，会转成 float2 类型，来做求和计算，这里就非常细致了，尽可能地保障了较高精度，值得学习借鉴。

 

针对访存带宽方面，笔者以为除 fp16 以外其它数据类型也可以进一步优化，比如可以自定义 pack 类型进行合并读写，尽量把带宽打满。

 

针对线程网络结构方面，源码中基本使用一个 block 处理一行数据的模式进行设计，这里笔者私以为针对 seq_len 和 latent_dim 已知比较小的情况下（不超过1024），完全可以一个线程束处理一行数据，束内同步的开销远小于块内同步。当然，这个要求确实有些苛刻了。

 

源码中提供了一个块内规约的代码，思路非常好，值得初学 cuda 的读者品读。

 

 

 

参考文献链接

https://mp.weixin.qq.com/s/Z74UUMXHBnkUMePzKyNm7w

https://mp.weixin.qq.com/s/VzDCfKUMZBB_cO7uLjOUig