CANN ops-transformer torch_extension 开发规范：PyTorch 风格 NPU 算子 API 的完整实现指南

算子库人工智能深度学习Ascend【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址https://gitcode.com/cann/ops-transformer点击查看免费下载导读本文是 CANN ops-transformer 仓库中 torch_extension_guidelines.md 的深度展开版系统讲解如何为仓库新增/修改一个 PyTorch 风格的 NPU 算子 API。cann_ops_transformer通过 PyTorch JITtorch.utils.cpp_extension.load在首次调用时即时编译 C Kernel Wrapper把 PyTorch 函数接口桥接到 CANN 的 aclnn 接口同时通过 GE Converter 支持 torchair 图模式。读完本文你将掌握算子 API 的目录组织、跨层命名约定、C/Python/图模式三层实现要点、文档编写规范并能以flash_attn为模板独立新增一个算子 API。1. 架构总览一个算子 API 的四个组成部分一个完整的算子 API 由「Python 前端、C 后端、torchair 图模式 Converter、文档」四部分组成Python 前端${op_api}.py负责 JIT 编译管理、schema/meta 注册、PrivateUse1 dispatcher 注册与对外函数封装C 后端csrc/${op_api}.cpp通过ACLNN_CMD宏把 PyTorch 张量桥接到 CANN 的 aclnn C-API是真正的 kernel 拉起层图模式 Convertergraph_convert_${op_api}.py可选在 torchair 图模式下把 aten 算子转换为 GE 节点文档docs/torchapi_${op_api}.md对外 API 的中文接口文档。cann_ops_transformer的 JIT 编译与调度核心在 OpBuilder 基类 中实现模块加载时_ensure_initialized()完成延迟初始化解析ASCEND_HOME_PATH、定位 torch_npu 与 CANN 路径、注册 schema 与 Meta真正的编译由首次load()触发。load()内部调用torch.utils.cpp_extension.load并自动携带从 include_paths() 解析出的 CANN/torch_npu 头文件路径、cxx_args() 中的编译选项-O3、-fPIC、栈保护、-fvisibilityhidden等与 extra_ldflags() 中的链接选项-lascendcl、-ltorch_npu编译产物按算子名缓存到OpBuilder._loaded_ops避免重复编译。1.1 目录组织规范新增算子 api 涉及的文件布局以算子${op_api}为例源码位于ops-transformer/category/op/ops-transformer/ ├── category/op/ │ ├── op_host/ # 原有算子实现代码 │ ├── op_kernel/ # 原有 kernel 代码 │ ├── tests/ # 原有测试 │ └── torch_extension/ # torch_extension 文件 │ ├── __init__.py # 导出算子接口与 Converter │ ├── ${op_api}.py # Python 前端 │ ├── graph_convert_${op_api}.py # torchair 图模式 Converter可选 │ └── csrc/ │ └── ${op_api}.cpp # C 后端 ├── torch_extension/ │ ├── setup.py # 自动收集分布式算子文件 │ ├── cann_ops_transformer/ │ │ ├── __init__.py # 包根入口动态发现 命名空间导出 │ │ ├── op_builder/ │ │ │ ├── __init__.py # 导出 OpBuilder, get_as_library │ │ │ └── builder.py # OpBuilder 基类支持 category 参数、延迟初始化 │ │ ├── common/ │ │ │ ├── aclnn_common.h # ACLNN_CMD 宏、类型转换等公共能力 │ │ │ └── hccl_common.h # 通信类算子公共能力 │ │ ├── csrc/ │ │ │ └── extension.cpp # C stub │ │ ├── ops/ │ │ │ └── __init__.py # 自动发现算子目录扫描 entry point │ │ └── docs/ # 开发规范等接口文档放在算子 docs/新增文件请放在对应算子的category/op/torch_extension/目录下import 路径统一以cann_ops_transformer为根。1.2 新增算子 API 的标准动作清单以mc2/mega_moe为例该算子的 torch_extension 目录结构已在仓库中验证mc2/mega_moe/torch_extension/下含__init__.py、mega_moe.py、graph_convert_mega_moe.py与csrc/mega_moe.cpp在mc2/mega_moe/torch_extension/csrc/mega_moe.cpp中实现 C kernel wrapper调用ACLNN_CMD拉起 aclnn 接口在mc2/mega_moe/torch_extension/mega_moe.py中编写OpBuilder子类定义sources/schema/register_meta注册 dispatcher 实现并提供对外的 Python 函数在mc2/mega_moe/torch_extension/graph_convert_mega_moe.py中编写图模式 Converter若需支持图模式在mc2/mega_moe/torch_extension/__init__.py中定义__all__并导出算子接口与 Converter算子导入后自动被ops/__init__.py的自动发现机制加载无需手动注册在category/op/docs/torchapi_${op_api}.md中补充算子文档并同步更新 docs/zh/menu_torch_api.md 与 docs/zh/torch_api_list.md。第 5 步的自动发现有两套机制见 ops/init.py一是目录扫描_discover_ops_from_dir遍历ops/category/op/__init__.py二是entry point_discover_ops_from_entry_points读取cann_ops_transformer.ops分组由 setup.py 的TORCH_EXTENSION_OPS环境变量按需生成。此外还维护了一份_legacy_map旧名→新模块的兼容映射。2. 命名规范2.1 API 命名一个算子从 schema 注册到对外导出涉及多个层级的命名需保持一致且各司其职。对外 api 接口及算子名一律不带npu_前缀直接采用算子语义的小写蛇形名如flash_attn层级命名约定示例Library 名DEF 域固定为cann_ops_transformerget_as_library()创建的 Library 实例import 路径从cann_ops_transformer.op_builder导入from cann_ops_transformer.op_builder import OpBuilder, get_as_libraryC wrapper 函数名与 schema 算子名一致置于namespace op_api内op_api::flash_attnPYBIND11_MODULE导出名与 schema 算子名一致m.def(flash_attn, flash_attn, flash_attn);Meta 实现函数名schema 算子名 _meta后缀flash_attn_metaPrivateUse1 dispatcher 函数名下划线前缀 schema 算子名_flash_attnOpBuilder 子类名算子名的大驼峰 OpBuilder后缀内部专用可加_前缀FlashAttnOpBuilder、_FlashAttnOpBuilder对外 Python 接口名用户直接调用的函数名体现使用语义不带npu_前缀flash_attn图模式 GE op 函数名与 GE 算子op_type一致的大驼峰FlashAttentionScore图模式 Converter 函数名convert_ schema 算子名convert_flash_attn命名要点不带npu_前缀对外算子名与 api 接口统一使用算子语义名小写蛇形不加npu_等后端前缀schema 名、C 函数名、pybind 导出名三者必须与该名字完全一致否则 JIT 编译产物无法被正确调用。接口名体现语义对外函数名应贴近业务语义。无论是纯透传 aclnn 接口的算子如flash_attn还是封装了结构体构造、参数整理等额外逻辑的接口均采用语义化命名。aclnn 接口名独立底层 aclnn 接口沿用 CANN 命名大驼峰如aclnnFlashAttentionScore与对外算子名解耦C wrapper 内通过ACLNN_CMD(aclnnFlashAttentionScore, ...)调用。版本后缀同一算子的不同迭代版本以_v2、_v3等后缀区分schema 名、文件名、Converter 名需同步带上版本后缀如flash_attn_v2、graph_convert_flash_attn_v2.py。辅助/工具接口与主算子配套的工具函数采用动宾语义命名如get_flash_attn_workspace_size。2.2 文件命名统一使用小写蛇形命名snake_case单词以_连接禁止使用大写、驼峰或连字符。同一算子的各层文件主名保持一致放在category/op/torch_extension/下仅靠目录和前缀区分职责Python 前端category/op/torch_extension/${op_api}.py如flash_attn.py算子__init__.pycategory/op/torch_extension/__init__.py定义__all__并导出算子接口C 后端category/op/torch_extension/csrc/${op_api}.cpp主名与 Python 前端一致图模式category/op/torch_extension/graph_convert_${op_api}.py统一加graph_convert_前缀文档category/op/docs/torchapi_${op_api}.md与 aclnn 文档同目录文件名加torchapi_前缀主名与对外 api 一致。公共头文件放在torch_extension/cann_ops_transformer/common/下按能力域命名如aclnn_common.h、hccl_common.h。2.3 标识符命名Python 标识符函数/变量/参数小写蛇形snake_case如head_num、scale_value、input_layout。类名大驼峰PascalCase如OpBuilder、FlashAttnOpBuilder。模块级常量全大写蛇形UPPER_SNAKE_CASE如ASCEND_HOME_PATH、TORCH_DTYPE_ENUM_VALUE_TO_SCALAR_TYPE_MAP。模块内部私有符号以单下划线_前缀标识如_flash_attn_op_builder、_flash_attn。类型注解对外接口与关键内部函数应带类型注解from typing import Optional, Tuple, List可选参数统一用Optional[...]例如def flash_attn( query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, atten_mask: Optional[torch.Tensor] None, scale_value: float 1.0, head_num: int 1, input_layout: str BSH, ) - Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:参数命名一致性同一算子在 schema、meta、dispatcher、对外函数、Converter 中的同义参数应使用相同的名字如head_num、scale_value、input_layout避免在不同层出现不一致写法。C 标识符函数/局部变量/参数小写蛇形如head_num、scale_value、input_layout_ptr。命名空间算子实现统一置于namespace op_api内。常量const/constexpr常量使用全大写蛇形或大驼峰如const int DIM_THREE 3;、kATenScalarTypeToAclDataTypeTable。类型别名/结构体大驼峰如TensorWrapper、TensorListWrapper。入参类型约定必选 Tensor 用const at::Tensor 可选 Tensor 用const c10::optionalat::Tensor Tensor 列表用const std::vectorat::Tensor 可选列表用const c10::optionalstd::vectorat::Tensor 整型属性用int64_t可选整型属性用c10::optionalint64_t浮点属性用double字符串属性用std::string。Schema 标识符算子签名参数名采用小写蛇形与 Python/C 层一致。用*分隔位置参数与关键字参数*之前为必选的位置参数之后为可选的关键字参数带默认值。可选参数以?标注并给出默认值如Tensor? atten_maskNone、int? head_num1列表用Tensor[]可选列表用Tensor[]?。多输出用元组表示如- (Tensor, Tensor, Tensor, Tensor)。以flash_attn为例的 schemaflash_attn(Tensor query, Tensor key, Tensor value, *, Tensor? atten_maskNone, float scale_value1.0, int head_num1, str input_layoutBSH) - (Tensor, Tensor, Tensor, Tensor)仓库中flash_attn的真实 schema 要更复杂见 flash_attn.py它注册了两个算子——flash_attn_metadata(...) - Tensormetadata 生成接口与flash_attn(Tensor q, Tensor k, Tensor v, Tensor?block_tableNone, ..., bool return_softmax_lseFalse) - (Tensor, Tensor)主算子返回 attention_out 与 softmax_lse可作为多输入、多输出 schema 的完整参考。3. 各层实现规范3.1 C 后端category/op/torch_extension/csrc/${op_api}.cpp负责把 PyTorch 张量桥接到 aclnn C-API规范要点文件头部包含#include torch/extension.h与#include aclnn_common.h实现置于namespace op_api。函数签名与 schema 严格对应必选/可选参数类型按上文 C 入参类型约定 选择。入参校验使用TORCH_CHECK(cond, msg...)校验 shape、dtype、维度、取值范围等错误信息要可读且包含实际值例如TORCH_CHECK((head_num 0), The head_num should be greater than 0, current is: , head_num); TORCH_CHECK((query.scalar_type() key.scalar_type()), query and key should have the same dtype.);设置 DeviceGuard关键在申请输出张量之前必须先根据输入张量设置c10::OptionalDeviceGuard把当前 NPU 设备切到输入张量所在设备并用{}作用域把「DeviceGuard 输出申请」包在一起否则非默认卡调用时输出张量会落到错误设备导致 device 不一致at::Tensor attention_out{nullptr}; { auto local_device c10::Device(query.device()); const c10::OptionalDeviceGuard device_guard(local_device); attention_out at::empty(query.sizes(), query.options()); // ... 其余输出 ... }输出张量手动申请在 DeviceGuard 生效的作用域内按 meta 推导的 shape/dtype 用at::empty(...)申请输出标准 PyTorch 实践dtype 通过query.options().dtype(...)指定。拉起 kernel使用ACLNN_CMD(aclnn接口名, 入参..., 出参...)宏调用 aclnn 接口如ACLNN_CMD(aclnnFlashAttentionScore, ...)入参顺序需与 aclnn 接口定义一致该宏自动完成类型转换、workspace 申请、stream 下发与资源释放。导出绑定通过PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)将 C 函数绑定为与 schema 同名的 Python 接口m.def(flash_attn, flash_attn, flash_attn);。魔数如维度数3、默认 dtype 枚举应以具名常量表达避免裸写字面量。C wrapper 的典型骨架std::tupleat::Tensor, at::Tensor, at::Tensor, at::Tensor flash_attn( const at::Tensor query, const at::Tensor key, const at::Tensor value, const c10::optionalat::Tensor atten_mask, double scale_value, int64_t head_num, std::string input_layout) { // 3. 入参校验 TORCH_CHECK((head_num 0), The head_num should be greater than 0, current is: , head_num); at::Tensor attention_out{nullptr}; { // 4. DeviceGuard必须在申请输出之前作用域包住输出申请 auto local_device c10::Device(query.device()); const c10::OptionalDeviceGuard device_guard(local_device); // 5. 申请输出张量 attention_out at::empty(query.sizes(), query.options()); // ... 其余输出 ... } // 6. 拉起 aclnn kernel ACLNN_CMD(aclnnFlashAttentionScore, query, key, value, atten_mask, scale_value, head_num, input_layout.data(), /* outputs */ attention_out); return std::make_tuple(/* ... */); }源码印证ACLNN_CMD到底做了什么。仓库中的真实宏定义在 aclnn_common.h其内部执行了从参数元组中解码第一个 Tensor 的 deviceDecodeDevice构造OptionalDeviceGuard通过dlsym在libopapi.so/libopapi_transformer.so/libcust_opapi.so中查找aclnn_apiGetWorkspaceSize与aclnn_api的函数地址动态加载顺序见GetOpApiFuncAddr调用GetWorkspaceSize获取 workspace 大小若不为 0 则at::empty申请 workspace 张量调用ApplyDeterministicConfig()同步 PyTorch 的deterministicAlgorithms配置到 ACL 侧走 torch_npu 的ApplyDeterministicLevel弱符号缺失时回退到aclrtCtxSetSysParamOpt(ACL_OPT_DETERMINISTIC)将参数通过ConvertTypes/ConvertType批量转换为 acl 描述符aclTensor/aclScalar/aclIntArray等调用ConvertToOpApiFunc构造类型化函数指针后执行最后ReleaseConvertTypes释放描述符把 acl 调用包装进at_npu::native::OpCommand的SetCustomHandler在 NPU stream 上异步下发。类型转换方面ConvertToAclDataType通过kATenScalarTypeToAclDataTypeTable完成at::ScalarType → aclDataType的映射覆盖 float16/bfloat16/float8_e5m2/float8_e4m3fn/float4_e2m1 等全部 dtypeConvertType(const at::Tensor )还处理了 NPU 非基础格式如ACL_FORMAT_FRACTAL_NZ与 4-bit dtypeFP4_IN_INT8的 shape/stride 换算。这些公共能力由所有算子复用无需在单算子中重复实现。3.2 Python 前端category/op/torch_extension/${op_api}.py负责 JIT 编译管理、schema/meta 注册与对外接口封装OpBuilder 子类继承OpBuilder在__init__中以super().__init__(schema算子名, categorycategory)传入算子名和 category并实现三个抽象方法sources()返回相对cann_ops_transformer包根的 C 源文件路径列表如[csrc/attention/flash_attn.cpp]schema()返回算子 schema 字符串见上文 Schema 标识符register_meta()用impl(get_as_library(), self.name, Meta)注册 Meta 实现仅做 shape/dtype 推导不触碰真实 NPU 计算FakeTensor/图模式必需。Meta 中同样可用torch._check(...)做约束校验。源码印证OpBuilder.__init__的category参数会被resolve_source用于把flash_attn.cpp解析为csrc/attention/flash_attn.cpp见 builder.py与sources()返回的相对路径拼接后由get_absolute_paths转绝对路径交给 JIT。实例化与编译模块加载时实例化 builder 并调用_ensure_initialized()完成延迟初始化编译由首次load()触发_flash_attn_op_builder _FlashAttnOpBuilder() _flash_attn_op_builder._ensure_initialized()PrivateUse1 dispatcher用impl(get_as_library(), builder.name, PrivateUse1)注册 NPU 后端实现函数体透传到编译产物builder.load().算子名(...)。PrivateUse1是 PyTorch 为自定义 NPU 后端预留的 dispatch key。对外接口提供面向用户的函数flash_attn(...)负责参数整理、默认值处理等最终调用 dispatcher 实现。对外 api 必须书写注释docstring每个对外导出的接口都要有 docstring至少覆盖「功能说明、各参数含义/shape/dtype/取值范围、返回值说明」必要时给出简短调用示例。docstring 内容应与category/op/docs/torchapi_${op_api}.md保持一致便于 IDE 提示与help()查看。例如def flash_attn( query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, atten_mask: Optional[torch.Tensor] None, scale_value: float 1.0, head_num: int 1, input_layout: str BSH, ) - Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: FlashAttention 前向计算封装 aclnnFlashAttentionScore。 Args: query (Tensor): 查询张量shape 由 input_layout 决定如 BSHdtype 支持 float16/bfloat16。 key (Tensor): 键张量dtype 与 query 一致。 value (Tensor): 值张量dtype 与 query 一致。 atten_mask (Tensor, optional): 注意力掩码默认 None 表示不使用。 scale_value (float): 缩放系数默认 1.0。 head_num (int): 单卡 head 数即 query 的 N 轴长度默认 1。 input_layout (str): 输入数据排布支持 BSH/BNSD 等默认 BSH。 Returns: Tuple[Tensor, Tensor, Tensor, Tensor]: softmax_max、softmax_sum、softmax_out、attention_out。 Meta 实现、dispatcher、对外函数三者的参数顺序与默认值必须与 schema 一致。源码印证真实的flash_attnPython 前端。仓库中 flash_attn.py 的实现展示了几个进阶模式FlashAttenOpBuilder(OpBuilder)子类在__init__中super().__init__(flash_attn, categoryattention)schema()返回列表OpBuilder.register_schema支持 str 或 list逐个get_as_library().define(schema)注册register_meta()使用torch.library.register_fake(cann_ops_transformer::flash_attn)注册 fake 实现内部用_validate_layout_dims/torch._check做 layout 与维度约束校验并按layout_qTND/BSND/BNSD与layout_out推导attention_out、softmax_out的 shape对外接口通过flash_attn torch.ops.cann_ops_transformer.flash_attn直接指向命名空间算子dispatcher 层负责把 Python 侧Optional默认值规范化为 C 侧约定的-1/BSND等哨兵值metadata 生成算子还注册了torch.library.register_kernel(..., None)的 fallback处理全 Tensor 为 None 的调用场景。3.3 图模式 Convertercategory/op/torch_extension/graph_convert_${op_api}.py负责在 torchair 图模式GE下把 aten 算子转换为 GE 节点可选依赖保护torchair 相关 import 统一包在try/except ImportError中用_TORCHAIR_AVAILABLE标志位控制避免在无 torchair 环境下导入失败。GE op 函数定义与op_type同名的大驼峰函数如FlashAttentionScore通过 docstring 写明REG_OP的 IR 定义INPUT/DYNAMIC_INPUT/OPTIONAL_INPUT/OUTPUT/ATTR 等并组织inputs/attrs/outputs后调用ge_op(...)IR 通过IrDef(...)链式声明。Converter 注册用register_fx_node_ge_converter(torch.ops.cann_ops_transformer.flash_attn.default)装饰convert_flash_attn函数其参数顺序与 schema 完全一致函数体调用上面的 GE op 函数。在ops/__init__.py中导出 Converter如convert_flash_attn确保注册逻辑被执行。源码印证真实的 graph_convert_flash_attn.py 展示了完整的可选依赖保护写法_TORCHAIR_AVAILABLE True/False分支并以auto_convert_to_tensor标注哪些参数需要张量化。注意当前仓库中flash_attn的 Converter 函数体是raise AssertionError(GE not supported!)——图模式对flash_attn尚未开放这提示开发者若图模式暂不支持Converter 可仅保留注册骨架并在函数体内显式报错同时必须在算子文档中标注「图模式调用暂不支持」。3.4 对外导出算子__init__.py与包根__init__.py对外导出分两级包根已无需手动维护 import 列表category/op/torch_extension/__init__.py算子层每个新增算子的对外接口与 Converter 都需在__all__中声明并显式 import 导出。导入即触发 schema/meta/converter 注册__all__ [flash_attn, convert_flash_attn] from .flash_attn import flash_attn from .graph_convert_flash_attn import convert_flash_attn若算子有额外的辅助导出如flash_attn_metadata、get_symm_buffer_for_mega_moe也一并加入__all__和 import。仓库中真实的算子层__init__.py见 attention/flash_attn/torch_extension/init.py导出flash_attn、flash_attn_metadata并 importgraph_convert_flash_attn。ops/__init__.py子包层已改造为自动发现模式目录扫描 entry point无需手动添加 import。新增算子只要按规范放置文件并定义__init__.py的__all__就会被自动加载。cann_ops_transformer/__init__.py包根层通过from . import ops触发注册并通过__getattr__和__dir__动态导出算子接口使用户可直接通过cann_ops_transformer.接口名访问见 cann_ops_transformer/init.py。4. 文档规范category/op/docs/torchapi_${op_api}.md每个对外算子 api 需配套一份中文文档放在对应算子的docs/目录下与aclnn${OpName}.md并列文件名形如torchapi_${op_api}.md建议章节顺序与已有算子文档如 attention/flash_attn/docs/torchapi_flash_attn.md对齐标题算子名特殊字符如_需转义为\_。产品支持情况表格列出支持的产品形态如Ascend 950PR/Ascend 950DT及是否支持。仓库中真实文档使用!-- npu950 id1 --注释块逐产品标注支持/不支持状态便于自动化同步。功能说明API 功能概述 计算公式数学表达用 LaTeX并说明各符号与参数的对应关系。如flash_attn文档给出了Attention(Q,K,V)Softmax(QK^T/√d)V、引入 sink 后的 softmax_max/softmax_sum 修正公式、return_softmax_lse时softmax_lse的计算式。函数原型代码块给出完整函数签名含默认值与*分隔。参数说明逐个参数说明「必选/可选、语义、shape、dtype、数据格式如$ND$、是否支持非连续 Tensor、取值范围/约束」可选参数标注默认值与「暂不支持」说明。输出说明逐个输出说明 shape、dtype、格式等。约束说明分类列出参数一致性约束、shape/取值范围约束、量化场景约束等通信类算子还需列出通信域约束。配套接口说明若算子需与其他接口配套使用补充其原型、参数与输出说明。例如flash_attn文档明确要求先调用前置接口flash_attn_metadata完成负载均衡计算再调用主算子并传入 metadata。调用示例给出单算子模式必要时含多卡/通信初始化的完整可运行示例图模式若暂不支持需明确标注「图模式调用暂不支持」。5. 编码通用约束许可证头所有新增源文件.py/.cpp/.h必须包含 Huawei 版权与 CANN Open Software License Agreement Version 2.0 许可证头年份填当年。Python/脚本用#注释C 用//或/* */。仓库所有源文件顶部均可看到该头如 builder.py。接口注释对外 api 接口必须书写 docstring功能、参数、返回值见上文 3.2 Python 前端C wrapper 关键逻辑校验、DeviceGuard、aclnn 调用也应有简要注释。C 层 DeviceGuard关键调用 aclnn 的 C wrapper 中必须在申请输出张量之前用c10::OptionalDeviceGuard构造自c10::Device(输入张量.device())把设备切到输入张量所在设备详见 3.1 C 后端。参数校验前置Python 侧用torch._check(cond, lambda: f...{var}...)C 侧用TORCH_CHECK(cond, msg...)错误信息需包含变量实际值便于定位。源码中flash_attn的 C 侧CheckLayoutDims即按 layout 校验张量维度并输出实际 shape见 flash_attn.cpp。错误码Python 侧可结合torch_npu.utils._error_code的ErrCode/ops_error输出规范错误码如f... {ops_error(ErrCode.VALUE)}.。避免魔数维度数、dtype 枚举值等以具名常量表达并在文档/注释中说明枚举含义如23 → float8_e5m2、24 → float8_e4m3fn。flash_attn.cpp中DIM_ONE~DIM_FIVE、MAX_DIM_SIZE即为此类具名常量。公共能力复用类型转换、ACLNN_CMD、通信域处理等优先复用common下的公共头不在各算子中重复实现。一致性自检提交前确认 schema、C wrapper、Meta、dispatcher、对外函数、Converter、文档七处的算子名、参数名、参数顺序、默认值保持一致。6. 结语与自查清单新增一个算子 API 时建议按以下顺序完成并逐项自查步骤产物一致性检查点1csrc/${op_api}.cpp函数名与 schema 一致、DeviceGuard 包裹输出申请、TORCH_CHECK校验、ACLNN_CMD参数顺序与 aclnn 接口一致2${op_api}.pyOpBuilder子类三方法齐全、_ensure_initialized()已调用、dispatcher 透传参数与 schema 一致3graph_convert_${op_api}.py_TORCHAIR_AVAILABLE保护、Converter 参数顺序与 schema 一致4torch_extension/__init__.py__all__覆盖对外接口与 Converter5自动发现文件按规范放置、__init__.py定义__all__无需手动注册6docs/torchapi_${op_api}.md 索引与 aclnn 文档同目录同步更新 docs/zh/menu_torch_api.md 与 docs/zh/torch_api_list.md作为模板attention/flash_attn是当前仓库中 torch_extension 结构最完整的参考实现Python 前端、C 后端、图模式 Converter、算子层导出与接口文档齐全通信类算子可参考mc2/mega_moe的目录组织。遵循本文规范可以保证新增算子 API 与现有体系在命名、调度、文档与自动发现机制上完全对齐并规避多卡 device 错位、schema 不一致导致 JIT 产物无法调用等常见问题。赞分享算子库人工智能深度学习Ascend【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址https://gitcode.com/cann/ops-transformer点击查看免费下载相关推荐CANN ops-transformer torch_extension 开发指南基于 JIT 的 PyTorch NPU 算子扩展库CANN ops transformer torch_extension 开发指南基于 JIT 的 PyTorch NPU 算子扩展库 导读 本文围绕 tor算子库人工智能深度学习AscendCANN ops-nn torch_extension 接口全景指南从 PyTorch 调用到 NPU 算子加速CANN ops nn torch_extension 接口全景指南从 PyTorch 调用到 NPU 算子加速 CANN 开源神经网络算子库 ops nn人工智能算子库深度学习CANNAscend基于 Ascend C 与 PyTorch Extension 开发自定义 NPU 算子的完整指南CANN ops-nn 实战基于 Ascend C 与 PyTorch Extension 开发自定义 NPU 算子的完整指南CANN ops nn 实战 本文以 CANN 开源算子库人工智能算子库深度学习CANNAscend创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考