尧图网站设计 尧图网站设计YAOTU DESIGN
ARTICLE DETAIL

资讯详情

深耕网站设计与一线实操的经验洞察。

CANN 平台 SuperKernel 算子二进制融合:模型图调度优化原理、开启方式与 LongCat-Flash 实战

CANN 平台 SuperKernel 算子二进制融合:模型图调度优化原理、开启方式与 LongCat-Flash 实战 CANN 平台 SuperKernel 算子二进制融合模型图调度优化原理、开启方式与 LongCat-Flash 实战【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer本篇围绕 cann-recipes-infer 仓库中的 SuperKernel 技术文档展开系统讲解这一面向模型计算图的调度优化技术它的二进制融合原理、四类算子级优化与两类网络级优化机制以及如何在npugraph_ex和 GE 图模式两种后端中实际开启。读完本文你将掌握 SuperKernel 融合范围的标定方法并能结合仓库中 LongCat-Flash 优化样例理解“按分核与分流范围标定多个 SuperKernel”的工程实践。1. 背景调度层面还剩多少性能空间提升 LLM 推理性能通常会叠加多种手段算子级优化算子融合、通信-计算重叠、Weight 预取等网络执行优化Continuous Batching、Pipeline Parallelism 等。但这些手段主要作用于算子内部或任务编排模型调度层面仍有明显空间在图模式下即使已经做了算子融合也受融合规则限制融合覆盖范围有限融合后的算子之间仍需逐个调度执行存在调度开销与等待时间。SuperKernel 正是针对这一层的优化——通过更彻底的任务调度进一步压缩计算过程中的等待时间。SuperKernel 的核心思想是基于计算图中算子的先验信息算子类型、前后序依赖关系等结合即时编译JIT能力突破传统算子融合规则的限制把网络中的多个算子编译为一个 SuperKernel 整体调度执行从而显著降低算子间调度开销。2. 原理算子二进制融合SuperKernel 是一种算子二进制融合技术。它与源码融合不同聚焦于内核函数Kernel的二进制调度优化基于已编译的二进制代码创建一个超级 Kernel 函数SuperKernel把多个其他内核函数作为子函数来调用。相比单算子逐个下发SuperKernel 带来三方面收益减少调度开销降低任务调度的等待时间和调度开销利用 Task 间隙资源进一步摊薄算子启动开销编译期全局视野编译阶段可获取全部子算子的先验信息从而实施更深层次的算子级与网络级优化。下图展示了融合前后的算子序列变化以 MoE 计算子图为例从源码结构看这种“按依赖切分融合段”的行为在实现层有明确体现编译器按算子执行顺序依次识别可融合性遇到不可融合算子如 TBE 算子时就切段详见约束章节。2.1 算子级优化SuperKernel 编译期拿到的是二进制 Kernel及其依赖关系因此优化必须建立在二进制语义之上。文档列出了四类典型优化2.1.1 ICache Preload 优化SuperKernel 执行时运行系统通常只预取入口点的指令子 Kernel 的代码段往往无法被硬件预取机制捕获导致指令缓存ICache命中率下降产生 ICache Miss。ICache Preload的解法是在当前子 Kernel 执行完毕前以2KB 对齐方式提前将下一个子 Kernel 的代码段预取进 ICache。指令加载延迟被隐藏在当前子 Kernel 的执行过程中后续算子的 ICache Miss 随之减少。2.1.2 Early-Start 优化常规调度要求前序算子全部完成后才启动后续算子。但观察指令构成可以发现并发机会多数前序算子的末尾指令是 MTEMemory Transfer Engine内存传输引擎数据搬运指令后续算子的起始指令通常是与输入数据无关的初始化标量指令。两类指令分属不同计算单元具备并发条件。Early-Start 的做法是在前序算子的搬运指令前插入Set 同步点在后续算子的初始化指令后插入Wait 同步点让两个子算子的部分指令重叠执行。2.1.3 同步优化为保证执行顺序正确SuperKernel 在各子算子调度之间会插入全核同步。对于 Kernel Type 为Mix 1:21 个 Cube 核对应 2 个 Vector 核的混合类型算子完整的全核同步需要等待所有 AI Core 的 Vector 核与 Cube 核都到达同步点——开销并不小。由于编译时能够识别每个子算子的 Kernel TypeSuperKernel 可以定制同步范围例如连续两个 Vector 算子之间只需执行全 Vector 核同步即可。通过细粒度控制同步范围可显著降低子算子间的同步开销。2.1.4 子 Kernel 复制多核系统下多个计算核心执行同一段代码时会并发访问同一指令地址在共享 L2 Cache 层面形成串行化访问队列引发资源争用削弱多核并行的性能增益。SuperKernel 的解法是将子 Kernel 代码复制多份不同核心按核 ID 映射到不同物理地址执行从而缓解同一指令地址的争用提升算子执行效率。2.2 网络级优化2.2.1 Tiling 下沉与 Weight 预取SuperKernel 支持基于内存语义的Notify与Wait事件以此适配两类典型网络级场景Tiling 下沉Tiling 下沉算子指 Tiling 计算依赖前序算子的输出结果为避免主机与设备间频繁交互这类 Tiling 计算被部署在 AICPU 上执行。SuperKernel 融合它时需要区分两种情况若融合的是其前序算子前序算子执行完成后通过Notify事件通知 AICPU 启动 Tiling 计算若融合的就是 Tiling 下沉算子本身需先通过Wait事件等待 AICPU 完成 Tiling 计算再执行 Device 侧计算。Weight 预取借助 CMOCache Management Operation任务调用专用硬件单元SDMA把权重数据提前加载到 L2 Cache。SDMA 与 AI Core 的协作正是通过内存语义的Notify/Wait事件实现的。这一点在仓库代码中有直接呼应LongCat-Flash 样例中与 SuperKernel 协同使用的npu_prefetch权重预取逻辑同样依赖 Event 与独立 Prefetch 流完成 SDMA 搬运与计算的重叠见 模型实现。2.2.2 双流并发融合在通过多流实现 Cube/Vector 并发之后如果 SuperKernel 转换时不感知依赖关系、仅按执行顺序处理SuperKernel 内部执行会退化为串行性能收益不及预期。解决方式将算子按类型分类在 SuperKernel 内部按 Cube 和 Vector 的不同特性分配执行队列结合流的属性和 Event精准插入同步点实现 Cube 与 Vector 的高效并发执行。该小节整理自 graph-autofusion 项目的能力描述具体实现与能力以对应版本源码和文档为准。3. 实现方式两种后端如何开启当前 SuperKernel 特性通过PyTorch 图模式开启主要支持npugraph_ex后端和GE图模式两种方式。3.1 npugraph_ex 后端在npugraph_ex后端中SuperKernel 融合能力通过torch.compile的options参数配置设置super_kernel_optimizeTrue即可开启compiled_model torch.compile( model, backendnpugraph_ex, # 省略部分参数 options{ static_kernel_compile: True, super_kernel_optimize: True, # 省略部分参数 }, )其中super_kernel_optimizeTrue表示开启 SuperKernel 融合优化。如需进一步控制融合范围可通过范围标定接口手动圈定torch.npu.super_kernel_scope_begin(scope_name: str) torch.npu.super_kernel_scope_end(scope_name: str)在super_kernel_scope_begin/end标定的范围内满足条件的算子将参与 SuperKernel 融合。仓库内的 编译工具函数 展示了这套配置在真实推理框架中的组装方式从推理配置的custom_params中读取开关并一并传入编译选项——enable_superkernel model_config.custom_params.get(enable_superkernel, False) if exe_mode npugraph_ex: compile_options { frozen_parameter: True, static_kernel_compile: enable_static_kernel, super_kernel_optimize: enable_superkernel, super_kernel_optimize_options: {dcci_disable_on_kernel: [.*]} } ... compiled torch.compile(model_forward, dynamicenable_dynamic_graph, fullgraphTrue, backendnpugraph_ex, optionscompile_options)从源码结构看这里除了文档提到的super_kernel_optimize外还额外设置了super_kernel_optimize_options中的dcci_disable_on_kernel选项用于按 Kernel 正则控制 DCCI 行为属于仓库对编译选项的进一步工程化配置。3.2 GE 图模式后端在 GE 图模式中SuperKernel 通过 TorchAir 提供的作用域接口进行标定并配合 TorchAir 的CompilerConfig开启。用户需要先分析模型脚本中可被融合的算子范围再使用torchair.scope.super_kernel标定融合区域with语句块内的算子会被融合为一个 SuperKernel 计算。接口形式with torchair.scope.super_kernel(scope: str, options: str ): ...参数说明scope当前上下文中算子融合后的 SuperKernel 名称。相同的scope表示属于同一个融合范围由用户指定optionsSuperKernel 的编译选项。示例import torchair with torchair.scope.super_kernel(super_kernel_0): y op1(x) z op2(y)上述示例中op1和op2位于同一个super_kernel_0作用域内满足融合条件时会被融合为一个 SuperKernel。当scope为None时该作用域内的算子不进行 SuperKernel 融合。仓库为这一用法封装了带开关的上下文管理器见 superkernel_scopedef superkernel_scope(enable: bool, scope: str, options: str None): if enable: return tng.scope.super_kernel(scope, options) else: return FakeContextManager()这样同一份模型代码可以在“开启/关闭 SuperKernel”两种配置下运行无需修改模型结构便于做 A/B 对比。4. 实战样例LongCat-Flash 中的 SuperKernel 标定以仓库中的 LongCat-Flash 模型优化样例 为例。该样例在不同流上采用了不同的分核策略按照分核与分流的范围在整网中共标定了三个 SuperKernel对应源码中的标定位置清晰可查1每层解码器标定的三段作用域。在 multi_stream_forward 中每层被划分为三段 SuperKernel 范围# scope_{layer}_part1输入 LayerNorm 第一组注意力含 O 投影 with superkernel_scope(self.enable_superkernel, fscope_{self.layer_idx}_part1, ): hidden_states, residual self.input_layernorm0 attn_ret self.self_attn[0].forward_page_attention_absorb(...) ... # scope_{layer}_part2_moe独立 MoE 流上的 FFN配合 limit_core_num 限制核数 with limit_core_num(True, self.aic_num1, self.aiv_num1): with superkernel_scope(self.enable_superkernel, fscope_{self.layer_idx}_part2_moe, ): shortcut_mlp_output self.mlp(hidden_states_norm, is_prefill, cur_topk_listcur_topk_list) # scope_{layer}_part2_main主计算流上的第二组 MLP 与注意力 with limit_core_num(not self.enable_afd, self.aic_num2, self.aiv_num2): with superkernel_scope(self.enable_superkernel, fscope_{self.layer_idx}_part2_main, ): ...2AFDAttention/FFN 解耦部署下的独立 FFN 子图。在 FFNModel.forward 中decode 阶段每层 MoE 计算被整层包进一个作用域for i in range(self.moe_layer_num): ... else: with superkernel_scope(self.enable_superkernel, fscope_{i}_moe, ): hidden_states, gmm2_out, gmm2_event self.layersi # 层间在独立预取流上执行 npu_prefetch预取下层 router 权重与 w13 权重 if i self.moe_layer_num - 1: with npu_stream_switch(on_stream, self.npugraph_prefetch_stream): ... npu_prefetch(self.enable_prefetch, self.layers[i 1].mlp.router.classifier.weight.data, ...) npu_prefetch(self.enable_prefetch, self.layers[i 1].mlp.experts.w13_weight.data, ...)3配置入口。开关通过推理配置文件的custom_params控制例如 longcat_flash_densetp8_ep128_gegraph_mtp_eplb_w8a8.yamlmodel_config: exe_mode: ge_graph # [eager, npugraph_ex, ge_graph] ... custom_params: enable_multi_streams: 1 # [0, 1, 2] enable_prefetch: True # [False, True] enable_superkernel: True # [False, True]仓库中对这一能力的适用性还有明确的源码级防护与版本限制说明实践时需注意从 模型入口 看eager模式不支持 cache compile 与 superkernelLongCat-Flash 的npugraph_ex分支当前也不支持 superkernel开启会直接抛出ValueError因此该样例中 SuperKernel 实际运行在ge_graph模式下LongCat-Flash README 说明当前 CANN 软件版本下SuperKernel 标记范围内的部分算子尚不支持完全融合该限制将在后续社区版本中解决。因此实际收益取决于所用 CANN 版本对范围内算子的融合覆盖程度。5. 约束与注意事项开启 SuperKernel 前需确认以下约束源自文档约束章节切段规则编译器按网络中算子的执行顺序依次识别可融合性。若遇到 TBETensor Boost Engine等不可融合算子会将其前面已识别的连续可融合算子组成一段 SuperKernel同时跳过该不可融合算子继续向后识别并组成下一段 SuperKernel通信算子支持范围目前支持 SuperKernel 融合的通信类算子包括 AllReduce、ReduceScatter、AllGather 和 AlltoAll静态编译前置依赖在npugraph_ex图模式下开启 SuperKernel 融合优化时需要同时开启静态 Kernel 编译功能对应 3.1 节示例中的static_kernel_compile: TrueGE 图模式要求静态图且with语句块内不支持断图功能取舍开启 SuperKernel 融合优化后将禁用算子 Data Dump功能调试手段需相应调整。6. 小结SuperKernel 把“逐算子调度”升级为“按融合段整体调度”通过 ICache Preload、Early-Start、细粒度同步与子 Kernel 复制等算子级手段叠加 Tiling 下沉、Weight 预取、双流并发等网络级手段系统性压缩图模式下的调度开销。仓库提供了完整的落地参照executor中 superkernel_scope 封装 与 npugraph_ex 编译选项组装 展示了两种后端的开启路径LongCat-Flash 样例 则演示了在真实 MoE 推理模型中按分核/分流范围标定多个 SuperKernel、并与权重预取流协同工作的完整工程实践。【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
返回列表