最近在使用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
