大模型自定义算子优化方案学习笔记：CUDA算子定义、算子编译、正反向梯度实现

01算子优化的意义

随着大模型应用的普及以及算力紧缺，下一步对于计算性能的追求一定是技术的核心方向。因为目前大模型的计算逻辑是由一个个独立的算子或者说OP正反向求导实现的，底层往往调用的是GPU提供的CUDA的驱动程序。如果不能对于整个计算过程学习并了解，对于性能优化领域无非是隔靴搔痒，今天也是抽一点时间读了下网上的一些文档和CUDA的文档，整理了学习材料。

首先说下为什么要自定义算子，无非是两个原因，

（1）TF、PyTorch等提供的原生算子不满足需求，需要通过底层接口，比如CUDA层更灵活的实现个性化算子

（2）目前提供的算子性能不足，没有很好的利用到GPU的并行计算优势，有优化空间

接着性能优化的问题说，因为GPU与CPU最大的区别是计算单元占据了绝大部分的体积（图中绿色部分），而控制单元较少。

自定义手写算子可以更好地利用绿色的计算单元的并行化优势。大的思路是Grid包含Block，Block包含Thread。于是首先算子需要把计算逻辑拆分成Thread，让程序可以并行化的运行起来，然后有机的管理各个Block的执行节奏，解决好异步和同步问题，就可以让芯片的计算效率最大化。

02实现流程

整个自定义算子优化其实可以分为CUDA算子定义、算子编译、正方向梯度实现几个步骤。

1、CUDA算子定义

需要定义以下几个文件：

（1）function.cpp：python层和CUDA层的衔接，实现手写的C++ CUDA算子可以被python代码调用

（2）function.h：CUDA函数声明文件

（3）function.cu：CUDA函数的逻辑实现

这里比较核心的文件就是.cu文件，构建的时候主要做两个事：一个是建设Kernel函数，因为只有Kernel函数是在GPU端执行，执行完之后要将控制权给到控制函数，这里要控制好异步、同步的问题。二是在kernel函数中需要通过循环函数定义每个Thread以及每个Block的工作，真正让计算并行化

.cpp文件可以通过pybind函数实现，这个函数主要解决的是C++代码和Python绑定的问题，项目地址：GitHub - pybind/pybind11: Seamless operability between C++11 and Python

2.编译和执行

import torch
from torch.utils.cpp_extension import load
cuda_module = load(name="function",
                   extra_include_paths=["include"],
                   sources=["function.cpp", "function.cu"],
                   verbose=True)

接着就是执行过程中的编译，通过load函数底层会执行JIT（Just in time）的动态编译模式调用.cpp和.cu文件，底层运行的是Ninjia编译器，通过调用nvcc编译.so文件

[1/2] nvcc -c function.cu -o function.cuda.o
[2/3] c++ -c function.cpp -o function.o
[3/3] c++ function.o function.cuda.o -shared -o function.so

3.正反向梯度实现

以上两步实现了自定义算子的逻辑，可以通过手写CUDA算子并在python框架层调用，如果要满足建模需求，还需要实现正方向梯度。具体的做法是在建模的函数中实现forward和backward函数，定义导数作为输出。

以上大概就是手写算子优化的简单流程，仅当学习笔记。

参考:

【1】熬了几个通宵，我写了份CUDA新手入门代码 - 知乎

【2】CUDA C++ Programming Guide