添加自定义 aclnn 算子

添加自定义 aclnn 算子#

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

自定义 aclnn 算子在 vllm-ascend 中是如何工作的？#

在构建 vllm-ascend 的过程中，自定义 aclnn 算子会被编译并安装到 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 算子？#

  • 在 csrc 目录下创建一个新的算子文件夹

  • 分别创建 op_host 和 op_kernel 目录，用于存放 Host 端和 Kernel 端的源代码

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

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

  • 在 csrc/torch_binding_meta.cpp 中为算子编写 meta 实现，以支持其被捕获到 aclgraph 中

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

目录