CANN ops-transformer Fast Kernel Launch 实战用 Ascend C PyTorch Extension 单文件开发自定义 NPU 算子【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer本篇指南基于 ops-transformer 仓库的 Fast Kernel Launch 示例讲解如何用 Ascend C 与 PyTorch Extension 在“单个 C 文件”中完成自定义 NPU 算子的开发、编译与框架适配。读完本文你将掌握该示例从环境部署、Wheel 构建到新增算子的完整流程并能理解 Schema 注册、Meta 函数、Ascend C Kernel 与 NPU 调用注册四层实现背后的构建系统原理。一、核心思想单交付件 直接启动核函数Fast Kernel Launch快速内核启动是 CANN 提供的一种轻量算子开发模式示例工程位于 examples/fast_kernel_launch_example。与传统 AI Core 算子需要 op_hostTiling/ op_kernelTik/ op_graph 等多目录交付件不同该模式的两个核心优势是单交付件一个.cpp文件完成算子开发和 PyTorch 框架适配同时包含算子 Schema 注册、Meta 函数实现、Ascend C Kernel 实现与 NPU 调用注册四个模块高效调用直接使用 CUDA 风格的numBlocks, nullptr, stream语法启动核函数无需编写 Tiling 数据序列化与 op_host 适配层流程简单高效。该模式非常适合快速验证算子思路、原型开发仓库中已用它承载了包括分组矩阵乘grouped_matmul、增量 Flash Attentionincre_flash_attention、MoE 分布式 combine/dispatch 等多个较复杂的算子说明该模式并不局限于“简单算子”。二、环境要求在开始之前需要先完成基础环境CANN 套件搭建仓库内参考文档为 环境部署。Fast Kernel Launch 示例自身的额外要求如下依赖要求说明gcc9.4.0宿主机 C 编译Python3.8构建与运行环境torch 2.6.0PyTorch 框架TorchNPU对应版本需与 torch/CANN 版本匹配依赖清单见 requirements.txt包含build、pyyaml、numpy2、pytest并配置了 PyTorch CPU 版的--extra-index-url。三、安装步骤3.1 安装依赖cd examples/fast_kernel_launch_example python3 -m pip install -r requirements.txt3.2 构建 Wheel 包# NPU_SOC_VERSION 设置编译款型 # Atlas A2 系列产品使用 ascend910b默认值 # Atlas A3 系列产品使用 ascend910_93 # Ascend 950PR / Ascend 950DT 产品使用 ascend950 export NPU_SOC_VERSIONascend910b # -n: non-isolated build (uses existing environment) python3 -m build --wheel -n构建完成后产物位于当前目录的dist文件夹下产物名为ascend_ops-1.0.0-${python_version}-abi3-${arch}.whl。其中${python_version}表示当前环境中的 Python 版本如 python3.8.3 对应cp38${arch}表示 CPU 架构平台标签。关于构建过程从 setup.py 的CMakeBuildCommand可以看到它实际向 CMake 传递的关键参数-DCMAKE_BUILD_TYPERelease-DTorch_DIR取自动态导入的torch.utils.cmake_prefix_path即当前环境 PyTorch 的安装路径-DTORCH_NPU_PATH取自动态导入的torch_npu.__file__所在目录-DNPU_SOC_VERSION读取环境变量NPU_SOC_VERSION未设置时默认ascend910b对应 setup.py#L117-DCANN_3RD_LIB_PATH指向仓库根目录下的third_party示例工程相对路径../../third_party用于引入仓库内的第三方 CMake 模块。setup.py 中的ABI3Wheel强制使用abi3标签打包配合 CMake 中Py_LIMITED_API0x03080000的编译定义见 CMakeLists.txt#L100使同一个 Wheel 可支持 Python 3.8 的多个版本。3.3 安装 Wheel 包python3 -m pip install dist/*.whl --force-reinstall --no-deps3.4 可选清理编译缓存再次构建前建议先执行python setup.py clean从 setup.py#L27-L53 可以看到clean命令会删除build、dist、ascend_ops.egg-info三个目录并清理所有.pyc/.pyo文件确保增量构建不会引入陈旧产物。此外仓库还提供一个一键脚本 build_and_test.sh它按“安装依赖 →setup.py clean→ 构建并安装 Wheel → 遍历tests/下各子目录逐一运行 pytest”的完整流程自动执行适合在 CI 或新环境中快速验证。四、快速开始像普通 PyTorch 算子一样调用安装完成后你可以像使用普通 PyTorch 操作一样使用 NPU 算子。以 add 算子调用为例import torch import torch_npu import ascend_ops # 构建出的python包 # Initialize data on NPU x torch.randn(10, 32, dtypetorch.float32).npu() y torch.randn(10, 32, dtypetorch.float32).npu() # Call the custom NPU operator npu_result torch.ops.ascend_ops.add(x, y) # PyTorch Custom Operator Dispatch机制: torch.ops.library_name.operator_name # Verify against CPU ATen implementation cpu_x x.cpu() cpu_y y.cpu() cpu_result cpu_x cpu_y assert torch.allclose(cpu_result, npu_result.cpu(), rtol1e-6) print(Verification successful!)调用约定为torch.ops.library_name.operator_name。这里library_name是 CMake 中的EXTENSION_MODULE_NAME固定为ascend_ops见 CMakeLists.txt#L32operator_name则是 C 侧注册的算子名如add。Python 包入口 ascend_ops/__init__.py 会先导入编译产物_C即 CMake 生成的_C.abi3.so导入失败会抛出带提示的ImportError随后导入 ascend_ops/ops.py 等模块其中groupedmatmul函数展示了如何用一层 Python 包装把torch.ops.ascend_ops.groupedmatmul暴露为更易用的接口。五、开发指南新增一个算子以开发add算子为例只需提供一个 C 实现共五步。5.1 建立目录结构在csrc目录下使用算子名add建立文件夹并在其内按当前要开发的 SoC 款型建立子文件夹ascend910bcsrc/ └── add/ └── ascend910b/ ├── CMakeLists.txt └── add.cpp这个“算子名/SoC 款型”两级目录约定被构建系统自动识别cmake/func.cmake 中的recursive_add_subdirectory()宏会遍历csrc下的每个子目录仅当存在算子名/NPU_SOC_VERSION/CMakeLists.txt时才执行add_subdirectory。也就是说同一个 Wheel 工程可以为不同 SoC 款型维护不同实现编译时只构建当前NPU_SOC_VERSION对应的子目录。仓库中csrc下已有add/ascend910b、grouped_matmul/ascend910b、incre_flash_attention/ascend910_93、moe_distribute_combine_v2/ascend910_93等多款示例可参考。5.2 编写 CMakeLists.txt在 SoC 目录下新建CMakeLists.txtadd_sources(--npu-archdav-2201)这里dav-2201为 ascend910b 芯片对应的--npu-arch编译参数其他款型请使用对应的 arch 编码可从 CANN 的 NpuArch 说明中查询。add_sources宏定义于 cmake/func.cmake#L24-L71会将传入参数作为编译 flags、追加-xasc以 Ascend C 方式编译、递归收集当前目录下所有.cpp源文件、创建算子名_obj目标库并汇入全局OBJECTS_LIST。如果算子内核参数较多还可以像 grouped_matmul 的 CMakeLists.txt 那样追加--cce-aicore-input-parameter-size4096等 CCE 编译参数来扩大 AI Core 内核入参上限。5.3 编写算子实现文件 add.cpp在 SoC 目录下新建add.cpp建议使用算子名作为文件名。这个文件包含开发一个 AI Core 算子所需的全部模块完整实现见 csrc/add/ascend910b/add.cpp结构如下#include ATen/Operators.h #include torch/all.h #include torch/library.h #include torch_npu/csrc/core/npu/NPUStream.h #include torch_npu/csrc/framework/OpCommand.h #include kernel_operator.h #include platform/platform_ascendc.h #include type_traits namespace ascend_ops { // 当前项目为一个命名空间 namespace Add { // 建议每个算子自己有一个独立的namespace防止全局变量污染 /** * 将算子schema注册给PyTorch框架 * 框架知道有这样一个算子 */ // Register the operators schema TORCH_LIBRARY_FRAGMENT(EXTENSION_MODULE_NAME, m) { m.def(add(Tensor x, Tensor y) - Tensor); } /** * 实现算子的Meta函数即InferShapeInferDtype * 根据输入推导出这个算子的输出是什么样子需要多少空间不需要实际计算这个算子 */ // Meta function implementation of Add torch::Tensor add_meta(const torch::Tensor x, const torch::Tensor y) { TORCH_CHECK(x.sizes() y.sizes(), The shapes of x and y must be the same.); auto z torch::empty_like(x); return z; } /** * 将算子的Meta函数注册给框架 * 框架可以调用这个Meta函数在真正执行这个算子计算前知道需要多大空间 * 后续可以支持torch.compile/AutoGrad/AclGraph等图加速 */ // Register the Meta implementation TORCH_LIBRARY_IMPL(EXTENSION_MODULE_NAME, Meta, m) { m.impl(add, add_meta); } /** * NPU算子Kernel实现使用AscendC API面向当前的soc编写 */ template typename T __global__ __aicore__ void add_kernel(GM_ADDR x, GM_ADDR y, GM_ADDR z, int64_t totalLength, int64_t blockLength, uint32_t tileSize) { // kernel implementation } /** * 实现算子调用接口 * 在这个接口中需要完成NPU Kernel的调用 * 1. 计算出输出的Tensor的个数/Shape/Dtype(可以调用Meta函数实现也可以直接实现) * 2. 计算Tiling根据Shape得到如何分块计算 * 3. 调用NPU Kernel */ torch::Tensor add_npu(const torch::Tensor x, const torch::Tensor y) { // OptionalDeviceGuard确保后续操作在正确的设备上下文执行 // 它会记录当前设备状态执行完作用域代码后自动恢复 const c10::OptionalDeviceGuard guard(x.device()); auto z add_meta(x, y); auto stream c10_npu::getCurrentNPUStream().stream(false); int64_t totalLength, numBlocks, blockLength, tileSize; totalLength x.numel(); std::tie(numBlocks, blockLength, tileSize) calc_tiling_params(totalLength); auto x_ptr (GM_ADDR)x.data_ptr(); auto y_ptr (GM_ADDR)y.data_ptr(); auto z_ptr (GM_ADDR)z.data_ptr(); auto acl_call []() - int { AT_DISPATCH_SWITCH( x.scalar_type(), add_npu, // 根据不同的数据类型调用不同的NPU Kernel AT_DISPATCH_CASE(torch::kFloat32, [] { using scalar_t float; add_kernelscalar_tnumBlocks, nullptr, stream(x_ptr, y_ptr, z_ptr, totalLength, blockLength, tileSize); }) AT_DISPATCH_CASE(torch::kFloat16, [] { using scalar_t half; add_kernelscalar_tnumBlocks, nullptr, stream(x_ptr, y_ptr, z_ptr, totalLength, blockLength, tileSize); }) AT_DISPATCH_CASE(torch::kInt32, [] { using scalar_t int32_t; add_kernelscalar_tnumBlocks, nullptr, stream(x_ptr, y_ptr, z_ptr, totalLength, blockLength, tileSize); }) ); return 0; }; // 需要使用RunOpApi/RunOpApiV2接口调用保证时序与TorchNPU调用aclnn接口一致。 at_npu::native::OpCommand::RunOpApi(Add, acl_call); return z; } /** * 将算子的调用函数注册给框架Device为PrivateUse1 * 框架知道当输入均在NPU Device上时Dispatch到这个算子实现 */ // Register the NPU implementation TORCH_LIBRARY_IMPL(EXTENSION_MODULE_NAME, PrivateUse1, m) { m.impl(add, add_npu); } } // namespace Add } // namespace ascend_ops四个模块的职责可以概括为Schema 注册TORCH_LIBRARY_FRAGMENT让 PyTorch 框架“知道”有这么一个算子其签名参数名、类型、默认值会暴露到 Python 侧的torch.ops.ascend_opsMeta 函数TORCH_LIBRARY_IMPL(..., Meta, ...)只做 InferShape InferDtype不实际计算。注册 Meta 实现后torch.compile/ AutoGrad / AclGraph 等图模式才能在真正执行前完成 shape/dtype 推导与内存规划Kernel 实现以__global__ __aicore__模板函数编写 Ascend C 内核面向当前 SoC 款型编译NPU 调用注册TORCH_LIBRARY_IMPL(..., PrivateUse1, ...)声明“当输入均在 NPU Device 上时 dispatch 到该实现”并在其中完成输出推导、Tiling 计算与内核启动。5.4 重新构建并安装按第三节“安装步骤”重新执行构建与安装命令即可新增算子无需修改顶层 CMake 文件。5.5 基于 pytest 测试算子 API参考 tests/add/test_add.py 的实现。该测试包含两类用例值得借鉴接口存在性测试test_add_interface_exist仅断言torch.ops.ascend_ops下能发现add。它可以守护一个常见故障——算子在 C 侧实现并注册了但因 Schema 与 C 注册签名不匹配而未暴露到 Pythontorch.ops命名空间参数化功能测试test_add_operator用 19 种 shape从(1,)到(1000, 1000)× 3 种 dtypefloat32/float16/int32的组合对拍 CPU 结果浮点类型使用rtol1e-4, atol1e-4整型要求严格相等并通过torch.npu.is_available()判断无 NPU 环境时自动跳过。六、源码纵深构建系统与内核实现细节6.1 顶层构建Bisheng 编译器与_C.abi3.so顶层 CMakeLists.txt 的关键点通过find_package(ASC REQUIRED)、find_package(AICPU REQUIRED)定位 CANN 的 ASC/AICPU 工具链工程语言声明为LANGUAGES ASC AICPU CXXcmake/ascend.cmake 负责定位 Ascend toolkit 路径优先读环境变量ASCEND_HOME_PATH否则依次检查/usr/local/Ascend/或~/Ascend/下的latest路径并把Bisheng 编译器${ASCEND_DIR}/${系统架构}-linux/ccec_compiler/bin/bisheng同时设为 C/C 编译器与链接器——Ascend C 源码正是由它编译为目标板核函数cmake/torch.cmake 与 cmake/torch_npu.cmake 分别解析 PyTorch 与 TorchNPU 的头文件、库路径顶层 CMakeLists.txt#L89-L111 把 csrc/extension.cpp一个最小的PyInit__C入口与所有算子名_obj目标库链接成_C.abi3.so动态库编译选项包含-DEXTENSION_MODULE_NAMEascend_ops -DTORCH_MODE链接库包括torch_npu、ascendcl、platform、register、tiling_api、runtime、hccl、hcomm等POST_BUILD 阶段会自动把产物拷贝回ascend_ops/包目录随 Wheel 一起分发。从源码结构看-DTORCH_MODE与 TorchNPU 的OpCommand框架配合使自定义算子的执行时序与 TorchNPU 调用 aclnn 接口保持一致——这正是 5.3 节中at_npu::native::OpCommand::RunOpApi(Add, acl_call)的作用把内核启动 lambda 包进 TorchNPU 的统一调用时序中而不是绕过框架直接裸调 ACL。6.2 内核实现Tiling 参数推导与双缓冲流水线add内核add.cpp#L64-L146是一个完整的 Ascend C 事件流流水线示例Tiling 推导calc_tiling_paramsadd.cpp#L48-L62通过platform_ascendc::PlatformAscendCManager动态查询板卡 UB 大小与 AI Core 数量numBlocks min(coreNum, ceil(totalLength/1024))每核至少 1024 个元素blockLength为每核分摊长度tileSize ubSize / PIPELINE_DEPTH(2) / BUFFER_NUM(3)。也就是说Fast Kernel Launch 模式下 Tiling 计算可以写在 host 侧代码里用平台 API 动态获得硬件参数流水线内核用TPipe 三条TQueVECIN两个输入队列、VECOUT一个输出队列PIPELINE_DEPTH2双缓冲组织 CopyIn → ComputeAscendC::Add→ CopyOut 三段流水越界安全每个核先按AscendC::GetBlockIdx()偏移全局缓冲区并单独处理最后一块不足一个 tile 的尾部数据tail tile保证totalLength不是 tile 整数倍时结果依然正确。6.3 多 SoC 与复杂算子的工程实践csrc下的其他示例展示了该模式的工程化用法可作为进阶参考csrc/incre_flash_attention/ascend910_93/npu_fused_infer_attention_score.cpp增量 Flash Attention 算子的 Schema 拥有 20 多个可选张量/标量参数量化 scale/offset、Rope、block_table、learnable_sink 等并通过m.def原始字符串语法声明复杂 Schema内核参数多到需要自定义LAUNCH_INCRE_FA宏统一传参同时配套独立的incre_flash_attention_meta算子目录见 csrc/incre_flash_attention_meta/ascend910_93/CMakeLists.txt与 AclGraph 图模式测试tests/incre_flash_attention/test_aclgraph.py验证 Meta 注册后可被图执行模式消费csrc/moe_distribute_combine_v2/ascend910_93/ 与 csrc/moe_distribute_dispatch_v2/ascend910_93/将 kernel 头文件放到 SoC 目录下的op_kernel/子目录、把 torch 适配与校验逻辑拆成多个.cpp由add_sources递归收集编译并各自附带 README 与 pytest 测试说明单交付件模式同样支持“主文件 头文件 多编译单元”的组织方式。6.4 一键构建与测试仓库提供的 build_and_test.sh 完整流程为pip install -r requirements.txt python3 setup.py clean python3 -m build --wheel --no-isolation python3 -m pip install dist/*.whl --force-reinstall --no-deps # 遍历 tests/ 下每个算子目录逐个运行 pytest -v七、要点小结Fast Kernel Launch 通过“一个.cpp文件 Schema Meta Kernel NPU 调用”的四段式结构把自定义 NPU 算子开发压缩到最小交付面目录约定csrc/算子名/SoC款型/NPU_SOC_VERSION环境变量ascend910b/ascend910_93/ascend950共同实现多 SoC 款型的按款编译内核用blocks, nullptr, stream直接启动Tiling 可在 host 侧借助PlatformAscendCManager动态推导注册 Meta 实现对torch.compile/AutoGrad/AclGraph 等图加速能力是必要前提启动内核应包在OpCommand::RunOpApi中以保证与 TorchNPU 调用时序一致测试上建议同时做“接口存在性 shape/dtype 参数化对拍”两层参考 tests/add/test_add.py。适用前提提醒本示例要求已正确安装 CANN 基础环境与匹配的 torch 2.6.0 / TorchNPU 组合--npu-arch参数必须与NPU_SOC_VERSION所指的 SoC 款型一致否则编译的核函数无法在目标硬件上运行。【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考