graph-autofusion SuperKernel JIT 基础示例:用 torchair.scope.super_kernel 将 MoE 算子融合为超级 Kernel

graph-autofusion SuperKernel JIT 基础示例:用 torchair.scope.super_kernel 将 MoE 算子融合为超级 Kernel graph-autofusion SuperKernel JIT 基础示例用 torchair.scope.super_kernel 将 MoE 算子融合为超级 Kernel【免费下载链接】graph-autofusionGraph-autofusion 是一个面向昇腾Ascend芯片的轻量级、解耦式组件集合旨在通过自动融合技术加速模型执行。 目前已开源 SuperKernel 组件和 Autofuse 组件未来将持续开放更多自动融合相关模块。项目地址: https://gitcode.com/cann/graph-autofusion本篇围绕 SuperKernel JIT 基础示例 展开演示如何通过 TorchAir 的super_kernel作用域将多个算子融合为一个超级 Kernel 并在昇腾 NPU 上编译执行。读完后你可以掌握torchair.scope.super_kernel的标定语法与使用边界、示例脚本中模型输入与算子流的完整构造方式、torch.compile接入 NPU 后端的执行链路以及从仓库源码角度理解多个子算子如何被合并为单一 Kernel 并插入同步/预取等编译优化的底层机制。示例功能与模型结构该示例jit/example01_super_kernel_base的核心功能是使用super_kernel范围scope将sk1 片段内的 GroupedMatmul GroupedMatmul MoeGatingTopK 三个算子融合为一个超级 Kernel 进行计算覆盖 MoEMixture of Experts模型中典型的专家计算与门控 TopK 选择链路。从 示例脚本 中的注释图可以看出完整模型实际包含两个super_kernel 片段结构如下|o data |o | |o sk1:GroupedMatmulGroupedMatmulMoeGatingTopK |o | |o reshape-square-concat |o | |o sk2:DequantSwigluQuantQuantBatchMatmulV3 |o | |o netoutput其中sk1两组 GroupedMatmul分组矩阵乘串联后接 MoeGatingTopK即 README 中声明的三算子融合片段中间段不融合reshape → square → concat → reshape四个 PyTorch 原生算子保持在 super_kernel 之外独立调度sk2DequantSwigluQuant反量化 SwiGLU 激活 再量化与 QuantMatmul量化矩阵乘脚本中对应npu_quant_matmul两个量化链路算子的融合片段。这种融合段 非融合段交替的组织方式体现了 super_kernel 的典型使用形态按需标定融合范围而非整图融合。环境准备与依赖运行本示例前需按 SuperKernel 样例使用指导 完成环境准备Python 依赖安装 requirements.txt 中的依赖核心版本约束为protobuf3.13,4 torch2.6.0 torch_npu2.6.0.post5 PyYAML5.1即需要torch 2.6.0与匹配的torch_npu示例代码依赖torchairTorchAir 编译器前端与torchair.configs.compiler_config中的CompilerConfig。硬件架构JIT 类 SuperKernel 示例目前适用于dav-2201架构的 NPU统一运行入口 run_example.sh 中dav-2201会依次运行全部 JIT 示例而dav-3510会跳过 JIT 示例仅运行 AOT 示例因此运行本示例前请确认当前可见 NPU 属于dav-2201系列。解释器选择run_example.sh会优先使用已激活的VIRTUAL_ENV或仓库根目录下的venv/bin/python否则回退到系统python3见 run_example.sh。super_kernel 作用域语法如何标定融合范围示例使用的核心语法是一个with语句块with torchair.scope.super_kernel(sk1): # 块内所有算子均被融合为一个超级 Kernel 进行计算字符串参数sk1是融合片段的命名标识脚本中另一个片段命名为sk2。作用域块内的所有算子在 TorchAir 捕获计算图时会被标记为同一个 SuperKernel 范围最终由编译器合并为一个__global__设备函数。该标定方式属于 TorchAir 图内 SuperKernel 范围配置与模型代码以普通 Python 控制流混写无需单独构造融合图。下面给出示例中两个作用域的原始代码摘自 superkernel_scope.pydef 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) reshape_01 torch.reshape(moe_gating_top_k_01[1], (2, 8, 128)) square_01 torch.square(reshape_01) concat_01 torch.cat((square_01, data1), dim0) reshape_02 torch.reshape(concat_01, (-1, 128)) with tng.scope.super_kernel(sk2): dequant_swiglu_quant_01 torch_npu.npu_dequant_swiglu_quant(xreshape_02, weight_scaledsq_input[0], activation_scaledsq_input[1], biasNone, quant_scaledsq_input[2], quant_offsetNone, group_indexNone, activate_leftFalse, quant_mode1) quant_matmul_01 torch_npu.npu_quant_matmul(x1dequant_swiglu_quant_01[0], x2data2, scalescale, offsetNone, biasNone, pertoken_scaleNone, output_dtypetorch.float16) return quant_matmul_01几个值得注意的实现细节列表型输入npu_grouped_matmul的x、weight、bias均为张量列表分组输入sk1 中第一组形状为(64,128)×(128,64)第二组为(256,180)×(180,320)多输出算子的取用npu_grouped_matmul返回多输出元组第二个 GroupedMatmul 显式取grouped_matmul_01整体传入MoeGatingTopK 取grouped_matmul_02[1]低精度/整型混合链路sk2 链路中data2为int8张量(6,64,64)、scale为int64张量体现了量化算子融合的典型场景量化 scale 张量均为float32。完整示例代码逐段解析superkernel_scope.py 的完整执行流程可分为四步。第 1 步设定设备并构造输入脚本 L37-L65。脚本通过torch.npu.set_device(0)绑定 0 号 NPU随后用np.random.uniform在 CPU 上生成随机数经torch.from_numpy(...).to(dtype).npu()搬运至设备。各输入的语义与形状如下输入变量形状类型用途gmm1_x2 个(64,128)、(256,180)fp16第一组 GroupedMatmul 分组激活gmm1_weight2 个(128,64)、(180,320)fp16第一组 GroupedMatmul 分组权重gmm1_bias2 个(64,)、(320,)fp16第一组 GroupedMatmul 偏置gmm2_weight2 个(64,56)、(320,256)fp16第二组 GroupedMatmul 权重moe1_bias(256,)fp16MoeGatingTopK 偏置dsq_input3 个(128,)、(48,1)、(1,64)fp32DequantSwigluQuant 的 weight/activation/quant scaledata1(4,8,128)int32中间段 concat 的第二个输入data2(6,64,64)int8量化 MatMul 的第二个输入scale(1,)int64量化 MatMul 缩放因子第 2 步定义带 super_kernel 作用域的模型脚本 L68-L91即上一节展示的Network类forward内以两个with tng.scope.super_kernel(...)块标定 sk1/sk2 融合范围。第 3 步接入 TorchAir NPU 后端编译脚本 L93-L98config CompilerConfig() npu_backend tng.get_npu_backend(compiler_configconfig) model Network().npu() # 在 npu 上执行有 superkernel 配置的模型 model torch.compile(model, fullgraphTrue, backendnpu_backend, dynamicFalse)这里通过torchair.configs.compiler_config.CompilerConfig构造编译器配置tng.get_npu_backend返回 TorchAir NPU 后端再以fullgraphTrue、dynamicFalse交给torch.compile做整图捕获与编译——super_kernel 作用域正是在这一编译阶段被解析并生成融合 Kernel。第 4 步执行并校验脚本 L99-L102_npu_output model(gmm1_x, gmm1_weight, gmm1_bias, gmm2_weight, moe1_bias, dsq_input, data1, data2, scale) print(execute sample success)执行命令与预期结果在示例目录下直接运行脚本python3 superkernel_scope.py执行成功后终端打印execute sample success也可以通过 统一运行入口 运行需显式指定 NPU 架构bash super_kernel/examples/run_example.sh --npu-archdav-2201run_example.sh按顺序执行三个 JIT 示例example01_super_kernel_base本示例、example02_super_kernel_profiling、example03_super_kernel_runtime_ascendc_only全部通过后打印Run all examples success。源码级原理super_kernel 编译链路做了什么理解示例为什么能一个作用域 一个 Kernel可以参考仓库中 JIT 侧的编译实现 super_kernel/src/jit/superkernel/super_kernel.py。编译入口。所有子算子的合并由compile(kernel_infos, called_kernel_name, compile_infos)触发super_kernel.py#L1044-L1075输入是子算子的op_list每个子算子提供bin_path与json_path以及super_kernel_options编译选项入口先重置全局变量存储、初始化特性管理器校验当前 SoC 是否支持 SuperKernel随后构建SuperOperatorInfos并调用gen_super_kernel_file生成融合后的 AscendC 源文件再交给compile_super_kernel完成设备侧编译。融合后的单 Kernel 结构。从gen_super_kernel_file的代码生成逻辑super_kernel.py#L904-L1041可以确认生成的融合 Kernel 是一个extern C __global__ __aicore__函数其内部按子算子顺序完成四件事插入算子间同步屏障gen_sync_and_event_code在前序算子与当前算子之间生成屏障代码当 Early-Start 关闭时走gen_inter_ops_barrier按超级 Kernel 的类型选择同步原语super_kernel.py#L136-L150——Mix 1:1/1:2 类型使用AscendC::SyncAllfalse()纯 AIC 类型使用ffts_cross_core_sync(PIPE_FIX, ...SYNC_AIC_FLAG)AIV 类型使用ffts_cross_core_sync(PIPE_MTE3, ...SYNC_AIV_ONLY_ALL)。这正是 SuperKernel 原理介绍 所述按前后子算子 Kernel Type 定制同步范围的实现落点指令预取ICache Preload按preload_mode三档生成不同代码——整段预取时生成AscendC::PreLoad(8)逐级预取时在当前函数前插入preload_call_block提前一级预取时同时预取当前与下一个子算子super_kernel.py#L972-L993预取列表长度按核类型区分AIV 侧为 8 段、AIC 侧为 16 段每段 2KB见 super_kernel.py#L210-L213Early-Start 优化开启时由gen_early_start_config根据前序/当前子算子的设备类型AIC/AIV/MIX编码 2 位配置位并生成g_super_kernel_early_start_configsuper_kernel.py#L57-L87配合前后子算子的 Set/Wait 同步点实现跨算子指令级并发可选的子 Kernel 拆分与 Profilingsplit_mode 1时按split{i}生成多份函数地址变量以支持子 Kernel 代码副本profiling开启时在每个子算子调用前后插入RecordProfiling(index, 0x8, start/stop)打点super_kernel.py#L1012-L1024为 example02 profiling 对比示例 提供逐子算子耗时数据。编译选项体系。示例脚本本身未显式传选项使用编译器默认值但 JIT 管线支持的选项集可从 super_kernel_option_parse.py 中注册的解析器完整枚举选项取值说明对应源码行为func-align2 的幂或 0函数对齐大小默认 512 字节经CodeTextAlignParser校验preload-codemax/none/per-func对应PreLoadByWhole/PreloadNA/PreloadByAdanvanceStep三档指令预取early-start0/1/2关闭 / EnableV2 / V2 禁用子 Kernel 拆分stream-fusion0/1双流融合开关开启时走gen_2_real_stream_super_kernel_file生成 AIC/AIV 双路代码feed-sync-all0/1是否启用跨核自动 SyncAll 的 workspace 初始化与计数逻辑profiling0/1是否生成RecordProfiling打点split-mode(0, 64] 整数子 Kernel 拆分的副本数debug-dcci-all/debug-sync-all0/1调试用每算子结束后 DCCI 逐出 / 强制全同步debug-aic-num/debug-aiv-num非负整数调试用核数配置compile-options、strict-scope-check、dcci-before-kernel-start、dcci-after-kernel-end、dcci-disable-on-kernel非空字符串附加编译控制项选项字符串以keyvalue:keyvalue形式传入并逐键校验parse_super_kernel_optionssuper_kernel_option_parse.py#L171-L197非法取值会直接抛出编译错误。AOT 场景下这些选项通过 TorchAir 的 SuperKernel options 传递其用法见 AOT 示例二。相关示例延伸本示例是 JIT 系列的基础功能验证同一目录下的其他示例可按需深入example02_super_kernel_profiling构造含 6 个 sk 片段部分复用缓存、部分在线编译的模型分别以启用/未启用 super_kernel 两种配置编译在prof_result目录下产出两份op_statistic.csv按OP_Type / Core Type / Total Time对比融合前后耗时文档给出的对比结果为约 5.61% 的收益example03_super_kernel_runtime_ascendc_onlyAscendC Runtime 的极简样例aot/目录下的 双流与 Event 控制边示例、SuperKernel options 示例、Pybind 自定义算子融合示例 则演示通过npugraph_ex静态编译路径启用 SuperKernel。小结该示例以最短路径展示了 SuperKernel 在 JIT 路径下的完整用法在torch.compile捕获的模型中用with torchair.scope.super_kernel(sk1):标定融合范围块内算子GroupedMatmul ×2 MoeGatingTopK将被合并为单一超级 Kernel块外算子保持独立调度执行成功以execute sample success为验收标准。结合仓库中 super_kernel.py 与 super_kernel_option_parse.py 的源码可以进一步理解融合并非简单串联而是在单 Kernel 内按子算子插入类型感知的同步屏障、分档的 ICache 预取、Early-Start 并发以及可选的 profiling 打点与子 Kernel 拆分这些编译期先验优化正是 SuperKernel 相对普通算子调度的核心增量。【免费下载链接】graph-autofusionGraph-autofusion 是一个面向昇腾Ascend芯片的轻量级、解耦式组件集合旨在通过自动融合技术加速模型执行。 目前已开源 SuperKernel 组件和 Autofuse 组件未来将持续开放更多自动融合相关模块。项目地址: https://gitcode.com/cann/graph-autofusion创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考