梳理cuda算子编译与python调用的流程_以vllm为例

最近在使用vllm框架，部署大语言模型的时候。发现吞吐量提升比较明显。对里面用到的技术比较感兴趣。后来发现vllm使用了一些新的技术，如kv cache，page attention等。其中很多是用cuda编写加速的。并且对cuda算子如何应用到python服务中比较感兴趣，现在就自己的了解，用文章说明一下。

下面就以paged_attention_v2。这个算子为例进行说明一下。

1.cuda算子代码实现

./vllm/csrc/attention/attention_kernels.cu 中

2.python中如何调用cuda算子

./vllm/vllm/model_executor/layers/attention.py 中：

from vllm._C import ops                # 注意这里的包的名字：vllm._C。后面会补充说明怎么来的。
from vllm._C import cache_ops

# 省略部分

ops.paged_attention_v2(
    output,
    exp_sums,
    max_logits,
    tmp_output,
    query,
    key_cache,
    value_cache,
    num_kv_heads,
    scale,
    input_metadata.block_tables,
    input_metadata.context_lens,
    block_size,
    input_metadata.max_context_len,
    alibi_slopes,
)

看了cuda算子的实现，以及python中如何调用cuda算子。现在最大的问题就是：一个是python代码，一个是cuda代码，如何能调用成功呢？

这个就需要先编译。具体就要用到steup.py来编译。

3.steup.py

基于预编译的扩展由于需要编译，而setup.py文件正是基于setuptools的编译脚本。因此一个 Python package 的扩展是可以在setup.py文件中找到其蛛丝马迹的。这里我们截取一段 vllm 的 setup.py 文件。

这里可以看到 setup 函数中一个主要的参数 ext_modules，该参数需要指定为一个 Extension 列表：ext_modules，代表实际需要编译的扩展。

ext_modules的相关代码如下：

ext_modules 里的元素的类是 CUDAExtension。

这里补充说一下：生成扩展的函数会随系统环境不同而有所区别。例如：当系统中没有 CUDA 时会调用 CppExtension，且只编译所有 .cpp文件，反之则调用 CUDAExtension。其实 CppExtension 与 CUDAExtension 都是基于setuptools.Extension的扩展，这两个函数都额外将系统目录中的 torch/include 加入到 C++ 编译时的include_dirs中，另外 CUDAExtension 会额外将CUDA相关的库以及头文件加到默认编译搜索路径中。 

在上述代码中我们终于看到了vllm._C，该名字正是新定义的扩展的名字。由此我们便知道上文 python 中的：from vllm._C import ops，实际上是在 setup.py 文件中指定其模块名字的。 

4.执行编译

python setup.py sdist bdist_wheel

部分编译的日志，如下：

[1/10] /usr/local/cuda/bin/nvcc  -I/root/miniconda3/envs/vllm/lib/python3.10/site-packages/torch/include -I/root/miniconda3/envs/vllm/lib/python3.10/site-packages/torch/include/torch/csrc/api/include -I/root/miniconda3/envs/vllm/lib/python3.10/site-packages/torch/include/TH -I/root/miniconda3/envs/vllm/lib/python3.10/site-packages/torch/include/THC -I/usr/local/cuda/include -I/root/miniconda3/envs/vllm/include/python3.10 -c -c /data/caiyueliang/vllm/csrc/cuda_utils_kernels.cu -o /data/caiyueliang/vllm/build/temp.linux-x86_64-cpython-310/csrc/cuda_utils_kernels.o -D__CUDA_NO_HALF_OPERATORS__ -D__CUDA_NO_HALF_CONVERSIONS__ -D__CUDA_NO_BFLOAT16_CONVERSIONS__ -D__CUDA_NO_HALF2_OPERATORS__ --expt-relaxed-constexpr --compiler-options ''"'"'-fPIC'"'"'' -O2 -std=c++17 -D_GLIBCXX_USE_CXX11_ABI=0 -gencode arch=compute_86,code=sm_86 --threads 8 -DTORCH_API_INCLUDE_EXTENSION_H '-DPYBIND11_COMPILER_TYPE="_gcc"' '-DPYBIND11_STDLIB="_libstdcpp"' '-DPYBIND11_BUILD_ABI="_cxxabi1011"' -DTORCH_EXTENSION_NAME=_C -D_GLIBCXX_USE_CXX11_ABI=0

参数说明：

1. -I: 添加包含（Include）目录到搜索路径。这些路径是编译器查找头文件的位置。在你的命令中，它们指向了PyTorch库的头文件和CUDA头文件的位置。

/usr/local/cuda/include、/root/miniconda3/envs/vllm/lib/python3.10/site-packages/torch/include、/root/miniconda3/envs/vllm/include/python3.10

1. -c : 指示编译器进行编译和汇编，但不链接。 指定要编译的CUDA源文件。在你的命令中，这个文件是/data/caiyueliang/vllm/csrc/cuda_utils_kernels.cu。
2. -o : 指定输出的目标文件。在这个例子中，编译后的输出文件是/data/caiyueliang/vllm/build/temp.linux-x86_64-cpython-310/csrc/cuda_utils_kernels.o。
3. -D: 定义宏。这些选项用于定义预处理器宏，常用于条件编译。

-DTORCH_API_INCLUDE_EXTENSION_H、-DTORCH_EXTENSION_NAME=_ext、-D_GLIBCXX_USE_CXX11_ABI=0: 这些是针对PyTorch扩展的特定宏定义。

1. --expt-relaxed-constexpr: 允许在CUDA设备代码中使用更宽松的constexpr。
2. --compiler-options '-fPIC': 指定额外的编译器选项。-fPIC表示生成位置无关代码（与目录无关），这对于动态链接库是必要的。
3. -O2: 指定优化级别。-O2表示编译器应使用第二级优化。
4. -std=c++17: 使用C++17标准进行编译。
5. -D_GLIBCXX_USE_CXX11_ABI=0: 控制对GCC 5及更高版本中的C++11 ABI的使用。
6. -gencode arch=compute_86,code=sm_86: 指定CUDA代码的生成目标。这里是为计算能力8.6的设备生成代码。
7. --threads 8: 指定编译时使用的线程数。
8. -DTORCH_API_INCLUDE_EXTENSION_H: 可能是为PyTorch扩展定义的特定宏。
9. -DPYBIND11_COMPILER_TYPE, -DPYBIND11_STDLIB, -DPYBIND11_BUILD_ABI: 这些是Pybind11相关的宏定义，用于控制与Python绑定的一些方面。
10. -DTORCH_EXTENSION_NAME=_C: 定义一个用于PyTorch扩展的宏。将TORCH_EXTENSION_NAME宏赋值为：_C。则才可以使用from vllm._C import ops。
11. -D_GLIBCXX_USE_CXX11_ABI=0: 这个宏控制是否使用新的C++11 ABI。0表示使用旧的ABI。

这条命令的主要目的是编译一个CUDA源文件，其中包含了多个用于指定编译配置、处理PyTorch扩展、控制ABI兼容性等的参数。这种复杂的命令行通常用于在具有特定要求的开发环境中编译CUDA代码。

编译目录的内容：

./vllm/build/temp.linux-x86_64-cpython-310/csrc/ 这个目录下的输出文件：

g++ -pthread -B /root/miniconda3/envs/vllm/compiler_compat -shared -Wl,-rpath,/root/miniconda3/envs/vllm/lib -Wl,-rpath-link,/root/miniconda3/envs/vllm/lib -L/root/miniconda3/envs/vllm/lib -Wl,-rpath,/root/miniconda3/envs/vllm/lib -Wl,-rpath-link,/root/miniconda3/envs/vllm/lib -L/root/miniconda3/envs/vllm/lib /data/caiyueliang/vllm/build/temp.linux-x86_64-cpython-310/csrc/activation_kernels.o /data/caiyueliang/vllm/build/temp.linux-x86_64-cpython-310/csrc/attention/attention_kernels.o /data/caiyueliang/vllm/build/temp.linux-x86_64-cpython-310/csrc/cache_kernels.o /data/caiyueliang/vllm/build/temp.linux-x86_64-cpython-310/csrc/cuda_utils_kernels.o /data/caiyueliang/vllm/build/temp.linux-x86_64-cpython-310/csrc/layernorm_kernels.o /data/caiyueliang/vllm/build/temp.linux-x86_64-cpython-310/csrc/pos_encoding_kernels.o /data/caiyueliang/vllm/build/temp.linux-x86_64-cpython-310/csrc/pybind.o /data/caiyueliang/vllm/build/temp.linux-x86_64-cpython-310/csrc/quantization/awq/gemm_kernels.o /data/caiyueliang/vllm/build/temp.linux-x86_64-cpython-310/csrc/quantization/gptq/q_gemm.o /data/caiyueliang/vllm/build/temp.linux-x86_64-cpython-310/csrc/quantization/squeezellm/quant_cuda_kernel.o -L/root/miniconda3/envs/vllm/lib/python3.10/site-packages/torch/lib -L/usr/local/cuda/lib64 -lc10 -ltorch -ltorch_cpu -ltorch_python -lcudart -lc10_cuda -ltorch_cuda -o build/lib.linux-x86_64-cpython-310/vllm/_C.cpython-310-x86_64-linux-gnu.so

copying build/lib.linux-x86_64-cpython-310/vllm/_C.cpython-310-x86_64-linux-gnu.so -> build/bdist.linux-x86_64/wheel/vllm

在./build/lib.linux-x86_64-cpython-310/vllm/ 这个目录中，会编译出一个_C.xxx.so的库：

5.PYBIND11_MODULE

上面说到通过 setup.py 我们编译了扩展文件。可是目前仍然有个疑问，为什么编译出来的 C++ / CUDA 二进制文件可以在 Python 中直接被调用呢？

使用pybind 构建python和cpp的桥梁  pybind.cpp 在ops.h中具体做了函数声明。

这里PYBIND11_MODULE是一个宏，定义在 pybind11 库中(见https://link.zhihu.com/?target=https%3A//github.com/pybind/pybind11/blob/master/include/pybind11/pybind11.h)。而 pybind11 是一个用来在 C++ 代码中创建 Python的连接的库。找到了源头，我们进一步分析。

这里PYBIND11_MODULE 的作用是为 C++ 代码接入 Python 解释器提供入口。以上述代码为例， TORCH_EXTENSION_NAME 正是在上文 gcc编译过程中出现的宏，对应为extension的 name 变量。因此在这里会被解释成 _C（注意没有双引号） 。m 则代表 TORCH_EXTENSION_NAME 所对应的模块实例。ops 是 m 的子模块名字。{}中的每个 ops.def 都定义了一个 _C 的成员函数，其一般形式为 ops.def("函数名"，具体 C++ 实现的函数指针, "文档", 参数列表)。

通过这种形式，paged_attention_v2 也就顺利地成为了 vllm._C.ops 的成员函数。在 Python 中也就可以运行 from vllm._C import ops了。