CANN npugraph_ex图模式优化【免费下载链接】cann-learning-hubCANN 学习中心仓支持在线互动运行、边学边练提供教程、示例与优化方案一站式助力昇腾开发者快速上手。项目地址: https://gitcode.com/cann/cann-learning-hub前言随着人工智能的飞速发展大模型推理场景的“低时延高吞吐”诉求推动了PyTorch图模式的快速发展。torch.compile是PyTorch 2.0推出的核心特性通过即时编译JIT将PyTorch代码转换为计算图支持inductor等多种后端编译器并利用图捕获和重放能力减少Host下发头开销以优化大模型推理时延。为此CANN提供了具备图捕获和重放能力的aclGraph技术并基于torch.compile提供高性能的图后端npugraph_ex融合了aclGraph的调度能力和亲和NPU的图优化能力以满足大模型推理场景下的性能加速需求。HighLights快速接入CANN提供aclGraph有效的消除了算子调度时延并已经集成在PyTorch中可快速使能高性能在复用原有aclGraph调度能力前提下借助于PyTorch的FX图进一步叠加NPU亲和的优化可集成面向大模型推理的服务化框架可以做到快速集成无缝对接到vLLM/SGLang等服务化框架中。1. 快速接入基于aclGraph的调度加速能力aclGraph是CANN推出的一种基于图下发的优化机制核心目标是降低任务提交的开销。aclGraph提供了图捕获及重放能力的C接口为方便用户使用在Pytorch层也封装了对等的接口torch.npu.graph开发者可以使用相关的接口平滑的迁移到NPU上极大简化了使用成本将现有PyTorch中的使用代码仅改造接口调用部分的代码即可享受NPU加速带来的性能提升。这种开箱即用的设计大大减少了开发难度和时间成本。使用示例# capture g torch.npu.NPUGraph() with torch.npu.graph(g, streams): static_output model(static_input) # ...data preparation... # replay g.replay()aclGraph的接口文档https://www.hiascend.com/document/detail/zh/CANNCommunityEdition/850alpha001/appdevg/acldevg/aclcppdevg_000519.htmlaclGraph所取得的收益来自于调度层收益基于此前CANN图模式在NPU优化上的经验积累面向PyTorch框架下的大模型训推场景提供了轻量化、高性能、可集成的PyTorch FX图优化器***npugraph_ex***。2.高性能NPU图优化样板间--npugraph_ex在不增加开发者使用成本的前提下基于aclGraph为了更充分的发挥NPU的图模式性能以torch.compile为入口我们提供了基于FX图的NPU优化组件npugraph_ex我们将其定位为一个基于CANN的图模式样板间并开源了相关代码https://gitcode.com/Ascend/torchair 以方便开发者复用或者参考构建自己的图优化器整体逻辑架构如下图所示在保持开发者行为一致的前提下npugraph_ex通过多项创新优化充分利用CANN的能力挖掘了NPU潜能构建了其在图计算方面独特的优势。整体优化内容分为三个层次Common类主要是能力增强相关优化内存复用FA刷新输入copy优化等融合pass类基于图算子的一些融合优化部分复用inductor的基础pass能力NPU高阶优化面向性能的NPU特有优化静态化Kernel、多流等相关优化特性介绍如下2.1Common类Host tilling参数刷新aclGraph支持在Replay前刷新内核kernels的功能能够在类FIA算子执行前根据实际计算出的actual_seq_len值动态更新分块tiling结果从而确保此类算子能够被aclGraph成功捕获并消除tiling计算所带来的调度间隙。 如下图所示主要通过以下两个阶段实现Capture阶段通过update接口在capture阶段记录需要更新参数的算子并插入event wait控制更新时的时序Replay阶段replay接口调用后下发需要更新参数的算子并插入event record保证参数更新后算子执行。2.2Common类内存复用nupgraph_ex采用静态下沉执行模式在图捕获capture阶段固化图相关的内存地址因此当存在多个 aclGraph实例时将导致多份独立的内存分配造成资源浪费。为解决多图场景下的内存重复占用问题我们提供了内存复用能力帮助用户在多个 aclGraph之间高效共享内存资源显著降低整体内存开销。 npugraph_ex提供了2种内存复用机制单张FX图之间共享内存池的内存复用默认开启一般是在dynamic为True时产生效果每次出现新的shape就会触发一次新的aclGraph的捕获同一张动态的FX图捕获的多张aclGraph之间可进行内存复用由于动态shape变化产生多张FX图——capture多张aclGraph场景的内存池复用默认关闭需要用户在确认无内存踩踏的情况下开启。通过为不同的FX图之间指定相同的内存池达成内存复用一般在dynamicFalse时使用。2.3FX pass 优化npugraph_ex内置了多种常用的优化Pass如pattern pass、冗余算子消除、reinplace优化和自定义pass等。通过这些优化能有效减少计算图中的冗余计算降低内存占用提高计算吞吐量。pattern pass开发者可以自定义算子替换规则使用融合算子替换FX图中多个算子有效减少部分场景下不必要的下发开销提高模型执行效率通过register_replacement接口实现默认开启。冗余算子消除自动识别并消除计算图中不影响程序逻辑或数据计算的冗余操作减少非必要的计算开销通过config.experimental_config.remove_noop_ops True配置方式实现默认开启。reinplace优化通过优化FX图安全地把out-placecopy_算子替换回脚本的in-place算子提升性能对于中间节点的reinplace直接调用原生FX实现对于原地修改输入的reinplace参考inductor实现逻辑增加输入原地修改reinplace的实现。自定义pass开发者能够使用PyTorch FX的图变换能力实现算子替换、子图重构、多流并行等通用优化。通过开放特有的API支持开发者在FX Graph中直接表达原生框架所不具备的多流并行与流间执行时序控制等硬件级优化这使多流并行的实现得以自动化、结构化显著降低了适配成本与代码冗余。2.4NPU高阶优化静态化编译单个算子在使能静态Kernel编译后其在 Device上的执行时间显著缩短。这一性能提升主要源于基于静态shape编译之后算子Kernel内部的计算逻辑基于静态数值得到了化简从而消除运行时判断与动态计算开销实现更高效执行。静态编译通过编译期确定tensor shape、适配硬件优化、标量预计算标量数据和消除冗余计算实现内存高效利用、针对性能提升及并行指令流优化。2.5NPU高阶优化多流并行大模型推理场景下为了提高硬件利用率从而提升推理性能对于一些可以并行的场景在图模式下做分流并行操作可以提升端到端性能从而提升图模式竞争力。在aclGraph中通过Event机制可以实现多流并行捕获与执行。 具体而言在单图多流场景下利用aclGraph的Event推导能力可以在不同Stream之间建立同步关系当在Stream1中下发Record Event并在Stream2中下发对应的Wait Event时能够将Stream2上的任务一并捕获到模型运行时实例中从而实现多流并发执行。为了确保多流捕获的正确性加入捕获状态的Stream最终需要通过直接或间接的Record Event回到主流否则在结束捕获阶段可能会触发错误。这种机制使得多个异步执行流可以在满足依赖关系的前提下协同工作提高硬件资源利用率。 此外Event机制不仅用于控制执行顺序还用于内存管理中的多流复用确保在Device侧计算完成后再进行内存复用避免数据竞争和错误。3. 可集成性拥抱主流社区生态热门推理框架如vLLM/SGLang等提供了其默认的图模式后端以加速大模型推理性能npugraph_ex同样可以无缝集成到用户现有的工作流中以加速服务化框架在NPU上的执行性能当前相关代码已经合入至对应社区。vLLM Ascend社区接入示例https://github.com/vllm-project/vllm-ascend/pull/4700SGLang 社区接入示例https://github.com/sgl-project/sglang/pull/134104.性能表现 使用体验与其他图编译器后端相比npugraph_ex的性能优势非常明显,能够提供更低的延迟和更高的吞吐量特别是在处理复杂图计算任务时表现尤为突出。以下是decode阶段执行耗时集成到vLLM框架后的性能提升效果模型eagerfullgraphfullgraphnpugraph_ex(force_eager)Qwen3-30B150.33ms21.9ms20.2msDeepSeekV3.1185.6ms45.0ms41.2ms实测数据表明在Qwen3-30B模型上npugraph_ex相较fullgraph实现7.8%的性能提升在DeepSeekV3.1模型npugraph_ex进一步达到8.5%的优化幅度借助于其性能与易用性的均衡npugraph_ex已经被来自多个头部模型厂商的开发者所借鉴并以其生态一致和易集成的特点被开发者所青睐开发者自主完成自建图模式的快速迭代。5.未来路标在保持易用性与可集成性的前提下持续进行性能的优化图内设置AI Core和Vector Core核数多流场景下会出现所有核Core都被一个流占用的情况导致算子执行并行度降低因此需要把核分给不同的流用从而保证算子并行执行收益Super Kernel提供可集成的Super-Kernel特性插件化接入npugraph_ex中减少核启动开销降低执行时延Graph-Autofusion集成graph-autofusion实现更加智能的自动融合结合硬件特性自动优化计算任务进一步提高图计算的性能。【免费下载链接】cann-learning-hubCANN 学习中心仓支持在线互动运行、边学边练提供教程、示例与优化方案一站式助力昇腾开发者快速上手。项目地址: https://gitcode.com/cann/cann-learning-hub创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考