添加自定义 aclnn 算子¶
本文档介绍如何向 vllm-ascend 添加自定义 aclnn 算子。
自定义 aclnn 算子在 vllm-ascend 中如何工作?¶
自定义 aclnn 算子在 vllm-ascend 构建过程中会被编译并安装到 vllm_ascend/cann_ops_custom 目录。然后,aclnn 算子会被绑定到 torch.ops._C_ascend 模块,使用户能够在 vllm-ascend 的 Python 代码中调用它们。
要启用自定义算子,请使用以下代码:
如何添加自定义 aclnn 算子?¶
- 在
csrc目录下创建一个新的算子文件夹。 - 创建
op_host和op_kernel目录,分别存放主机端和内核端源代码。 - 在
csrc/build_aclnn.sh中为支持的 SOC 添加编译选项。注意,多个算子应使用;分隔,例如:CUSTOM_OPS="op1;op2;op3"。 - 在
csrc/torch_binding.cpp中将 aclnn 算子绑定到csrc/torch_binding_meta.cpp模块。 - 在