企业官网建设流程全解析

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_host（Tiling）/ op_kernel（Tik）/ op_graph 等多目录交付件不同，该模式的两个核心优势是：

• 单交付件：一个.cpp文件完成算子开发和 PyTorch 框架适配，同时包含算子 Schema 注册、Meta 函数实现、Ascend C Kernel 实现与 NPU 调用注册四个模块；
• 高效调用：直接使用 CUDA 风格的<<<numBlocks, nullptr, stream>>>语法启动核函数，无需编写 Tiling 数据序列化与 op_host 适配层，流程简单高效。

该模式非常适合快速验证算子思路、原型开发，仓库中已用它承载了包括分组矩阵乘（grouped_matmul）、增量 Flash Attention（incre_flash_attention）、MoE 分布式 combine/dispatch 等多个较复杂的算子，说明该模式并不局限于“简单算子”。

二、环境要求

在开始之前，需要先完成基础环境（CANN 套件）搭建，仓库内参考文档为 环境部署。Fast Kernel Launch 示例自身的额外要求如下：

依赖清单见 requirements.txt，包含build、pyyaml、numpy<2、pytest，并配置了 PyTorch CPU 版的--extra-index-url。

三、安装步骤

3.1 安装依赖

cd examples/fast_kernel_launch_example python3 -m pip install -r requirements.txt

3.2 构建 Wheel 包

# NPU_SOC_VERSION 设置编译款型： # Atlas A2 系列产品使用 "ascend910b"（默认值） # Atlas A3 系列产品使用 "ascend910_93" # Ascend 950PR / Ascend 950DT 产品使用 "ascend950" export NPU_SOC_VERSION=ascend910b # -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_TYPE=Release；
• -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_API=0x03080000的编译定义（见 CMakeLists.txt#L100），使同一个 Wheel 可支持 Python >= 3.8 的多个版本。

3.3 安装 Wheel 包

python3 -m pip install dist/*.whl --force-reinstall --no-deps

3.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, dtype=torch.float32).npu() y = torch.randn(10, 32, dtype=torch.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(), rtol=1e-6) print("Verification successful!")

调用约定为torch.ops.<library_name>.<operator_name>。这里library_name是 CMake 中的EXTENSION_MODULE_NAME（固定为ascend_ops，见 CMakeLists.txt#L32），operator_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 款型建立子文件夹ascend910b：

csrc/ └── 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.txt：

add_sources("--npu-arch=dav-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-size=4096等 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 operator's schema TORCH_LIBRARY_FRAGMENT(EXTENSION_MODULE_NAME, m) { m.def("add(Tensor x, Tensor y) -> Tensor"); } /** * 实现算子的Meta函数，即InferShape+InferDtype * 根据输入推导出这个算子的输出是什么样子，需要多少空间，不需要实际计算这个算子 */ // 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_kernel<scalar_t><<<numBlocks, nullptr, stream>>>(x_ptr, y_ptr, z_ptr, totalLength, blockLength, tileSize); }) AT_DISPATCH_CASE(torch::kFloat16, [&] { using scalar_t = half; add_kernel<scalar_t><<<numBlocks, nullptr, stream>>>(x_ptr, y_ptr, z_ptr, totalLength, blockLength, tileSize); }) AT_DISPATCH_CASE(torch::kInt32, [&] { using scalar_t = int32_t; add_kernel<scalar_t><<<numBlocks, 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

四个模块的职责可以概括为：

1. Schema 注册（TORCH_LIBRARY_FRAGMENT）：让 PyTorch 框架“知道”有这么一个算子，其签名（参数名、类型、默认值）会暴露到 Python 侧的torch.ops.ascend_ops；
2. Meta 函数（TORCH_LIBRARY_IMPL(..., Meta, ...)）：只做 InferShape + InferDtype，不实际计算。注册 Meta 实现后，torch.compile/ AutoGrad / AclGraph 等图模式才能在真正执行前完成 shape/dtype 推导与内存规划；
3. Kernel 实现：以__global__ __aicore__模板函数编写 Ascend C 内核，面向当前 SoC 款型编译；
4. 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 种 dtype（float32/float16/int32）的组合对拍 CPU 结果，浮点类型使用rtol=1e-4, atol=1e-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 CXX；
• cmake/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_NAME=ascend_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_params（add.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+ 三条TQue（VECIN两个输入队列、VECOUT一个输出队列，PIPELINE_DEPTH=2双缓冲）组织 CopyIn → Compute（AscendC::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

七、要点小结

1. Fast Kernel Launch 通过“一个.cpp文件 = Schema + Meta + Kernel + NPU 调用”的四段式结构，把自定义 NPU 算子开发压缩到最小交付面；
2. 目录约定csrc/<算子名>/<SoC款型>/+NPU_SOC_VERSION环境变量（ascend910b/ascend910_93/ascend950）共同实现多 SoC 款型的按款编译；
3. 内核用<<<blocks, nullptr, stream>>>直接启动，Tiling 可在 host 侧借助PlatformAscendCManager动态推导；
4. 注册 Meta 实现对torch.compile/AutoGrad/AclGraph 等图加速能力是必要前提；启动内核应包在OpCommand::RunOpApi中以保证与 TorchNPU 调用时序一致；
5. 测试上建议同时做“接口存在性 + 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