添加自定义aclnn算子

本文档介绍如何向vllm-ascend添加自定义aclnn算子。

自定义aclnn算子在vllm-ascend中的工作原理？

自定义aclnn算子在vllm-ascend构建过程中会被编译并安装到vllm_ascend/cann_ops_custom目录。随后这些aclnn算子会被绑定到torch.ops._C_ascend模块，使用户能够在vllm-ascend的Python代码中调用它们。

要启用自定义算子，请使用以下代码：

from vllm_ascend.utils import enable_custom_op

enable_custom_op()

如何添加自定义aclnn算子？

1. 在csrc目录下创建新的算子文件夹。

2. 创建op_host和op_kernel目录，分别存放主机端和内核端源代码。

3. 在csrc/build_aclnn.sh中为支持的SOC添加构建选项。注意多个算子之间需用;分隔，例如CUSTOM_OPS="op1;op2;op3"。

4. 在csrc/torch_binding.cpp中将aclnn算子绑定到torch.ops._C_ascend模块。

5. 在csrc/torch_binding_meta.cpp中为算子编写元实现，使其能够被捕获到aclgraph中。

成功构建vllm-ascend后，即可在Python代码中调用自定义aclnn算子。