添加自定义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目录。在
csrc/build_aclnn.sh中为支持的SOC添加构建选项。注意多个算子应用;分隔，例如CUSTOM_OPS="op1;op2;op3"。在
csrc/torch_binding.cpp中将aclnn算子绑定到torch.ops._C_ascend模块。在
csrc/torch_binding_meta.cpp中为算子编写一个元实现，以便其能被捕获到aclgraph中。
成功构建vllm-ascend后，即可在Python代码中调用自定义的aclnn算子。