CANN SuperKernel 样例实战指南JIT 与 AOT 两类 Python 样例的运行、编译与调优【免费下载链接】graph-autofusionGraph-autofusion 是一个面向昇腾Ascend芯片的轻量级、解耦式组件集合旨在通过自动融合技术加速模型执行。 目前已开源 SuperKernel 组件和 Autofuse 组件未来将持续开放更多自动融合相关模块。项目地址: https://gitcode.com/cann/graph-autofusion本篇指南围绕 super_kernel/examples/README.md 展开系统讲解 CANN 开源仓库中 SuperKernel 组件的两类 Python 样例——基于 JIT 接口的运行时融合样例与基于 TorchAirnpugraph_ex的 AOT 静态编译样例并结合 run_example.sh、common.sh 及各样例源码说明--npu-arch架构选择逻辑、JIT/AOT 样例的具体实现路径、SuperKernel 静态 kernel 的退出清理机制帮助你在 Ascend NPU 上完整跑通、复现并改造这些样例。一、样例体系概览JIT 与 AOT 两条技术路线super_kernel/examples/目录提供两类 SuperKernel Python 样例分别对应 SuperKernel 的两种落地方式jit/通过 SuperKernel JIT 接口完成融合、编译和运行。样例直接在 Python 代码中用tng.scope.super_kernel上下文管理器标注融合范围配合torch.compile走 TorchAir 的 NPU 后端编译执行适合在训练/推理脚本中快速验证融合效果。aot/通过 TorchAirnpugraph_ex静态编译并启用 SuperKernel 优化。样例将图以静态方式编译为 kernel 包再加载执行覆盖双流控制边、SuperKernel options 调优、自定义 AscendC 算子融合等更完整的工程场景。两类样例统一由仓库根目录下的super_kernel/examples/run_example.sh入口脚本驱动公共 Bash 函数集中放在 aot/scripts/common.sh 中。二、目录结构examples/ ├── run_example.sh # 样例统一运行入口 ├── jit/ │ ├── example01_super_kernel_base/ # SuperKernel 基础功能 │ ├── example02_super_kernel_profiling/ # SuperKernel profiling 对比 │ └── example03_super_kernel_runtime_ascendc_only/ # AscendC Runtime 极简样例 └── aot/ ├── scripts/ # AOT 样例公共 Bash 函数 ├── example01_dual_stream/ # 双流与 NPU Event 控制边 ├── example02_sk_options/ # SuperKernel options └── example03_kernel_pybind/ # Pybind 自定义算子融合各样例目录下均提供中英文README.md可作为深入阅读的起点例如 JIT 基础样例说明、双流样例说明、Pybind 样例说明。三、前置条件在运行样例前需要完成两项准备环境构建先参考 源码构建指南仓库文档路径docs/zh/build.md对应文档中的相对链接../../docs/zh/build.md完成 CANN 环境准备确保当前机器有可见的 Ascend NPU。所有样例均使用当前可见 NPU代码中通过torch.npu.set_device(0)或aclrt_set_device(DEVICE_ID)指定为 0 号设备。Python 依赖安装 requirements.txt 中声明的依赖protobuf3.13,4 torch2.6.0 torch_npu2.6.0.post5 PyYAML5.1其中torch/torch_npu版本对齐是 JIT 样例能正常torch.compile到 NPU 后端的关键前提。四、统一运行入口与--npu-arch架构选择样例运行必须显式传入--npu-arch参数用于选择适用当前硬件的样例集合bash super_kernel/examples/run_example.sh --npu-archdav-2201 bash super_kernel/examples/run_example.sh --npu-archdav-3510从入口脚本 run_example.sh 的源码可以看出其执行逻辑脚本首先解析 Python 解释器优先使用VIRTUAL_ENV环境其次回退到仓库根目录的venv/bin/python最后回退到系统python3通过common.sh中的sk_parse_npu_arch解析参数该函数用getopt仅接受-h/--help与--npu-arch两个参数未传--npu-arch会报ERROR: --npu-arch is required.取值非法既非dav-2201也非dav-3510会报unsupported NPU architecture并以退出码 2 失败随后按架构分支执行样例。架构差异说明这一点在源码中体现得很直接dav-2201会依次运行全部三个 JIT 样例superkernel_scope.py、superkernel_compare.py、superkernel_runtime_ascendc_basic.pydav-3510不支持这些 JIT 样例脚本打印[INFO] Skipping SuperKernel JIT examples on dav-3510.并跳过仅运行 AOT 样例。一个需要注意的现状当前仓库版本中run_example.sh对三个 AOT 样例的调用行处于注释状态并附有注释暂时禁用因当前 CI 环境限制无法运行的 AOT examples。因此在统一入口下dav-3510分支目前实际会打印跳过信息后直接进入Run all examples success。如需运行 AOT 样例应改用各样例自带的独立入口见第五节例如bash super_kernel/examples/aot/example01_dual_stream/run.sh --npu-archdav-2201此外run_example.sh中的 JIT/AOT 分支逻辑印证了 README 中样例内部的编译参数由各样例处理的说明入口脚本只负责架构分流和解释器选择真正的编译配置SOC 版本、options 等由各样例内部完成。五、JIT 样例详解5.1 example01_super_kernel_baseSuperKernel 基础功能入口文件 superkernel_scope.py 演示了 SuperKernel 最核心的使用范式——用tng.scope.super_kernel上下文管理器在nn.Module.forward中标注融合范围import torchair as tng from torchair.configs.compiler_config import CompilerConfig class Network(nn.Module): def forward(self, gmm1_x, gmm1_weight, gmm1_bias, gmm2_weight, moe1_bias, dsq_input, data1, data2, scale): with tng.scope.super_kernel(sk1): grouped_matmul_01 torch_npu.npu_grouped_matmul(group_type-1, xgmm1_x, weightgmm1_weight, biasgmm1_bias) grouped_matmul_02 torch_npu.npu_grouped_matmul(group_type-1, xgrouped_matmul_01, weightgmm2_weight) moe_gating_top_k_01 torch_npu.npu_moe_gating_top_k(xgrouped_matmul_02[1], biasmoe1_bias, k8, k_group4, group_count8, group_select_mode1, norm_type1) # scope 外的普通算子 reshape_01 torch.reshape(moe_gating_top_k_01[1], (2, 8, 128)) ... with tng.scope.super_kernel(sk2): dequant_swiglu_quant_01 torch_npu.npu_dequant_swiglu_quant(...) quant_matmul_01 torch_npu.npu_quant_matmul(...) return quant_matmul_01 config CompilerConfig() npu_backend tng.get_npu_backend(compiler_configconfig) model torch.compile(model, fullgraphTrue, backendnpu_backend, dynamicFalse) _npu_output model(gmm1_x, gmm1_weight, gmm1_bias, gmm2_weight, moe1_bias, dsq_input, data1, data2, scale)要点解析tng.scope.super_kernel(sk1)的第一个参数是 scope 名称用于在编译日志中区分不同的 SuperKernel 融合组样例中定义了sk1GroupedMatmul GroupedMatmul MoeGatingTopK与sk2DequantSwigluQuant QuantMatmulV3两个 scope中间穿插的 reshape/square/concat 属于 scope 外算子tng.get_npu_backend(compiler_configconfig)获取 TorchAir 的 NPU 编译后端再通过torch.compile(..., fullgraphTrue, dynamicFalse)将其接入 PyTorch 编译体系fullgraphTrue要求整图编译、不允许图碎片化scope 内选择的是 MoE 场景的典型算子组合分组矩阵乘、门控 TopK、反量化 SwiGLU 量化、量化 MatMul贴近真实大模型推理负载便于观察融合收益。5.2 example02_super_kernel_profiling融合前后性能对比superkernel_compare.py 构造了包含 6 个 SuperKernel scopesp1~sp6的 MoE 网络并通过FakeContextManager实现同一网络开/关 SuperKernel两个版本def super_kernel_scope(enable_superkernel: bool, scope: str, options: str): if enable_superkernel: return tng.scope.super_kernel(scope, options) else: return FakeContextManager()随后分别对Network(False)与Network(True)做torch.compile并用torch_npu.profiler采集第二次执行的 traceexperimental_config torch_npu.profiler._ExperimentalConfig( export_type[torch_npu.profiler.ExportType.Text], profiler_leveltorch_npu.profiler.ProfilerLevel.Level1, aic_metricstorch_npu.profiler.AiCMetrics.AiCoreNone, l2_cacheFalse, op_attrFalse, record_op_argsFalse, data_simplificationFalse) with torch_npu.profiler.profile( activities[torch_npu.profiler.ProfilerActivity.CPU, torch_npu.profiler.ProfilerActivity.NPU], scheduletorch_npu.profiler.schedule(wait0, warmup1, active1, repeat1, skip_first0), on_trace_readytorch_npu.profiler.tensorboard_trace_handler(./prof_result/sk_model/), ...) as prof: for _ in range(4): sk_model torch.compile(sk_model, fullgraphTrue, backendnpu_backend, dynamicFalse) sk_model(...) prof.step()使用时的关键点schedule(wait0, warmup1, active1, ...)表示跳过 warmup 阶段只采集正式执行数据两组 trace 分别输出到./prof_result/no_sk_model/与./prof_result/sk_model/用 TensorBoard 打开即可对比算子时长、launch 开销等指标直观看到 SuperKernel 把多个算子合并为单 kernel 后的差异。注意该脚本循环内重复torch.compile的写法每次迭代重新编译是样例为演示 profiling 流程所保留的原始逻辑实际业务中应只编译一次。5.3 example03_super_kernel_runtime_ascendc_onlyAscendC Runtime 极简路径superkernel_runtime_ascendc_basic.py 是最贴近底层的样例不经过 PyTorch直接用 Python ctypes 驱动 ACL 接口把两个 AscendC 子 kernel 编译、组装、注册、下发并取回结果。整个流程分七步注释里写得很清楚初始化 ACLacl_init、aclrt_set_device(0)、创建 stream编译 subkernel 并声明复用关系核心在 compile_sk.pywith SkCompileContext(need_cleanTrue) as ctx: sub_kernels compile_subkernel(ctx) # 设置输入输出内存复用关系A 的输出就是 B 的输入 sub_kernels[0].set_input([arg_in1, arg_in2]) sub_kernels[0].set_output([out1]) sub_kernels[1].set_input([out1]) sub_kernels[1].set_output([out2]) super_kernel_result compile_superkernel(ctx, sub_kernels) super_kernel_result.output sub_kernels[1].outputset_input/set_output的目的是在 SuperKernel 阶段推导张量复用关系compile_sk.py中的注释说明比如 A-B 时A 的输出名称要设置为 B 的输入名称而 workspace 复用则与命名无关同 stream 下小 workspace 可复用大 workspace准备输入数据与内存aclrt_malloc_host分配 host 内存、aclrt_memcpy_asyncH2D 拷贝、按workspaces_size()分配 device 工作空间算子加载读取kernel_meta/name.json元数据blockDim、magic、kernelName、binFileName等字段与.o二进制构造aclrt_dev_binary并aclrt_dev_binary_register、aclrt_function_register、aclrt_get_function_by_namecompile_sk.py中维护了MAGIC_MAPPING表将元数据里的 magic 字符串如RT_DEV_BINARY_MAGIC_PLAIN_ELF映射为对应的数值魔数Launch 下发aclrt_get_c2c_ctrl_addr取 C2C 控制地址作为占位首参把各 subkernel 的输入/输出/workspace 指针依次拼进args_list最后aclrt_kernel_launch(stub_func, block_dim, args, ...)结果获取aclrt_synchronize_stream同步后 D2H 拷贝回 host 并打印浮点数组释放资源释放内存、aclrt_dev_binary_unregister、销毁 stream、重置设备。compile_sk.py还展示了子 kernel 编译时的关键构建配置current_build_config()[enable_super_kernel] 1、super_kernel_sub_loc middle、super_kernel_options early-start0等super_kernel_sub_info以及 SOC 版本解析优先级——依次检查环境变量SUPERKERNEL_COMPILE_SOC_VERSION、ASCEND_COMPILE_SOC_VERSION、ASCEND_SOC_VERSION再回退到aclrtGetSocName()动态查询默认Ascend910_9391。这解释了为什么该样例依赖asc_op_compile_base等 CANN 内部编译模块运行环境必须是完整 CANN 环境。六、AOT 样例详解AOT 样例通过 TorchAirnpugraph_ex将图静态编译为 kernel 包再由run.sh统一处理清理、日志与产物校验。6.1 example01_dual_stream双流与 NPU Event 控制边该样例使用两个 NPU stream 和 event 建立控制依赖并通过auto_op_parallel启用 SuperKernel 自动算子并行分别运行启用和关闭 SuperKernel 的网络并检查两侧输出一致性验证融合不改变计算语义。6.2 example02_sk_optionsSuperKernel options 调优该样例展示 SuperKernel 的 optimize/debug options 用法并使用 eager 输出校验静态编译结果的正确性。目录内按架构拆分了入口 main-dav-2201.py 与 main-dav-3510.pyrun.sh会根据--npu-arch选择对应架构的 attention 网络运行——这体现了 README 所说样例内部的编译参数由各样例处理不同架构的差异下沉到样例内部入口脚本只做架构分发。6.3 example03_kernel_pybindPybind 自定义算子融合该样例演示自定义算子进入 SuperKernel 融合的完整链路使用bisheng编译带SK_BIND标注的 AscendC custom add kernel源码位于 ops/aclgraph_add_ops/csrc/add_custom.asc 与 pybind11.asc通过 pybind 注册为 PyTorch 算子op_extension 包再由npugraph_ex静态编译并启用 SuperKernel。其安装流程有若干工程细节值得注意来自样例 READMErun.sh将架构参数显式传给扩展安装脚本安装脚本仅为本次 Python 构建设置SK_NPU_ARCH供 setup.py 配置bisheng --npu-arch扩展通过当前 Python支持 venv安装到样例的tmp/python_packages经PYTHONPATH导入退出时仅清理该目录不卸载当前环境中的同名包——这一设计保证样例运行不会污染用户的 Python 环境。七、AOT 退出清理与静态 kernel 安全卸载README 特别说明了 AOT 样例的退出清理策略其实现位于 common.sh 的sk_uninstall_static_kernel_from_log函数各 AOT 样例的run.sh通过trap ... EXIT挂载该逻辑。其安全约束可以概括为从运行日志中用正则提取形如/路径/uninstall.sh的候选卸载脚本仅接受本次编译产物对应脚本名在本次static_kernel_compile_outputs下生成的*.run包集合内、位于当前 CANN 安装目录$ASCEND_HOME_PATH/opp/static_kernel/ai_core下的预期路径且通过路径校验不含..、非符号链接、resolve()后路径完全一致、是普通文件的uninstall.sh日志中的其他路径一律拒绝执行并输出WARN: refusing unrelated or redirected uninstall script告警清理失败同样只输出告警而不中断。配套函数还有sk_cleanup_local运行前清理log、tmp、sk_meta、kernel_meta、profiling、static_kernel_compile_outputs等目录并重建tmp/、log/、sk_run_python_with_log以tee双路输出运行日志并用PIPESTATUS[0]取 Python 真实退出码、sk_check_static_kernel_outputs检查*_compile_error.log与*.run产物编译失败或缺少必需产物时报错退出。八、参数速查与运行前提参数/配置取值/说明出处--npu-arch必填dav-2201或dav-3510缺省或非法值报错退出common.sh 中sk_parse_npu_arch/sk_validate_npu_arch-h/--help打印用法Usage: bash $0 --npu-archdav-2201\|dav-3510同上JIT 样例支持dav-2201全量运行dav-3510跳过并打印 INFOrun_example.shAOT 样例支持三个 AOT 样例均支持dav-2201与dav-3510各样例 README 声明各aot/*/README.mdPython 解释器选择VIRTUAL_ENVvenv/bin/pythonpython3run_example.shPython 依赖protobuf3.13,4、torch2.6.0、torch_npu2.6.0.post5、PyYAML5.1requirements.txtSK_NPU_ARCH仅 example03 安装流程使用供setup.py配置bisheng --npu-archexample03 READMESOC 版本解析环境变量SUPERKERNEL_COMPILE_SOC_VERSION/ASCEND_COMPILE_SOC_VERSION/ASCEND_SOC_VERSIONaclrtGetSocName() 默认Ascend910_9391compile_sk.py运行前提归纳当前机器有可见的 Ascend NPU且已按 源码构建指南 完成环境准备JIT 样例还要求torch_npu与 CANN 环境匹配必须显式传入--npu-arch样例内部不做架构探测统一入口run_example.sh当前版本因 CI 环境限制暂未启用 AOT 样例调用AOT 样例需使用各自的run.sh独立运行SuperKernel 各 option 的完整含义以 TorchAir 官方《SuperKernel 使用说明》文档为准原文档给出的参考链接指向 TorchAir 仓库的 npugraph_ex 文档与昇腾社区 Ascend Extension for PyTorch 用户指南此处不重复外链。九、延伸阅读各样例的中英文说明文档如 JIT 基础样例、profiling 样例、Runtime 样例、sk_options 样例SuperKernel 组件整体设计与开发指南可继续参阅仓库中的 developer_guide.md从源码结构看examples中的 JIT 接口tng.scope.super_kernel与 AOT 静态编译最终都汇聚到 SuperKernel 的 scope 划分与 kernel 组装机制上compile_sk.py中super_kernel.compile(kernel_info, kernel_name)的op_list含stream_id、bin_path、json_path即是多子 kernel 组装为单个 super kernel 的直接输入形式理解它对自定义融合场景有直接参考价值。【免费下载链接】graph-autofusionGraph-autofusion 是一个面向昇腾Ascend芯片的轻量级、解耦式组件集合旨在通过自动融合技术加速模型执行。 目前已开源 SuperKernel 组件和 Autofuse 组件未来将持续开放更多自动融合相关模块。项目地址: https://gitcode.com/cann/graph-autofusion创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 CANN SuperKernel 样例实战指南：JIT 与 AOT 两类 Python 样例的运行、编译与调优
CANN SuperKernel 样例实战指南JIT 与 AOT 两类 Python 样例的运行、编译与调优【免费下载链接】graph-autofusionGraph-autofusion 是一个面向昇腾Ascend芯片的轻量级、解耦式组件集合旨在通过自动融合技术加速模型执行。 目前已开源 SuperKernel 组件和 Autofuse 组件未来将持续开放更多自动融合相关模块。项目地址: https://gitcode.com/cann/graph-autofusion本篇指南围绕 super_kernel/examples/README.md 展开系统讲解 CANN 开源仓库中 SuperKernel 组件的两类 Python 样例——基于 JIT 接口的运行时融合样例与基于 TorchAirnpugraph_ex的 AOT 静态编译样例并结合 run_example.sh、common.sh 及各样例源码说明--npu-arch架构选择逻辑、JIT/AOT 样例的具体实现路径、SuperKernel 静态 kernel 的退出清理机制帮助你在 Ascend NPU 上完整跑通、复现并改造这些样例。一、样例体系概览JIT 与 AOT 两条技术路线super_kernel/examples/目录提供两类 SuperKernel Python 样例分别对应 SuperKernel 的两种落地方式jit/通过 SuperKernel JIT 接口完成融合、编译和运行。样例直接在 Python 代码中用tng.scope.super_kernel上下文管理器标注融合范围配合torch.compile走 TorchAir 的 NPU 后端编译执行适合在训练/推理脚本中快速验证融合效果。aot/通过 TorchAirnpugraph_ex静态编译并启用 SuperKernel 优化。样例将图以静态方式编译为 kernel 包再加载执行覆盖双流控制边、SuperKernel options 调优、自定义 AscendC 算子融合等更完整的工程场景。两类样例统一由仓库根目录下的super_kernel/examples/run_example.sh入口脚本驱动公共 Bash 函数集中放在 aot/scripts/common.sh 中。二、目录结构examples/ ├── run_example.sh # 样例统一运行入口 ├── jit/ │ ├── example01_super_kernel_base/ # SuperKernel 基础功能 │ ├── example02_super_kernel_profiling/ # SuperKernel profiling 对比 │ └── example03_super_kernel_runtime_ascendc_only/ # AscendC Runtime 极简样例 └── aot/ ├── scripts/ # AOT 样例公共 Bash 函数 ├── example01_dual_stream/ # 双流与 NPU Event 控制边 ├── example02_sk_options/ # SuperKernel options └── example03_kernel_pybind/ # Pybind 自定义算子融合各样例目录下均提供中英文README.md可作为深入阅读的起点例如 JIT 基础样例说明、双流样例说明、Pybind 样例说明。三、前置条件在运行样例前需要完成两项准备环境构建先参考 源码构建指南仓库文档路径docs/zh/build.md对应文档中的相对链接../../docs/zh/build.md完成 CANN 环境准备确保当前机器有可见的 Ascend NPU。所有样例均使用当前可见 NPU代码中通过torch.npu.set_device(0)或aclrt_set_device(DEVICE_ID)指定为 0 号设备。Python 依赖安装 requirements.txt 中声明的依赖protobuf3.13,4 torch2.6.0 torch_npu2.6.0.post5 PyYAML5.1其中torch/torch_npu版本对齐是 JIT 样例能正常torch.compile到 NPU 后端的关键前提。四、统一运行入口与--npu-arch架构选择样例运行必须显式传入--npu-arch参数用于选择适用当前硬件的样例集合bash super_kernel/examples/run_example.sh --npu-archdav-2201 bash super_kernel/examples/run_example.sh --npu-archdav-3510从入口脚本 run_example.sh 的源码可以看出其执行逻辑脚本首先解析 Python 解释器优先使用VIRTUAL_ENV环境其次回退到仓库根目录的venv/bin/python最后回退到系统python3通过common.sh中的sk_parse_npu_arch解析参数该函数用getopt仅接受-h/--help与--npu-arch两个参数未传--npu-arch会报ERROR: --npu-arch is required.取值非法既非dav-2201也非dav-3510会报unsupported NPU architecture并以退出码 2 失败随后按架构分支执行样例。架构差异说明这一点在源码中体现得很直接dav-2201会依次运行全部三个 JIT 样例superkernel_scope.py、superkernel_compare.py、superkernel_runtime_ascendc_basic.pydav-3510不支持这些 JIT 样例脚本打印[INFO] Skipping SuperKernel JIT examples on dav-3510.并跳过仅运行 AOT 样例。一个需要注意的现状当前仓库版本中run_example.sh对三个 AOT 样例的调用行处于注释状态并附有注释暂时禁用因当前 CI 环境限制无法运行的 AOT examples。因此在统一入口下dav-3510分支目前实际会打印跳过信息后直接进入Run all examples success。如需运行 AOT 样例应改用各样例自带的独立入口见第五节例如bash super_kernel/examples/aot/example01_dual_stream/run.sh --npu-archdav-2201此外run_example.sh中的 JIT/AOT 分支逻辑印证了 README 中样例内部的编译参数由各样例处理的说明入口脚本只负责架构分流和解释器选择真正的编译配置SOC 版本、options 等由各样例内部完成。五、JIT 样例详解5.1 example01_super_kernel_baseSuperKernel 基础功能入口文件 superkernel_scope.py 演示了 SuperKernel 最核心的使用范式——用tng.scope.super_kernel上下文管理器在nn.Module.forward中标注融合范围import torchair as tng from torchair.configs.compiler_config import CompilerConfig class Network(nn.Module): def forward(self, gmm1_x, gmm1_weight, gmm1_bias, gmm2_weight, moe1_bias, dsq_input, data1, data2, scale): with tng.scope.super_kernel(sk1): grouped_matmul_01 torch_npu.npu_grouped_matmul(group_type-1, xgmm1_x, weightgmm1_weight, biasgmm1_bias) grouped_matmul_02 torch_npu.npu_grouped_matmul(group_type-1, xgrouped_matmul_01, weightgmm2_weight) moe_gating_top_k_01 torch_npu.npu_moe_gating_top_k(xgrouped_matmul_02[1], biasmoe1_bias, k8, k_group4, group_count8, group_select_mode1, norm_type1) # scope 外的普通算子 reshape_01 torch.reshape(moe_gating_top_k_01[1], (2, 8, 128)) ... with tng.scope.super_kernel(sk2): dequant_swiglu_quant_01 torch_npu.npu_dequant_swiglu_quant(...) quant_matmul_01 torch_npu.npu_quant_matmul(...) return quant_matmul_01 config CompilerConfig() npu_backend tng.get_npu_backend(compiler_configconfig) model torch.compile(model, fullgraphTrue, backendnpu_backend, dynamicFalse) _npu_output model(gmm1_x, gmm1_weight, gmm1_bias, gmm2_weight, moe1_bias, dsq_input, data1, data2, scale)要点解析tng.scope.super_kernel(sk1)的第一个参数是 scope 名称用于在编译日志中区分不同的 SuperKernel 融合组样例中定义了sk1GroupedMatmul GroupedMatmul MoeGatingTopK与sk2DequantSwigluQuant QuantMatmulV3两个 scope中间穿插的 reshape/square/concat 属于 scope 外算子tng.get_npu_backend(compiler_configconfig)获取 TorchAir 的 NPU 编译后端再通过torch.compile(..., fullgraphTrue, dynamicFalse)将其接入 PyTorch 编译体系fullgraphTrue要求整图编译、不允许图碎片化scope 内选择的是 MoE 场景的典型算子组合分组矩阵乘、门控 TopK、反量化 SwiGLU 量化、量化 MatMul贴近真实大模型推理负载便于观察融合收益。5.2 example02_super_kernel_profiling融合前后性能对比superkernel_compare.py 构造了包含 6 个 SuperKernel scopesp1~sp6的 MoE 网络并通过FakeContextManager实现同一网络开/关 SuperKernel两个版本def super_kernel_scope(enable_superkernel: bool, scope: str, options: str): if enable_superkernel: return tng.scope.super_kernel(scope, options) else: return FakeContextManager()随后分别对Network(False)与Network(True)做torch.compile并用torch_npu.profiler采集第二次执行的 traceexperimental_config torch_npu.profiler._ExperimentalConfig( export_type[torch_npu.profiler.ExportType.Text], profiler_leveltorch_npu.profiler.ProfilerLevel.Level1, aic_metricstorch_npu.profiler.AiCMetrics.AiCoreNone, l2_cacheFalse, op_attrFalse, record_op_argsFalse, data_simplificationFalse) with torch_npu.profiler.profile( activities[torch_npu.profiler.ProfilerActivity.CPU, torch_npu.profiler.ProfilerActivity.NPU], scheduletorch_npu.profiler.schedule(wait0, warmup1, active1, repeat1, skip_first0), on_trace_readytorch_npu.profiler.tensorboard_trace_handler(./prof_result/sk_model/), ...) as prof: for _ in range(4): sk_model torch.compile(sk_model, fullgraphTrue, backendnpu_backend, dynamicFalse) sk_model(...) prof.step()使用时的关键点schedule(wait0, warmup1, active1, ...)表示跳过 warmup 阶段只采集正式执行数据两组 trace 分别输出到./prof_result/no_sk_model/与./prof_result/sk_model/用 TensorBoard 打开即可对比算子时长、launch 开销等指标直观看到 SuperKernel 把多个算子合并为单 kernel 后的差异。注意该脚本循环内重复torch.compile的写法每次迭代重新编译是样例为演示 profiling 流程所保留的原始逻辑实际业务中应只编译一次。5.3 example03_super_kernel_runtime_ascendc_onlyAscendC Runtime 极简路径superkernel_runtime_ascendc_basic.py 是最贴近底层的样例不经过 PyTorch直接用 Python ctypes 驱动 ACL 接口把两个 AscendC 子 kernel 编译、组装、注册、下发并取回结果。整个流程分七步注释里写得很清楚初始化 ACLacl_init、aclrt_set_device(0)、创建 stream编译 subkernel 并声明复用关系核心在 compile_sk.pywith SkCompileContext(need_cleanTrue) as ctx: sub_kernels compile_subkernel(ctx) # 设置输入输出内存复用关系A 的输出就是 B 的输入 sub_kernels[0].set_input([arg_in1, arg_in2]) sub_kernels[0].set_output([out1]) sub_kernels[1].set_input([out1]) sub_kernels[1].set_output([out2]) super_kernel_result compile_superkernel(ctx, sub_kernels) super_kernel_result.output sub_kernels[1].outputset_input/set_output的目的是在 SuperKernel 阶段推导张量复用关系compile_sk.py中的注释说明比如 A-B 时A 的输出名称要设置为 B 的输入名称而 workspace 复用则与命名无关同 stream 下小 workspace 可复用大 workspace准备输入数据与内存aclrt_malloc_host分配 host 内存、aclrt_memcpy_asyncH2D 拷贝、按workspaces_size()分配 device 工作空间算子加载读取kernel_meta/name.json元数据blockDim、magic、kernelName、binFileName等字段与.o二进制构造aclrt_dev_binary并aclrt_dev_binary_register、aclrt_function_register、aclrt_get_function_by_namecompile_sk.py中维护了MAGIC_MAPPING表将元数据里的 magic 字符串如RT_DEV_BINARY_MAGIC_PLAIN_ELF映射为对应的数值魔数Launch 下发aclrt_get_c2c_ctrl_addr取 C2C 控制地址作为占位首参把各 subkernel 的输入/输出/workspace 指针依次拼进args_list最后aclrt_kernel_launch(stub_func, block_dim, args, ...)结果获取aclrt_synchronize_stream同步后 D2H 拷贝回 host 并打印浮点数组释放资源释放内存、aclrt_dev_binary_unregister、销毁 stream、重置设备。compile_sk.py还展示了子 kernel 编译时的关键构建配置current_build_config()[enable_super_kernel] 1、super_kernel_sub_loc middle、super_kernel_options early-start0等super_kernel_sub_info以及 SOC 版本解析优先级——依次检查环境变量SUPERKERNEL_COMPILE_SOC_VERSION、ASCEND_COMPILE_SOC_VERSION、ASCEND_SOC_VERSION再回退到aclrtGetSocName()动态查询默认Ascend910_9391。这解释了为什么该样例依赖asc_op_compile_base等 CANN 内部编译模块运行环境必须是完整 CANN 环境。六、AOT 样例详解AOT 样例通过 TorchAirnpugraph_ex将图静态编译为 kernel 包再由run.sh统一处理清理、日志与产物校验。6.1 example01_dual_stream双流与 NPU Event 控制边该样例使用两个 NPU stream 和 event 建立控制依赖并通过auto_op_parallel启用 SuperKernel 自动算子并行分别运行启用和关闭 SuperKernel 的网络并检查两侧输出一致性验证融合不改变计算语义。6.2 example02_sk_optionsSuperKernel options 调优该样例展示 SuperKernel 的 optimize/debug options 用法并使用 eager 输出校验静态编译结果的正确性。目录内按架构拆分了入口 main-dav-2201.py 与 main-dav-3510.pyrun.sh会根据--npu-arch选择对应架构的 attention 网络运行——这体现了 README 所说样例内部的编译参数由各样例处理不同架构的差异下沉到样例内部入口脚本只做架构分发。6.3 example03_kernel_pybindPybind 自定义算子融合该样例演示自定义算子进入 SuperKernel 融合的完整链路使用bisheng编译带SK_BIND标注的 AscendC custom add kernel源码位于 ops/aclgraph_add_ops/csrc/add_custom.asc 与 pybind11.asc通过 pybind 注册为 PyTorch 算子op_extension 包再由npugraph_ex静态编译并启用 SuperKernel。其安装流程有若干工程细节值得注意来自样例 READMErun.sh将架构参数显式传给扩展安装脚本安装脚本仅为本次 Python 构建设置SK_NPU_ARCH供 setup.py 配置bisheng --npu-arch扩展通过当前 Python支持 venv安装到样例的tmp/python_packages经PYTHONPATH导入退出时仅清理该目录不卸载当前环境中的同名包——这一设计保证样例运行不会污染用户的 Python 环境。七、AOT 退出清理与静态 kernel 安全卸载README 特别说明了 AOT 样例的退出清理策略其实现位于 common.sh 的sk_uninstall_static_kernel_from_log函数各 AOT 样例的run.sh通过trap ... EXIT挂载该逻辑。其安全约束可以概括为从运行日志中用正则提取形如/路径/uninstall.sh的候选卸载脚本仅接受本次编译产物对应脚本名在本次static_kernel_compile_outputs下生成的*.run包集合内、位于当前 CANN 安装目录$ASCEND_HOME_PATH/opp/static_kernel/ai_core下的预期路径且通过路径校验不含..、非符号链接、resolve()后路径完全一致、是普通文件的uninstall.sh日志中的其他路径一律拒绝执行并输出WARN: refusing unrelated or redirected uninstall script告警清理失败同样只输出告警而不中断。配套函数还有sk_cleanup_local运行前清理log、tmp、sk_meta、kernel_meta、profiling、static_kernel_compile_outputs等目录并重建tmp/、log/、sk_run_python_with_log以tee双路输出运行日志并用PIPESTATUS[0]取 Python 真实退出码、sk_check_static_kernel_outputs检查*_compile_error.log与*.run产物编译失败或缺少必需产物时报错退出。八、参数速查与运行前提参数/配置取值/说明出处--npu-arch必填dav-2201或dav-3510缺省或非法值报错退出common.sh 中sk_parse_npu_arch/sk_validate_npu_arch-h/--help打印用法Usage: bash $0 --npu-archdav-2201\|dav-3510同上JIT 样例支持dav-2201全量运行dav-3510跳过并打印 INFOrun_example.shAOT 样例支持三个 AOT 样例均支持dav-2201与dav-3510各样例 README 声明各aot/*/README.mdPython 解释器选择VIRTUAL_ENVvenv/bin/pythonpython3run_example.shPython 依赖protobuf3.13,4、torch2.6.0、torch_npu2.6.0.post5、PyYAML5.1requirements.txtSK_NPU_ARCH仅 example03 安装流程使用供setup.py配置bisheng --npu-archexample03 READMESOC 版本解析环境变量SUPERKERNEL_COMPILE_SOC_VERSION/ASCEND_COMPILE_SOC_VERSION/ASCEND_SOC_VERSIONaclrtGetSocName() 默认Ascend910_9391compile_sk.py运行前提归纳当前机器有可见的 Ascend NPU且已按 源码构建指南 完成环境准备JIT 样例还要求torch_npu与 CANN 环境匹配必须显式传入--npu-arch样例内部不做架构探测统一入口run_example.sh当前版本因 CI 环境限制暂未启用 AOT 样例调用AOT 样例需使用各自的run.sh独立运行SuperKernel 各 option 的完整含义以 TorchAir 官方《SuperKernel 使用说明》文档为准原文档给出的参考链接指向 TorchAir 仓库的 npugraph_ex 文档与昇腾社区 Ascend Extension for PyTorch 用户指南此处不重复外链。九、延伸阅读各样例的中英文说明文档如 JIT 基础样例、profiling 样例、Runtime 样例、sk_options 样例SuperKernel 组件整体设计与开发指南可继续参阅仓库中的 developer_guide.md从源码结构看examples中的 JIT 接口tng.scope.super_kernel与 AOT 静态编译最终都汇聚到 SuperKernel 的 scope 划分与 kernel 组装机制上compile_sk.py中super_kernel.compile(kernel_info, kernel_name)的op_list含stream_id、bin_path、json_path即是多子 kernel 组装为单个 super kernel 的直接输入形式理解它对自定义融合场景有直接参考价值。【免费下载链接】graph-autofusionGraph-autofusion 是一个面向昇腾Ascend芯片的轻量级、解耦式组件集合旨在通过自动融合技术加速模型执行。 目前已开源 SuperKernel 组件和 Autofuse 组件未来将持续开放更多自动融合相关模块。项目地址: https://gitcode.com/cann/graph-autofusion创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考