添加自定义 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 算子？

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

• 为宿主端和内核源代码创建 op_host 和 op_kernel 目录。

• Add build options in csrc/build_aclnn.sh for supported SOC. Note that multiple ops should be separated with ;, i.e. CUSTOM_OPS="op1;op2;op3".

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

• 在 csrc/torch_binding_meta.cpp 中为算子编写一个元实现，以便其能被捕获到 aclgraph 中。

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