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

本文介绍了CUDA算子在vllm框架中的应用,从cuda算子代码实现到python调用的流程,包括steup.py的编译过程和PYBIND11_MODULE的作用,详细阐述了CUDA扩展如何与Python交互。
摘要由CSDN通过智能技术生成

最近在使用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 --/data/caiyueliang/vllm/csrc/cuda_utils_kernels.cu -/data/caiyueliang/vllm/build/temp.linux-x86_64-cpython-310/csrc/cuda_utils_kernels.-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./data/caiyueliang/vllm/build/temp.linux-x86_64-cpython-310/csrc/attention/attention_kernels./data/caiyueliang/vllm/build/temp.linux-x86_64-cpython-310/csrc/cache_kernels./data/caiyueliang/vllm/build/temp.linux-x86_64-cpython-310/csrc/cuda_utils_kernels./data/caiyueliang/vllm/build/temp.linux-x86_64-cpython-310/csrc/layernorm_kernels./data/caiyueliang/vllm/build/temp.linux-x86_64-cpython-310/csrc/pos_encoding_kernels./data/caiyueliang/vllm/build/temp.linux-x86_64-cpython-310/csrc/pybind./data/caiyueliang/vllm/build/temp.linux-x86_64-cpython-310/csrc/quantization/awq/gemm_kernels./data/caiyueliang/vllm/build/temp.linux-x86_64-cpython-310/csrc/quantization/gptq/q_gemm./data/caiyueliang/vllm/build/temp.linux-x86_64-cpython-310/csrc/quantization/squeezellm/quant_cuda_kernel.-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了。

  • 16
    点赞
  • 18
    收藏
    觉得还不错? 一键收藏
  • 0
    评论
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值