已合并
add torchair super kernel docs #3124
hujiawen_kaven创建于 5月30日
add torchair super kernel docs #3124
已合并
共 7 个文件变更+153-0
| @@ -0,0 +1,151 @@ | |||
| 1 | +# SuperKernel融合优化功能 | ||
| 2 | + | ||
| 3 | +## 功能简介 | ||
| 4 | + | ||
| 5 | +SuperKernel是一种算子二进制融合技术。与源码融合不同,它聚焦于内核函数(Kernel)的二进制调度方案优化,在已编译的二进制代码基础上融合创建一个超级Kernel函数(简称SuperKernel),以子函数调用的方式整合多个内核函数,从而优化计算任务、提升性能和资源利用率。 | ||
| 6 | + | ||
| 7 | +与单算子下发相比,SuperKernel技术能够优化任务调度的等待时间和调度开销,并可利用Task间隙资源进一步优化算子头开销。 | ||
| 8 | + | ||
| 9 | +SuperKernel的实现原理如下: | ||
| 10 | + | ||
| 11 | +基于aclgraph捕获后的Model,SuperKernel模块识别可被融合的连续Task,将其合并为一个SuperKernel Task并替换原有Task,通过Runtime接口更新Model后基于新Model执行。 | ||
| 12 | + | ||
| 13 | +开启SuperKernel融合优化后,系统自动识别图内可被融合的算子,在SuperKernel内以子函数调用的方式依次执行。同时提供标定SuperKernel范围的能力,支持用户根据实际业务需求对融合范围内的算子进行标记和优化配置。 | ||
| 14 | + | ||
| 15 | +## 使用约束 | ||
| 16 | + | ||
| 17 | +- 本功能支持如下产品: | ||
| 18 | + | ||
| 19 | + - Atlas A3 训练系列产品/Atlas A3 推理系列产品 | ||
| 20 | + - Atlas A2 训练系列产品/Atlas A2 推理系列产品 | ||
| 21 | + | ||
| 22 | +- 需要注意的是,SuperKernel融合会按照网络中算子的顺序依次判断是否可被融合。当**识别到不可融合的算子时**,系统会生成第一段SuperKernel,并自动跳过该算子进行第二段SuperKernel融合。 | ||
| 23 | +- 目前支持SuperKernel融合的通信类算子包括AllReduce、ReduceScatter、AllGather、AlltoAll。 | ||
| 24 | +- 开启SuperKernel融合优化后,[算子Data Dump功能](../dfx/data_dump.md)将会失效。 | ||
| 25 | + | ||
| 26 | +## 使用方法 | ||
| 27 | + | ||
| 28 | +1. (可选)标定SuperKernel范围。 | ||
| 29 | + | ||
| 30 | + 使用如下语句块(super_kernel_scope_begin和super_kernel_scope_end),语句块内的算子将被融合为一个SuperKernel进行计算。 | ||
| 31 | + | ||
| 32 | + ```python | ||
| 33 | + torch.npu.super_kernel_scope_begin(scope_name: str) | ||
| 34 | + 待融合的算子操作 | ||
| 35 | + torch.npu.super_kernel_scope_end(scope_name: str) | ||
| 36 | + ``` | ||
| 37 | + | ||
| 38 | + scope_name:表示该范围内算子融合后的SuperKernel名称,相同的scope\_name代表相同的融合范围,由用户控制。若传入None,则该范围内的算子不进行SuperKernel融合。 | ||
| 39 | + | ||
| 40 | + >[!NOTE]说明 | ||
| 41 | + >- scope_name名称不超过256B字节。 | ||
| 42 | + >- 全局支持的scope_name(不同)数量不超过1024个。 | ||
| 43 | + >- 标定范围不代表SuperKernel的融合结果。 | ||
| 44 | + | ||
| 45 | + 当多个标定范围有交集时,按照标记集对算子进行划分。当交集不为子集时,如下图所示,op1和op2的标记集是\{sk1\},归入一组;op3和op4的标记集是\{sk1,sk2\},归入一组;op5、op6和op7的标记集是\{sk2\},归入一组。op8、op9没有被标记,不进行融合。 | ||
| 46 | + | ||
| 47 | +  | ||
| 48 | + | ||
| 49 | + 当交集为子集时,如下图所示,op1、op2、op6、op7、op8的标记集是\{sk1\},归入一组;op3、op4、op5的标记集是\{sk2\},归入一组。 | ||
| 50 | + | ||
| 51 | +  | ||
| 52 | + | ||
| 53 | + 当传入的标记为None,则被框定范围内的算子不进行融合,且该类型的标记具备最高优先级。如下图所示,标记为None的节点具备最高优先级,op3、op4、op5不进行融合。由于全局存在需要融合的标记sk1,所以不在标定范围内的op9也不进行融合。 | ||
| 54 | + | ||
| 55 | +  | ||
| 56 | + | ||
| 57 | + 当全局只存在为None的标记时,未被标记的算子都将进入SuperKernel的融合流程。如下图所示,op3、op4、op5标记为None,不进行融合,其他算子作为一组进入后续SuperKernel的融合流程。 | ||
| 58 | + | ||
| 59 | +  | ||
| 60 | + | ||
| 61 | +2. 通过npugraph\_ex的options配置开启SuperKernel融合优化,示例如下,该示例仅供参考,不支持直接拷贝运行,参数说明参见[表1](#table1)。 | ||
| 62 | + | ||
| 63 | + ```python | ||
| 64 | + import torch | ||
| 65 | + import torch_npu | ||
| 66 | + | ||
| 67 | + opt_model = torch.compile(model, backend="npugraph_ex", options={"super_kernel_optimize": True, "super_kernel_optimize_options": dict, "super_kernel_debug_options": dict}, dynamic=False, fullgraph=True) | ||
| 68 | + ``` | ||
| 69 | + | ||
| 70 | + **表 1** 参数说明<a id="table1"></a> | ||
| 71 | + | ||
| 72 | + |**参数名**|**参数说明**| | ||
| 73 | + |--|--| | ||
| 74 | + |super_kernel_optimize|布尔类型,是否开启SuperKernel融合优化。<br>•False(默认值):默认关闭。<br>•True:开启SuperKernel融合优化。| | ||
| 75 | + |super_kernel_optimize_options|字典类型,SuperKernel融合优化参数,参数说明参见[表2](#table2)。| | ||
| 76 | + |super_kernel_debug_options|字典类型,SuperKernel融合调试参数,参数说明参见[表3](#table3)。| | ||
| 77 | + | ||
| 78 | + **表 2** super_kernel_optimize_options参数说明<a id="table2"></a> | ||
| 79 | + | ||
| 80 | + |**参数名**|**参数说明**| | ||
| 81 | + |--|--| | ||
| 82 | + |dcci_before_kernel_start|通过本选项指定的算子,其内部调用GlobalTensor的GetValue/SetValue时不会自动插入缓存刷新指令,**而在SuperKernel调用该算子前**,会插入DataCacheCleanAndInvalid指令刷新ENTIRE_DATA_CACHE。<br>•配置格式形如:`dcci_before_kernel_start=[".*op1_type.*",".*op2_type.*"]`,仅支持匹配算子符号名而不是op_type,算子符号名为xx_op_type_xxx,带有前缀及后缀,仅支持`.*`格式的正则表达。<br>•若本选项指定的算子为<<<>>>算子或通信类算子,将无法去除算子内部调用GlobalTensor的GetValue/SetValue时插入的缓存刷新指令,仅支持控制SuperKernel调用算子前是否插入。SuperKernel默认会在所有子算子调用后插入DataCacheCleanAndInvalid指令刷新ENTIRE_DATA_CACHE。<br>•DataCacheCleanAndInvalid接口介绍参见《[CANN Ascend C API](https://hiascend.com/document/redirect/CannCommunityAscendCApi)》。<br>•本选项指定的算子op_type可通过Profiling查看,例如:`dcci_before_kernel_start=[".*GroupedMatmul.*",".*MoeGatingTopK.*"]`。| | ||
| 83 | + |dcci_after_kernel_end|通过本选项指定的算子,其内部调用GlobalTensor的GetValue/SetValue时不会自动插入缓存刷新指令,**而在SuperKernel调用该算子后**,会插入DataCacheCleanAndInvalid指令刷新ENTIRE_DATA_CACHE。<br>•配置格式形如:`dcci_after_kernel_end=[".*op1_type.*",".*op2_type.*"]`,仅支持匹配算子符号名而不是op_type,算子符号名为xx_op_type_xxx,带有前缀及后缀,仅支持`.*`格式的正则表达。<br>•若本选项指定的算子为<<<>>>算子或通信类算子,将无法去除算子内部调用GlobalTensor的GetValue/SetValue时插入的缓存刷新指令,仅支持控制SuperKernel调用算子后是否插入。SuperKernel默认会在所有子算子调用后插入DataCacheCleanAndInvalid指令刷新ENTIRE_DATA_CACHE。<br>•DataCacheCleanAndInvalid接口介绍参见《[CANN Ascend C API](https://hiascend.com/document/redirect/CannCommunityAscendCApi)》。<br>•本选项指定的算子op_type可通过Profiling查看,例如:`dcci_after_kernel_end=[".*GroupedMatmul.*",".*MoeGatingTopK.*"]`。| | ||
| 84 | + |dcci_disable_on_kernel|通过本选项指定的算子,其内部调用GlobalTensor的GetValue/SetValue会默认自动插入缓存刷新指令,**而在SuperKernel调用该算子前后**,不会插入任何DataCacheCleanAndInvalid指令。<br>•当某个算子同时配置了dcci_disable_on_kernel及dcci_before_kernel_start、dcci_after_kernel_end时,dcci_disable_on_kernel优先级低于其余两者,其余两者的功能将生效。为确保功能正确,暂不支持同时取消GlobalTensor的GetValue/SetValue中自动插入的缓存刷新指令以及SuperKernel调用算子前后插入的DataCacheCleanAndInvalid指令。<br>•配置格式形如:`dcci_disable_on_kernel=[".*op1_type.*",".*op2_type.*"]`,仅支持匹配算子符号名而不是op_type,算子符号名为xx_op_type_xxx,带有前缀及后缀,仅支持`.*`格式的正则表达。<br>•DataCacheCleanAndInvalid接口介绍参见《[CANN Ascend C API](https://hiascend.com/document/redirect/CannCommunityAscendCApi)》。<br>•本选项指定的算子op_type可通过Profiling查看,例如:`dcci_disable_on_kernel=[".*GroupedMatmul.*",".*MoeGatingTopK.*"]`。| | ||
| 85 | + |auto_op_parallel|多流融合场景下,SuperKernel默认按照算子下发顺序执行,相比单算子执行,可能无法达到最大化算子并行度。配置该选项后,在确保算子间依赖关系不变的前提下,可以将多条流上的算子按照Cube/Vector类型并行排布,获得相对较优的算子执行流水。<br>•0(默认值):不使能多流自动Cube/Vector自动并行排布。<br>•1:使能多流自动Cube/Vector自动并行排布。| | ||
| 86 | + |early_start|配置是否启用SuperKernel Early-Start优化。<br>•0:关闭本功能(默认值)。<br>•1:开启本功能。<br>开启后,SuperKernel会识别并生效算子内适配的AscendC::SetNextTaskStart()/AscendC::WaitPreTaskEnd()任务间同步接口(接口介绍参见《[CANN Ascend C API](https://hiascend.com/document/redirect/CannCommunityAscendCApi)》),达到算子头部Scalar指令提前执行的优化效果。| | ||
| 87 | + |aggressive_opt_strategies|字典类型,为保证功能正确性,SuperKernel默认采用保守融合策略,通过本选项可按需开启激进融合策略,实现更深度的优化。配置格式形如:`aggressive_opt_strategies={"task_breaker_bypass": 0, "value_breaker_bypass": 0b00}`,各子项说明如下:<br>1. task_breaker_bypass:<br>•0(默认值):SuperKernel融合过程中,遇到非AICore类型Task(如AICPU算子)会自动断开;<br>•1:SuperKernel融合范围内遇到了非AICore类型Task,且该流范围内没有其他AICore算子,则不做断开处理,仅将非AICore算子从融合范围内剔除。<br>2. value_breaker_bypass,比特位组合控制ValueWait节点的融合策略,各比特位含义如下:<br>•bit0(0b01):已配对的ValueWait节点(存在对应的ValueWrite节点且满足Flag匹配规则),在通过死锁检测后可进行融合。<br>•bit1(0b10):未配对的ValueWait节点(不存在对应的ValueWrite节点),允许进行融合。<br>取值示例:<br>•0b00(默认值):ValueWait节点不进行融合;<br>•0b01:仅允许已配对的ValueWait节点融合;<br>•0b10:仅允许未配对的ValueWait节点融合;<br>•0b11:同时允许已配对和未配对的ValueWait节点融合。<br>说明:ValueWait匹配规则详细可参考《[CANN Runtime API](https://hiascend.com/document/redirect/CannCommunityRuntimeApi)》。若ValueWait节点被多个ValueWrite匹配,则将会触发SuperKernel报错,整体网络融合失败。| | ||
| 88 | + | ||
| 89 | + **表 3** super_kernel_debug_options参数说明<a id="table3"></a> | ||
| 90 | + | ||
| 91 | + |**参数名**|**参数说明**| | ||
| 92 | + |--|--| | ||
| 93 | + |debug_sync_all|配置本选项后,SuperKernel内每个算子间会自动插入SyncAll全核同步指令,可辅助定界SuperKernel机制中的算子间同步问题。<br>•0:关闭本功能(默认值),SuperKernel子算子调用后不插入SyncAll全核同步。<br>•1:开启本功能。<br>注意:该模式会直接影响SuperKernel算子性能,建议仅在功能调试时使用。| | ||
| 94 | + |debug_op_exec_trace|配置本选项后,启用算子执行轨迹追踪功能。该功能会在Kernel侧记录每个子算子的执行状态(包括SK入口启动、子算子启动、子算子完成、SK入口完成),在异常发生时会在plog中打印所有核的算子执行进度信息(含当前执行到的算子ID及符号信息),辅助定位算子级执行异常。<br>•0:关闭本功能(默认值),不启用算子执行轨迹追踪功能。<br>•1:开启本功能。| | ||
| 95 | + |debug_cross_core_sync_check|配置本选项后,启用子算子Cube&Vector核间同步(参见《[CANN Ascend C API](https://hiascend.com/document/redirect/CannCommunityAscendCApi)》CrossCoreSetFlag接口mode 2介绍)校验功能,检测同步数量不匹配的问题。开启本功能时会隐式启用debug_op_exec_trace选项,用于辅助定位出错子算子。<br>•0:关闭本功能(默认值),不启用算子核间同步校验功能。<br>•1:开启本功能。| | ||
| 96 | + | ||
| 97 | +## 使用示例 | ||
| 98 | + | ||
| 99 | +```python | ||
| 100 | +import torch | ||
| 101 | +import numpy as np | ||
| 102 | +import torch_npu | ||
| 103 | + | ||
| 104 | +if __name__ == "__main__": | ||
| 105 | + | ||
| 106 | + # 定义模型model | ||
| 107 | + class ModelOrigin(torch.nn.Module): | ||
| 108 | + def __init__(self): | ||
| 109 | + super().__init__() | ||
| 110 | + def forward(self, x1, x2, scale, offset, bias, pertoken_scale, weight_scale): | ||
| 111 | + quant_matmul_res_origin = torch_npu.npu_quant_matmul(x1, x2, scale, offset=offset, bias=bias, pertoken_scale=pertoken_scale, output_dtype=torch.bfloat16) | ||
| 112 | + swiglu_res_origin = torch_npu.npu_dequant_swiglu_quant(quant_matmul_res_origin, weight_scale=weight_scale) | ||
| 113 | + return swiglu_res_origin, quant_matmul_res_origin | ||
| 114 | + class ModelSuperKernel(torch.nn.Module): | ||
| 115 | + def __init__(self): | ||
| 116 | + super().__init__() | ||
| 117 | + def forward(self, x1, x2, scale, offset, bias, pertoken_scale, weight_scale): | ||
| 118 | + # 将npu_quant_matmul和npu_dequant_swiglu_quant融合为SuperKernel,标记为sp1 | ||
| 119 | + torch.npu.super_kernel_scope_begin("sp1") | ||
| 120 | + quant_matmul_res_origin = torch_npu.npu_quant_matmul(x1, x2, scale, offset=offset, bias=bias, pertoken_scale=pertoken_scale, output_dtype=torch.bfloat16) | ||
| 121 | + swiglu_res_origin = torch_npu.npu_dequant_swiglu_quant(quant_matmul_res_origin, weight_scale=weight_scale) | ||
| 122 | + torch.npu.super_kernel_scope_end("sp1") | ||
| 123 | + return swiglu_res_origin, quant_matmul_res_origin | ||
| 124 | + | ||
| 125 | + m = 864 | ||
| 126 | + k = 7168 | ||
| 127 | + n = 4096 | ||
| 128 | + bias_flag = False | ||
| 129 | + cpu_x1 = torch.randint(-10, 10, (m, k), dtype=torch.int8) | ||
| 130 | + cpu_x2 = torch.randint(-10, 10, (n, k), dtype=torch.int8) | ||
| 131 | + cpu_x2 = torch_npu.npu_format_cast(cpu_x2.npu().transpose(1,0).contiguous(), 29) | ||
| 132 | + scale = torch.randn((n,), dtype=torch.float32) | ||
| 133 | + # print("scale:", scale) | ||
| 134 | + pertoken_scale = torch.randn((m,), dtype=torch.float32) | ||
| 135 | + # print("pertoken_scale:", pertoken_scale) | ||
| 136 | + bias = torch.randint(-1,1, (n,), dtype=torch.bfloat16) | ||
| 137 | + weight_scale = torch.randn((m, n), dtype=torch.float32).npu() | ||
| 138 | + # 使用图模式后端编译模型 | ||
| 139 | + model_sk = torch.compile(ModelSuperKernel(), backend="npugraph_ex", options={"static_kernel_compile": True, "super_kernel_optimize": True}, dynamic=False) | ||
| 140 | + print("-------------------- run sk -----------------------------------") | ||
| 141 | + swiglu_res_sk,quant_matmul_res_sk = model_sk(cpu_x1.npu(), cpu_x2, scale.npu(), None, None, pertoken_scale.npu(), weight_scale) | ||
| 142 | + model_no_sk = torch.compile(ModelOrigin(), backend="npugraph_ex", dynamic=False) | ||
| 143 | + print("-------------------- run no sk -----------------------------------") | ||
| 144 | + swiglu_res_origin,quant_matmul_res_origin = model_no_sk(cpu_x1.npu(), cpu_x2, scale.npu(), None, None, pertoken_scale.npu(), weight_scale) | ||
| 145 | + res = np.array_equal(swiglu_res_origin[0].cpu().numpy(), swiglu_res_sk[0].cpu().numpy()) | ||
| 146 | + res = res and np.array_equal(swiglu_res_origin[1].cpu().numpy(), swiglu_res_sk[1].cpu().numpy()) | ||
| 147 | + if res: | ||
| 148 | + print("Precision ====== Success!!!") | ||
| 149 | + else: | ||
| 150 | + print("Precision ====== Failed.") | ||
| 151 | +``` | ||
| @@ -83,6 +83,7 @@ torch.compile参数配置说明参见[表1](#fig1)。 | |||
| 83 | |[多流表达功能](./advanced/multi_stream.md)|大模型推理场景下,对于一些可并行的场景,可划分多个stream提升执行效率。| | 83 | |[多流表达功能](./advanced/multi_stream.md)|大模型推理场景下,对于一些可并行的场景,可划分多个stream提升执行效率。| |
| 84 | |[AI-Core和Vector-Core限核功能](./advanced/limit_cores.md)|提供Stream级核数配置,可调整最大AI Core数和Vector Core数,避免算子执行并行度降低。| | 84 | |[AI-Core和Vector-Core限核功能](./advanced/limit_cores.md)|提供Stream级核数配置,可调整最大AI Core数和Vector Core数,避免算子执行并行度降低。| |
| 85 | |[自定义FX图优化Pass功能](./advanced/post_grad_custom_pass.md)|传入自定义FX Pass函数,该配置可控制自定义Pass在框架内置Pass执行前/后生效。| | 85 | |[自定义FX图优化Pass功能](./advanced/post_grad_custom_pass.md)|传入自定义FX Pass函数,该配置可控制自定义Pass在框架内置Pass执行前/后生效。| |
| 86 | +|[SuperKernel功能](./advanced/superkernel.md)|将连续的可融合Task合并为一个SuperKernel Task,减少任务调度开销,提升执行性能。| | ||
| 86 | 87 | ||
| 87 | **表 4** npugraph\_ex DFX功能 <a id="fig4"></a> | 88 | **表 4** npugraph\_ex DFX功能 <a id="fig4"></a> |
| 88 | 89 | ||