NPU 图模式优化及其在 cann-recipes-infer 框架下的使能【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer1. 文档目标本文主要回答四个问题什么是图模式它解决什么问题NPU 图模式有哪些主要实现方式它们之间是什么关系在本仓的执行框架下如何为模型使能图模式适配和调试过程中最常见的问题是什么应该如何排查2. 图模式基础知识2.1 什么是图模式pytorch 默认按照eager模式可以理解为“边解释边执行”的模式。模型前向走到哪里框架就把当前位置对应的算子逐个下发到 Device 执行开发和调试都比较直接但在推理场景下也更容易暴露 Host 侧逐算子下发的开销。在图模式下前向逻辑会先被torch.compile捕获为图再交给 NPU 后端编译和执行。这样做的核心收益是减少 Host 逐算子下发开销缓解 host bound。让后端在更大范围上做算子融合、内存复用、调度优化。这通常可以显著降低时延。2.2 为什么图模式通常只用于 decode对 LLM 推理来说prefill和decode的特征差异很大阶段输入特征是否适合图模式原因prefill序列长度动态变化序列通常较长通常不建议shape 和控制流容易变化容易断图或重编译同时 prefill 阶段单次算子执行时间更长下发 bound 往往没有 decode 明显decode单 token 或固定小长度输入推荐shape 更稳定更容易形成可复用图2.3 什么是编译缓存图模式除了“如何编译”还有一个经常一起出现的概念编译缓存也就是cache compile。它解决的问题是重复启动或重复执行同一图时的启动开销。如果模型结构、输入 shape、dtype 和图模式配置等保持稳定把第一次编译的结果缓存下来后续运行可以直接复用缓存减少启动开销。 需要注意的是编译缓存是否命中强依赖模型代码、缓存目录和图模式配置是否一致。3. NPU 图模式的支持方式3.1 实现方式NPU 图模式实现方式可以归纳为两大类GE 图模式将 FX 图转换成 Ascend IR再由 GE 引擎编译执行npugraph_ex 图模式基于 npugraph capture replay强调下沉调度和低开销执行3.2 为什么文档里还会看到acl_graph仓内存在一部分较早的模型代码和文档仍然使用acl_graph这个历史命名实际阅读时通常可以把它对应到npugraph_ex这一路图模式。但注意新接入模型时优先使用执行框架的ge_graph/npugraph_ex配置。后续会逐步废弃acl_graph并清理对应实现和文档。3.3 两种方式如何选择目前建议优先功能稳定、性能更优选择ge_graph。当然npugraph_ex也在持续优化中本文档会持续更新。优先适配更轻量使用体验更接近 eager 模式选择npugraph_ex是本仓后续主推的模式。4. cann-recipes-infer 如何使能图模式4.1 使能链路配置层executor/core/config/inference_config.py通过model_config.exe_mode选择eager/ge_graph/npugraph_ex预热层executor/core/engine/execution_engine.pywarm_up()会先跑一次prefill再跑一次decode如果启用了图模式decode预热阶段会触发编译编译层executor/core/model_worker/model_worker.pycompile_model()最终调用executor/utils/graph_utils.py中的compile_model_forward()执行层ModelWorker.inference()中只有not is_prefill且exe_mode in [ge_graph, npugraph_ex]时才走self.model_compiled(**model_inputs)元数据层executor/utils/forward_metadata.pyForwardMetaData负责把is_prefill、kv_len、actual_seq_lengths_*、attention_mask等动态信息传给模型需要注意的是图模式开关只对 decode 阶段生效prefill 阶段统一走 eager 模式。4.2graph_utils.compile_model_forward()核心逻辑解析OfflineInference最终真正执行图编译的核心函数在../../executor/utils/graph_utils.py里的compile_model_forward()。npugraph_ex和ge_graph在这个函数里的主流程大体一致先做通用准备再根据exe_mode组织图编译配置最后根据enable_cache_compile选择普通编译还是cache_compile。差异主要在于配置承载方式以及dynamic的默认取值不同。通用准备逻辑import torchair as tng import torchair.ge_concrete_graph.ge_converter.experimental.patch_for_hcom_allreduce tng.patch_for_hcom() torch._dynamo.config.inline_inbuilt_nn_modules Falsetng.patch_for_hcom()是用来处理集合通信入图的在 PyTorch 2.6 及之后版本中这一步通常可以省略。inline_inbuilt_nn_modules False用于避免内建模块被过度内联减少部分图编译场景下的不确定性。两种图模式的主要差异真正需要开发者关注的差异主要有下面几项项目npugraph_ex模式ge_graph模式配置承载方式通过optionskwargs 传入通过CompilerConfig.experimental_config成员配置普通编译后端backendnpugraph_exbackendtng.get_npu_backend(...)缓存编译接口torch.npu.npugraph_ex.inference.cache_compile(...)tng.inference.cache_compile(...)dynamic设置TrueFalse配置项static_kernel_compile/frozen_parameterfrozen_parameter/tiling_schedule_optimize/topology_sorting_strategy补充说明当前npugraph_ex模式之所以保持dynamicTrue核心原因并不是图本身必须动态而是当前配套使用的 FIA 算子接口里actual_seq_lengths等入参还不支持 Tensor 输入只支持list[int]输入。在这种前提下如果强行使用dynamicFalse容易触发重编译。后续算子接口补齐 Tensor 输入支持这里的配置也会随之调整。4.3 图模式适配里最常见的两个问题Graph Break图捕获过程中断部分逻辑回退到 Python/Eager 执行。一般需要通过减少 Python 控制流、避免.item()、补齐入图适配来解决。Recompile虽然能入图但由于 shape、地址或 guard 条件变化导致反复重新编译性能变差。一般需要通过固定 shape、固定缓存地址并把动态量改为显式输入来解决。4.4 模型适配时必须满足的条件仅仅打开exe_mode不够。模型本身需要满足图模式约束。条件一模型必须先能在 eager 下稳定运行图模式不会修复 eager 下本来就存在的错误。条件二prefill 和 decode 要明确区分推荐做法prefill保持 eagerdecode使用图模式用forward_metadata.is_prefill或独立的prefill()/decode()方法区分两条路径条件三将动态信息以显式输入形式传入模型典型动态信息包括kv_lenposition_idsactual_seq_lengths_qactual_seq_lengths_kvis_prefill这些信息应该由框架构造后传给模型而不是在模型内部临时生成 Python 标量或依赖隐式状态推导。条件四KV Cache 和常驻 buffer 需要预分配并原地更新错误示例key torch.cat([past_key, new_key], dim1)推荐示例torch_npu.scatter_update_(k_cache, kv_len, key_states, -2) torch_npu.scatter_update_(v_cache, kv_len, value_states, -2)目标是避免 decode 场景下 KV Cache 的 shape 或地址发生变化以减少 dynamo guard 失败和重编译条件五避免典型的 Graph Break 写法尤其要避免tensor.item()基于 Tensor 值的 Pythonif/while在 forward 内部临时创建影响 shape 的控制分支根据 Python list/tuple 长度变化来切换图内控制流4.5 FIA 融合算子适配建议不同模式下FIA算子的接口入参有所区别推荐按照如下方式使用模式常见 FIA 接口actual_seq_lengths建议说明ge_graphtorchair.opsTensor只适合 GE 图模式npugraph_extorch_npu常见为list[int]当前以推理场景的 FIA 接口为主后续算子支持变化时再调整可以参考qwen3-moe的样例代码models/qwen3_moe/models/modeling_qwen3_moe.py在 decode 且enable_gegraph时会切到torchair.opsExecutionEngine在npugraph_exdecode 模式下会把actual_seq_lengths_*转成list5. 模型接入图模式的推荐步骤建议按下面的顺序推进而不是一次性把所有优化叠加上去。先跑通 eager 模式并确认精度正确。消除 graph break 和重编译风险去掉.item()和动态 Python 控制流固定 KV cache 与常驻 buffer 的地址和 shape并明确actual_seq_lengths的类型与来源。做功能、精度验证对比 eager 和 graph 输出至少覆盖一轮 prefill 多轮 decode如果模型有 MTP还要确认 main model 和 MTP model 的图模式输入 shape、dtype 和长度组织方式都稳定。最后再按需开启增强特性例如enable_cache_compile、enable_static_kernel仅npugraph_ex、model 自带的enable_superkernel目前仅ge_graph以及多流、限核等能力。6. 常见 Troubleshooting6.1 高频问题速查表现象常见根因处理建议编译前就报错eager 路径本身不正确或者模型输入的 shape、dtype、长度组织方式不稳定先单独验证 eager再检查图模式输入和前向参数组织图捕获中断Graph Break.item()、Tensor 驱动的 Python 分支、print、自定义算子未适配改写为 Tensor 逻辑或补齐入图适配decode 性能没有提升甚至变差发生重编译kv_len、actual_seq_lengths_*、缓存地址或输入 shape 不稳定打开torch._logging.set_logs(recompilesTrue)检查重编译原因重点关注 guard 变化来源actual_seq_lengths类型报错图模式与 FA 接口不匹配ge_graph优先 Tensor torchair.opsnpugraph_ex对齐本仓当前list[int]方案enable_static_kernel报错模式不对该选项只允许在npugraph_ex模式使用superkernel相关报错在不支持的模式上开启目前只在ge_graph模式尝试通信无法入图当前依赖版本要求的集合通信入图前置处理没有完成结合当前 PyTorch / TorchAir 版本检查是否需要torchair.patch_for_hcom()或其他等效前置配置cache compile 不生效缓存目录、输入 shape / dtype / 长度组织方式或配置变化固定 cache 目录避免模型代码或编译参数频繁变化图模式下精度异常KV cache / FA / 原地更新语义变化先回退到 eager 对齐再逐项恢复优化6.2 常用调试手段实践中常用的调试手段包括torch._logging.set_logs(recompilesTrue)出现重编译时打开相关日志eager / graph 输出比对看功能、精度是否一致cache compile 开关测试区分“图编译问题”还是“缓存命中问题”7. 参考资料仓内资料Offline Inference 执行机制设计文档InferenceConfig 类使用指南MTP 模型接入指南官方资料TorchAir 文档总览https://gitcode.com/Ascend/torchair/tree/master/docs/zhGE / Ascend IR 图模式https://gitcode.com/Ascend/torchair/tree/master/docs/zh/ascend_ir/ascend_ir.mdGE 图模式快速上手https://gitcode.com/Ascend/torchair/tree/master/docs/zh/ascend_ir/quick_start.mdnpugraph_ex 后端https://gitcode.com/Ascend/torchair/tree/master/docs/zh/npugraph_ex/npugraph_ex.mdnpugraph_ex 快速上手https://gitcode.com/Ascend/torchair/tree/master/docs/zh/npugraph_ex/quick_start.md常见案例与定位方法https://gitcode.com/Ascend/torchair/tree/master/docs/zh/appendix/casesFAQhttps://gitcode.com/Ascend/torchair/tree/master/docs/zh/appendix/faq.md【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考