添加自定义 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 添加构建选项。注意，多个 op 应使用 ; 分隔，例如 CUSTOM_OPS="op1;op2;op3"。
4. 在 csrc/torch_binding.cpp 中将 aclnn 算子绑定到 torch.ops._C_ascend 模块。
5. 在 csrc/torch_binding_meta.cpp 中为该 op 编写元实现，以便其被捕获到 aclgraph 中。

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