添加自定义 aclnn 操作

本文档描述了如何向 vllm-ascend 添加自定义 aclnn 操作。

vllm-ascend 中的自定义 aclnn 操作是如何工作的？

自定义 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 操作？

• 在 csrc 目录下创建一个新的操作文件夹

• 为 host 和 kernel 源代码创建 op_host 和 op_kernel 目录

• 在支持的 SOC 的 csrc/build_aclnn.sh 中添加构建选项。注意，多个操作应该用 ; 分隔，即 CUSTOM_OPS=op1;op2;op3

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

• 在 csrc/torch_binding_meta.cpp 中为将被捕获到 aclgraph 中的 op 编写元实现

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