CANN PyTorch自定义算子扩展

概述【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer本工程目标是基于torch原生支持的cpp extension方式提供用户自定义算子注册eager模式和torch.compile模式执行的kernel 实现。关于eager模式与torch.compile模式的介绍参考Link。编译与安装安装依赖编译自定义算子包之前请先完成 torch 包和torch_npu包的下载与安装。 取包地址xxxx源码下载克隆代码仓并更新三方库git clone https://gitcode.com/cann/cann-recipes-infer.git git submodule update --init --recursive编译安装源码编译与安装cd xxxxxx bash build_and_install.sh执行之后会完成源码编译与自定安装。编译完成之后会在当前路径新生成一个dist文件夹存放新编译的wheel包。快速上手用例执行执行examples路径下的测试用例验证自定义算子的注册和执行bash run.sh目录结构介绍主体功能代码在custom_pypto路径编译工程在setup.py, 测试用例在test路径。├── custom_pypto │ ├── __init__.py #custom_pypto包初始化文件用于导入自定义算子python实现和编译的so │ ├── csrc #自定义算子适配层c代码目录 │ │ ├── add_custom_autograd.cpp # 自定义add算子适配代码以及实现注册 │ │ ├── npu_lightning_indexer_pto.cpp # 自定义算子lightning indexer适配代码以及实现注册 │ │ ├── npu_sparse_attention_pto.cpp # 自定义算子sparse attention适配代码以及实现注册 │ │ ├── npu_lightning_indexer_prolog_pto.cpp # 自定义算子lightning indexer prolog适配代码以及实现注册 │ │ ├── ops_def_registration.cpp # 自定义算子新增定义 │ │ ├── ops_common.cpp # 自定义算子调用和下发框架实现 │ │ ├── ops_common.h # 自定义算子调用和下发框架公共接口定义 │ │ └── ops_function.h # 正反向接口定义头文件 │ └── converter # 自定义算子包python侧converter代码 │ │ ├── __init__.py # python初始化文件用于converter注册 │ │ ├── add_custom.py # 自定义add算子的converter注册 │ │ ├── npu_lightning_indexer_pto.py # 自定义lightning indexer算子的converter注册 │ │ ├── npu_sparse_attention_pto.py # 自定义sparse attention算子的converter注册 │ │ ├── npu_lightning_indexer_prolog_pto.py # 自定义lightning indexer prolog算子的converter注册 │ ├──setup.py # wheel包编译文件 │ ├──README.md #当前自定义算子仓库的使用指导 │ ├── build_and_install.sh # 自定义算子wheel包编译与安装脚本新增自定义算子指导基于上述代码目录结构介绍梳理新增自定义算子步骤step0新增pytorch算子前准备工作C extensions插件提供了将自定义算子映射到昇腾AI处理器的功能为使用PyTorch框架的开发者提供了便捷的NPU算子库调用能力基于PyTorch原生提供的自定义算子扩展功能用户可以编译安装自定义算子wheel包并运行。 本指导将以lightning_indexer_pto算子为例基于C extensions的方式介绍如何说明在昇腾NPU上完成自定义算子开发和适配。step1在算子适配层c代码目录csrc中完成C侧自定义算子代码适配注册自定义算子schemaPyTorch提供TORCH_LIBRARY宏来定义新的命名空间并在该命名空间里注册schema。注意命名空间的名字必须是唯一的。具体示例如下-- custom_pypto/csrc/ops_def_registration.cpp中添加如下代码 #include torch/library.h TORCH_LIBRARY_IMPL(custom_pypto, Meta, m) { m.impl(npu_lightning_indexer_pto, custom_pypto::npu_lightning_indexer_pto_meta); }注册Meta函数通过PyTorch的Meta后端帮助算子完成入图时所需要的shape和data type推导。具体示例如下-- 创建custom_pypto/csrc/npu_lightning_indexer_pto.cpp #include torch/library.h #include ops_common.h namespace custom_pypto { using namespace at_npu::native; // npu tensor max size const int SIZE 8; const int DIM_0 0; const int DIM_1 1; const int DIM_2 2; const int DIM_3 3; // 推导输出shape和data type at::Tensor construct_lightning_indexer_pto_output_tensor( const at::Tensor query, const at::Tensor key, const at::Tensor weights, const c10::optionalat::Tensor actual_seq_lengths_key, const c10::optionalat::Tensor block_table) { at::SmallVectorint64_t, SIZE output_size; int sparse_count 2048; output_size {query.size(DIM_0), query.size(DIM_1), key.size(DIM_2), sparse_count}; // BSND at::Tensor output at::empty(output_size, query.options().dtype(at::kInt)); return output; } // 为META设备实现前向接口 at::Tensor npu_lightning_indexer_pto_meta( const at::Tensor query, const at::Tensor key, const at::Tensor weights, const c10::optionalat::Tensor actual_seq_lengths_query, const c10::optionalat::Tensor actual_seq_lengths_key, const c10::optionalat::Tensor block_table, c10::string_view layout_query, c10::string_view layout_key, int64_t sparse_count, int64_t sparse_mode) { return construct_lightning_indexer_pto_output_tensor(query, key, weights, actual_seq_lengths_key, block_table); } } // 为Meta设备注册前向实现 TORCH_LIBRARY_IMPL(custom_pypto, Meta, m) { m.impl(npu_lightning_indexer_pto, custom_pypto::npu_lightning_indexer_pto_meta); }step2在算子适配层python代码目录converter中完成python侧自定义算子代码适配注册自定义算子converter完成自定义算子在NPU 设备上执行的功能接口内部会调用自定义算子的 C API如-- 创建custom_pypto/converter/npu_sparse_attention_pto.py from typing import ( Any, Callable, ContextManager, Iterator, List, Literal, NamedTuple, Optional, Sequence, Tuple, TypeVar, Union, overload, ) import torch import torch_npu import torchair from torch import Generator, contiguous_format, inf, strided, SymInt from torch.types import Device, Number, _bool, _complex, _device, _dtype, _float, _int, _layout, _qscheme, _size from torchair._ge_concrete_graph import ge_apis as ge from torchair._ge_concrete_graph.fx2ge_converter import declare_supported, register_fx_node_ge_converter from torchair.ge._ge_graph import Tensor, TensorSpec, DataType from torchair._ge_concrete_graph.supported_declaration import _TypedTensor, F32, F16, F64, I32, I16, I64, I8, U8, \ BOOL, Support from torchair._ge_concrete_graph.utils import dtype_promote from torchair.ge import attr # 注意命名空间和schema名称需要与前面注册保持一致 register_fx_node_ge_converter(torch.ops.custom_pypto.npu_lightning_indexer_pto.default) def convert_npu_lightning_indexer_pto( query: Tensor, key: Tensor, weights: Tensor, *, actual_seq_lengths_query: Tensor None, actual_seq_lengths_key: Tensor None, block_table: Tensor None, layout_query: str BSND, layout_key: str PA_BSND, sparse_count: int 2048, sparse_mode: int 3, meta_outputs: Any None ): NB: npu_lightning_indexer_pto(Tensor query, Tensor key, Tensor weights, Tensor actual_seq_lengths_key, *, Tensor? block_tableNone) - Tensor return torchair.ge.custom_op( LightningIndexerPto, # 和REG_OP中算子原型名称保持一致例如AddCustom inputs{x0: query, x1: key, x2: weights, x3: actual_seq_lengths_key, x4: block_table, }, outputs[y0] )更新说明时间更新事项2025/09/22完成readme初始版本【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考