CANN 推理加速实战npugraph_ex 图模式后端使用与 LLM Decode 适配指南【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infernpugraph_ex 是基于torch.compile的 aclgraph 图模式加速方案通过捕获模式Capture Replay将模型前向固化为可在昇腾 NPU 上低开销回放的计算图。本指南以其完整使用手册为主体结合 cann-recipes-infer 开源仓库中的框架接入实现graph_utils.py、推理配置inference_config.py与真实模型配置样例讲解 npugraph_ex 的适用场景、使用约束、快速上手、options 配置速查、问题定界流程与常见问题排查帮助读者在 LLM 在线推理尤其是 Decode 阶段中快速完成图模式适配。一、npugraph_ex 是什么npugraph_ex 是昇腾 NPU 上一种基于torch.compile的图模式后端。与 GE 图模式将 FX 图转换为 Ascend IR 后由 GE 引擎编译执行不同npugraph_ex 采用捕获模式Capture Replay先把一段稳定的 NPU 执行过程kernel 序列与内存布局捕获到 Device 侧并固化为可回放的执行序列后续运行在 Guards 校验命中时直接低开销回放省去逐算子 Host 下发与重复编译开销。从使用体验看npugraph_ex 比 GE 图模式更接近 eager 语义模型代码中的 Stream / Event 等对象也更接近运行时显式对象因此适配路径相对轻量便于从 eager 代码逐步迁移参见 NPU 图模式优化原理。在 cann-recipes-infer 的图模式选择建议中npugraph_ex 与 GE 图模式的取舍如下详见 model-infer-graph-mode 技能说明特性npugraph_ex 后端 (aclgraph)GE 图模式 (Ascend IR)启用方式backendnpugraph_extorchair.get_npu_backend()实现原理捕获模式 (Capture Replay)FX 图转换为 Ascend IRGE 引擎编译执行成熟度试验特性暂不支持商用更成熟稳定PyTorch 版本需要 2.6.0无特殊要求支持场景在线推理通用场景类似技术torch.cuda.CUDAGraph传统图编译配置承载optionskwargsCompilerConfig.experimental_config缓存编译接口torch.npu.npugraph_ex.inference.cache_compile(...)torchair.inference.cache_compile(...)当前建议是优先选择 npugraph_ex以降低适配成本、保留更接近 eager 的开发体验。捕获与回放的关键环节npugraph_ex 的图复用依赖以下环节详见 docs/cann/zh/npu_graph_optimization.mdDynamo compilePython 级 JIT 编译器运行时重写 Python 字节码把前向中的 PyTorch 操作序列提取到一张 FX 图中再交给npugraph_ex后端编译Guards编译时生成对输入 shape、dtype、部分标量的假设每次执行前先校验——全部命中则复用已有图否则触发重新捕获与编译aclgraph Capture将 Stream 上的任务捕获到 Device 侧把稳定的 NPU 执行过程固化为可低开销回放的执行序列Input 处理回放前将图内 input 类参数的输入地址更新为实际运行地址私有格式如FRACTAL_NZ信息会被保留ReplayDevice 基于给定输入执行已捕获的图并得到输出。二、适用场景npugraph_ex 不是万能的它的设计目标决定了它的最佳战场在线推理场景追求简单快速适配希望以最小代码改动获得图模式收益熟悉 CUDAGraph 模式使用习惯与torch.cuda.CUDAGraph类似上手成本低LLM Decode 阶段固定 shape 的单 token 输入天然满足图静态化要求。与之对应LLM 的 Prefill 阶段通常不适合图模式输入序列长度动态变化、首 token 生成逻辑复杂、shape 不固定强制图化容易导致重编译因此 cann-recipes-infer 的执行框架中图模式主要用于 DecodePrefill 默认走 eager 路径参见 executor/core/model_worker/model_worker.py 中not is_prefill才使用 compiled model 的判断。三、使用约束接入 npugraph_ex 前务必先确认以下约束是否满足约束项说明PyTorch 版本需要 2.6.0 及以上版本支持场景在线推理场景不支持反向流程 capture随机数算子不支持 capturerandn、dropout 等动态控制流不支持需保证图静态Stream 同步不支持 stream sync 操作成熟度试验特性暂不支持商用产品此外forward内部不能出现.item()调用——将 Tensor 转换为 Python 标量会强制 Graph Break。若模型包含随机数算子、基于 Tensor 值的if/while或.item()需要先改造模型代码见下文LLM 模型改造要点。四、快速上手import torch import torch_npu model YourModel().npu() opt_model torch.compile(model, backendnpugraph_ex, fullgraphTrue, dynamicFalse) output opt_model(input_tensor)三步即可完成基本接入导入torch_npu模型必须已运行在 NPU 上且torch_npu已导入NPU 不支持aot_eager、inductor、cudagraphs等后端torch.compile指定后端backendnpugraph_ex并设置fullgraphTrue要求整图编译、dynamicFalse固定静态若模型存在list[int]形态的动态长度入参需dynamicTrue详见下文调用编译后的模型opt_model与原始model调用方式一致用相同的关键字入参即可触发图执行 / replay。在 cann-recipes-infer 框架中上述过程被封装为统一的compile_model_forward()见 executor/utils/graph_utils.py并自动处理集合通信入图与 Dynamo 配置tng.patch_for_hcom() # 集合通信入图PyTorch 2.6 之后通常可省略 torch._dynamo.config.inline_inbuilt_nn_modules False # 避免内建模块过度内联五、options 配置速查npugraph_ex 的编译配置通过torch.compile的optionskwargs 传入。以下配置覆盖了调试、FX 图优化、内存优化、性能优化与捕获控制五个维度opt_model torch.compile( model, backendnpugraph_ex, fullgraphTrue, options{ # 调试 force_eager: False, # 强制 eager 模式调试 # FX图优化 inplace_pass: True, # 原地操作优化 input_inplace_pass: True, # 输入原地优化 pattern_fusion_pass: True, # 算子融合 # 内存优化 reuse_graph_pool_in_same_fx: True, # 图池复用 clone_input: True, # 克隆输入 clone_output: False, # 克隆输出 use_graph_pool: None, # 图池配置 # 性能优化 static_kernel_compile: False, # 静态Kernel编译 remove_noop_ops: True, # 移除空操作 frozen_parameter: False, # 冻结参数 # 捕获控制 capture_limit: 64, # 重捕获次数限制 } )各配置项的用途调试类force_eagerTrue时后端强制以 eager 方式执行用于快速隔离图编译问题与模型代码问题FX 图优化类inplace_pass/input_inplace_pass允许对中间结果与输入做原地写以省内存、提带宽pattern_fusion_pass开启图内算子融合扩大后端优化视野内存优化类reuse_graph_pool_in_same_fx在同一 FX 图内复用图池clone_inputTrue为输入做克隆以避免捕获后输入被外部修改破坏图语义use_graph_pool可显式指定图池性能优化类static_kernel_compile开启静态 Kernel 编译该能力当前仅用于 npugraph_ex 相关路径GE 静态图中算子默认为静态算子remove_noop_ops消除冗余空操作frozen_parameter冻结参数减少每次回放的参数处理开销捕获控制类capture_limit限制同一图的重捕获次数防止异常场景下无限重编译。框架中的 options 组装方式cann-recipes-infer 框架在 graph_utils.py 中组装了 npugraph_ex 的 options并自动挂接缓存编译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: [.*]} } if enable_cache_compile: compiled torch.npu.npugraph_ex.inference.cache_compile( model_forward, cache_dircache_dir, dynamicenable_dynamic_graph, optionscompile_options) else: compiled torch.compile(model_forward, dynamicenable_dynamic_graph, fullgraphTrue, backendnpugraph_ex, optionscompile_options)注意其中的关键差异框架默认frozen_parameterTrue与速查表中的False默认值不同——冻结参数对 LLM Decode 这类参数固定不变的推理场景收益明显static_kernel_compile由配置项enable_static_kernel驱动仅支持 npugraph_exinference_config.py 中会做校验enable_static_kernelTrue而exe_mode ! npugraph_ex时直接报错enable_cache_compileTrue时不再走torch.compile而是改用torch.npu.npugraph_ex.inference.cache_compile首次编译产物落盘到cache_dir后续启动命中缓存后跳过 Dynamo compile 与 Capture。六、核心 APIAPI用途说明compile_fx()自定义 backend通过torch.compile(backend...)机制注册自定义图编译后端register_replacement()自定义算子融合注册图内子图替换规则实现自定义融合模式cache_compile()编译缓存torch.npu.npugraph_ex.inference.cache_compile把编译产物持久化到磁盘跨进程复用limit_core_num()限核功能通过 scope 限定 AI Core / Vector Core 数量用于多流、负载均衡场景这些 API 的完整参数与行为说明可参考 TorchAir 的 npugraph_ex API 文档torch.npu.npugraph_ex、compile_fx、register_replacement、inference (cache_compile, readable_cache)、scope (limit_core_num)五个模块。编译缓存的使用时机cache_compile适合在图模式已跑通、输入 shape / guard / 通信域稳定后启用用于降低冷启动或多次拉起时的编译耗时它不用于修复 Graph Break 或非预期重编译。使用要点详见 model-infer-graph-mode 技能说明使用后原torch.compile编译流程不再需要被缓存的函数应是 module method、未被其他装饰器修饰、能形成 full graph且同一缓存函数只能触发一次 Dynamo tracePrefill / Decode 或 guard 不同的场景应拆分封装避免相互污染缓存若模型代码、输入规格、分布式 rank/world_size、CANN/torch_npu 版本发生变化需重新生成或清理缓存。七、问题定界流程接入图模式后若出现问题按以下流程逐级定界先排除用户脚本问题再怀疑图模式问题问题发生 │ ├─→ aot_eager 验证 ──失败──→ 修复用户脚本 │ ↓ 正常 │ ├─→ force_eagerTrue ──失败──→ 修复用户脚本 │ ↓ 正常 │ └─→ npugraph_ex/FX图问题 ├── 重编译问题 → 阅读 LLM 模型改造指南 npugraph_ex 指南 ├── Graph Break 问题 → 阅读 TorchAir 在线文档中的典型案例 └── 其他问题 → 阅读对应模式文档npugraph_ex-guide.md 或 ge-graph-guide.md定界要点aot_eager 验证aot_eager 是 PyTorch 的图编译前端绕过后端图编译直接执行。若 aot_eager 都失败说明是用户脚本模型代码问题force_eager 验证force_eagerTrue让 npugraph_ex 后端强制 eager 执行。若此时失败同样先修复用户脚本只有 eager 全部正常才需要排查 npugraph_ex 捕获、FX 图优化或重编译问题问题归属图模式后按重编译 / Graph Break / 其他分类参考下文常见问题逐项处理。八、在 cann-recipes-infer 框架中的接入方式8.1 通过 YAML 配置启用框架通过 inference_config.py 的ModelConfig统一管理图模式开关相关字段如下配置字段类型默认值说明exe_modestreager执行模式可选eager/ge_graph/npugraph_exenable_cache_compileboolFalse是否启用编译缓存落盘enable_static_kernelboolFalse是否开启静态 Kernel 编译仅支持 npugraph_exenable_dynamic_graphboolTrue是否使用动态图编译仓库真实样例 deepseek_v3.2_exp_rank_64_64ep_w8a8c8_decode_npugraphex_benchmark.yaml 展示了完整用法model_config: model_name: deepseek_v3_2_exp model_path: /data/models/DeepSeek-V3.2-Exp-W8A8C8 exe_mode: npugraph_ex # [eager, npugraph_ex, ge_graph], mode of decode with_ckpt: True # [False, True] next_n: 0 force_eplb: False # [False, True] enable_weight_nz: True enable_profiler: False # [False, True] enable_cache_compile: False # [False, True] enable_static_kernel: True # [False, True] only support npugraph_ex custom_params: enable_multi_streams: True # [False, True] enable_offload: False其中exe_mode: npugraph_ex表示Decode 阶段使用 npugraph_ex 图模式Prefill 阶段仍走 eager。从配置校验逻辑inference_config.py可以看到两条隐含约束enable_static_kernelTrue时exe_mode必须是npugraph_ex否则直接抛错exe_mode npugraph_ex时框架会自动设置环境变量TASK_QUEUE_ENABLE1npugraph_ex 仅支持 TASK_QUEUE_ENABLE 取 0 或 1eager 模式则默认设为 2。8.2 框架内部调用链框架中图模式的编译与调用链如下model_worker.pycompile_model()在 warm-up 阶段调用compile_model_forward()L470-L520并将编译结果绑定为self.model_compiled编译对象是模型的main_decode主模型或mtp_decodedraft 模型方法——框架要求模型必须实现对应方法否则报错。主模型与 draft 模型使用不同 compile 接口避免互相触发不必要的重编译正式推理时L445-L460仅在exe_mode in [ge_graph, npugraph_ex]且非 prefill时走compiled_model(**model_inputs)其余情况走 eager 的self.model(**model_inputs)使用缓存编译且非 warm-up 时调用被包在torch.compiler.set_stance(skip_guard_eval_unsafeTrue)中跳过 guard 求值以降低每次回放的 Host 开销。这种设计保证了warm-up 阶段完成 Dynamo compile 与 aclgraph Capture 并缓存图正式推理直接复用图、仅走 Guards 校验、Input 处理和 Replay图编译开销不落在正式推理关键路径上。九、LLM 模型改造要点图模式适配的本质是将动态变化的东西提取为模型输入模型内部尽量保证静态。对于 LLM 推理模型必须严格区分 prefill 和 decode 阶段详见 LLM 模型改造指南动态因素问题表现解决思路内存地址变化Guard 失败、重编译预分配固定大小原地更新Shape 变化图中断、多次编译固定 shape 或通过参数控制Python 控制流Graph Break使用 Tensor 操作或模式参数.item()调用强制 Graph Break保持 Tensor 或外部传入9.1 Prefill 与 Decode 阶段限制阶段是否支持图模式原因Prefill禁止使用输入长度动态变化、首 token 生成逻辑复杂、shape 不固定Decode推荐使用输入长度固定通常为 1、shape 稳定、适合图捕获实现建议将 prefill 和 decode 的 forward 逻辑分离成不同方法仅对 decode 方法应用torch.compile图模式prefill 保持 eagerclass YourModel: def prefill(self, input_ids, ...): Prefill 阶段使用 eager 模式 # 输入长度动态变化不适合图模式 return self._forward(input_ids, ...) def decode(self, input_ids, ...): Decode 阶段可使用图模式 # 输入长度固定通常为 1适合图捕获 return self._forward(input_ids, ...) # 仅对 decode 方法应用图模式model.prefill 保持 eager model.decode torch.compile(model.decode, backendnpugraph_ex, ...)9.2 KV Cache 模块改造改造目标消除 KV Cache 动态扩展导致的 shape 变化实现固定大小 cache 的原地更新。# 问题模式动态扩展 KV cache key torch.cat([past_key, new_key], dim1) # shape 变化 # 改造模式固定大小预分配 原地更新 def _init_kv_cache(self, batch_size, max_seq_len, device): cache_shape (batch_size, 1, max_seq_len, head_dim) self.kv_cache torch.zeros(cache_shape, dtypedtype, devicedevice) def forward(self, ..., kv_len, past_key_value): torch_npu.scatter_update_(past_key_cache, kv_len, new_key_states, dim-2)常见问题对照问题现象根因解决方案每次 decode 触发重编译torch.cat扩展 KV cache预分配固定大小原地更新内存占用过大预分配浪费结合 PagedAttention 按 block 管理返回 KV cache 开销大图模式下返回大量 tensor已原地更新无需返回9.3 动态信息外部化设计将动态变化的信息从模型内部移到输入参数是图模式改造的核心手法动态信息内部计算应避免外部传入推荐位置索引position_ids torch.arange(seq_len)作为参数传入序列长度seq_len hidden_states.size(1)actual_seq_lengths参数写入位置内部计算kv_lenkv_len参数模式切换内部判断is_prefill参数一个图模式友好的 forward 签名参考完整示例见 llm-model-guide.mddef forward( self, input_ids: torch.LongTensor, position_ids: Optional[torch.LongTensor] None, # 位置相关Tensor 形式支持图追踪 kv_len: Optional[torch.IntTensor] None, # KV 写入位置 actual_seq_lengths_kv: Optional[List[int]] None, # 序列长度List[int] 传给 NPU 算子 actual_seq_lengths_q: Optional[List[int]] None, is_prefill: bool False, # 模式控制 past_key_values: Optional[Tuple[torch.Tensor]] None, cos: Optional[torch.Tensor] None, # 预计算的 cos/sin sin: Optional[torch.Tensor] None, ... ): pass9.4 重编译问题定位与解决如果图模式性能劣化必须先定位是否发生了重编译# 开启重编译日志 torch._logging.set_logs(recompilesTrue) # 运行模型如果发生重编译会打印类似 # [recompiles] Recompiling function func_name for reason: reason解决思路dynamicFalse: 检测到重编译 │ └─→ 分析重编译原因 ├── 固定 shape 但仍重编译 → dynamicFalse skip_guard_eval_unsafeTrue └── 输入 shape 变化 → dynamicTrue框架中正式推理时的skip_guard_eval_unsafeTruemodel_worker.py正是这一思路的实现——确认 guard 稳定后跳过 guard 求值进一步压低 Host 开销。十、npugraph_ex 与 FIA 融合算子的配合图模式与 Flash Attention 融合算子结合时actual_seq_lengths参数的处理是最常见的出错点。npugraph_ex 模式下的推荐配置详见 model-infer-graph-mode 技能说明图模式FA 接口来源actual_seq_lengths 类型dynamic 设置npugraph_extorch_npu FA 接口list[int]dynamicTruenpugraph_extorch_npu FA 接口如有 Tensor 接口TensordynamicFalse需确认接口支持GE推荐torchair FA 接口TensordynamicFalseGE不推荐torch_npu FA 接口list[int]dynamicTruemark_staticnpugraph_ex 常用配置为list[int]类型 dynamicTrueimport torch import torch_npu # 使用 torch_npu 的 FA 接口actual_seq_lengths 为 list[int] 类型 attn_output torch.ops.npu.npu_fused_infer_attention_score( query, key, value, actual_seq_lengths[seq_len], # list[int] 类型 actual_seq_lengths_kv[kv_len], # ... 其他参数 ) # 必须配置 dynamicTrue opt_model torch.compile(model, backendnpugraph_ex, dynamicTrue)需要理解的是npugraph_ex 常保持dynamicTrue并不是因为图本身必须动态而是与当前推理场景中部分 FIA 算子接口有关——部分actual_seq_lengths入参仍以list[int]形式传入若强行静态化容易触发重编译。后续算子接口补齐 Tensor 输入后这类配置可以继续收敛参见 NPU 图模式优化原理。仓库中 deepseek_v4 等不存在这类 list 输入的模型建议选择静态图即enable_dynamic_graphFalse以获得更稳定的图复用和更低的下发开销。此外LLM Decode 场景下框架约定npugraph_ex 的 decode 路径使用 Host list 长度字段即从ForwardMetaData取actual_seq_lengths_list_kv/q/actual_seq_lengths_cu_list_kv/qList[int]形态传给静态图路径普通 eager / GE 路径仍用对应的 Tensor 字段——接入时务必对齐这一约定避免类型不匹配。十一、常见问题1. 如何判断是否应该使用 npugraph_ex适合LLM decode 阶段、固定 shape 推理、简单快速适配不适合需要动态 shape、生产环境稳定性优先、训练场景。2. 报错不支持 capture怎么办检查代码中是否包含随机数算子randn、dropout动态控制流基于 tensor 值的if/while.item()调用。这三类写法都会破坏图的静态性需要按动态信息外部化思路改写为 Tensor 逻辑或外部传入。3. 性能劣化怎么办开启重编译日志torch._logging.set_logs(recompilesTrue)检查是否发生重编译分析原因固定 shape 仍重编译 → 用skip_guard_eval_unsafeTrue输入 shape 变化 →dynamicTrue固定输入 shape、KV Cache 地址和常驻 buffer观察重编译是否消失参考 LLM 模型改造指南 中的改造指南逐项核对。4. 高频问题速查表现象常见根因处理建议编译前报错eager 路径本身不正确或输入 shape、dtype、长度组织方式不稳定先单独验证 eager再检查图模式输入和前向参数图捕获中断.item()、Tensor 驱动的 Python 分支、print、自定义算子未适配改写为 Tensor 逻辑或补齐入图适配Decode 性能没有提升甚至变差发生重编译或图执行没有命中预期优化打开重编译日志检查 guard 变化、shape、地址和缓存命中情况actual_seq_lengths类型报错图模式与 FIA 接口不匹配ge_graph优先 Tensor torchair.opsnpugraph_ex对齐list[int]方案enable_static_kernel报错开启模式不匹配该选项只在npugraph_ex模式下使用cache compile 不生效缓存目录、输入规格或编译配置发生变化固定 cache 目录保持模型代码、函数名称、Tensor shape、dtype 和配置稳定图模式下精度异常KV Cache、FA 接口或原地更新语义变化回退到 eager 对齐再逐项恢复图模式优化十二、相关文档LLM 模型改造指南llm-model-guide.mdLLM 适配优先阅读GE 图模式指南ge-graph-guide.md图模式适配技能总览model-infer-graph-mode/SKILL.md方案设计、图中断修复、FA 参数配置、验证测试全流程NPU 图模式优化原理docs/cann/zh/npu_graph_optimization.mdeager 与图模式对比、两类图模式关系、编译缓存原理框架图编译实现executor/utils/graph_utils.py图模式相关配置定义executor/core/config/inference_config.pynpugraph_ex 真实配置样例deepseek_v3.2_exp_rank_64_64ep_w8a8c8_decode_npugraphex_benchmark.yaml。图模式通常不是单独存在的优化开关实际项目中常与缓存编译、静态 Kernel、多流NPU 多流原理、预取NPU Prefetch 原理、superkernelNPU superkernel 原理等能力组合使用。推荐的推进顺序是先跑通 eager 并完成功能与精度验证 → 打开图模式消除 graph break 和 recompile → 对比 eager 与 graph 输出至少覆盖一轮 Prefill 和多轮 Decode→ 再按需开启缓存编译、静态 Kernel、多流、限核或模型自带的 superkernel 开关。【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 CANN 推理加速实战：npugraph_ex 图模式后端使用与 LLM Decode 适配指南
CANN 推理加速实战npugraph_ex 图模式后端使用与 LLM Decode 适配指南【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infernpugraph_ex 是基于torch.compile的 aclgraph 图模式加速方案通过捕获模式Capture Replay将模型前向固化为可在昇腾 NPU 上低开销回放的计算图。本指南以其完整使用手册为主体结合 cann-recipes-infer 开源仓库中的框架接入实现graph_utils.py、推理配置inference_config.py与真实模型配置样例讲解 npugraph_ex 的适用场景、使用约束、快速上手、options 配置速查、问题定界流程与常见问题排查帮助读者在 LLM 在线推理尤其是 Decode 阶段中快速完成图模式适配。一、npugraph_ex 是什么npugraph_ex 是昇腾 NPU 上一种基于torch.compile的图模式后端。与 GE 图模式将 FX 图转换为 Ascend IR 后由 GE 引擎编译执行不同npugraph_ex 采用捕获模式Capture Replay先把一段稳定的 NPU 执行过程kernel 序列与内存布局捕获到 Device 侧并固化为可回放的执行序列后续运行在 Guards 校验命中时直接低开销回放省去逐算子 Host 下发与重复编译开销。从使用体验看npugraph_ex 比 GE 图模式更接近 eager 语义模型代码中的 Stream / Event 等对象也更接近运行时显式对象因此适配路径相对轻量便于从 eager 代码逐步迁移参见 NPU 图模式优化原理。在 cann-recipes-infer 的图模式选择建议中npugraph_ex 与 GE 图模式的取舍如下详见 model-infer-graph-mode 技能说明特性npugraph_ex 后端 (aclgraph)GE 图模式 (Ascend IR)启用方式backendnpugraph_extorchair.get_npu_backend()实现原理捕获模式 (Capture Replay)FX 图转换为 Ascend IRGE 引擎编译执行成熟度试验特性暂不支持商用更成熟稳定PyTorch 版本需要 2.6.0无特殊要求支持场景在线推理通用场景类似技术torch.cuda.CUDAGraph传统图编译配置承载optionskwargsCompilerConfig.experimental_config缓存编译接口torch.npu.npugraph_ex.inference.cache_compile(...)torchair.inference.cache_compile(...)当前建议是优先选择 npugraph_ex以降低适配成本、保留更接近 eager 的开发体验。捕获与回放的关键环节npugraph_ex 的图复用依赖以下环节详见 docs/cann/zh/npu_graph_optimization.mdDynamo compilePython 级 JIT 编译器运行时重写 Python 字节码把前向中的 PyTorch 操作序列提取到一张 FX 图中再交给npugraph_ex后端编译Guards编译时生成对输入 shape、dtype、部分标量的假设每次执行前先校验——全部命中则复用已有图否则触发重新捕获与编译aclgraph Capture将 Stream 上的任务捕获到 Device 侧把稳定的 NPU 执行过程固化为可低开销回放的执行序列Input 处理回放前将图内 input 类参数的输入地址更新为实际运行地址私有格式如FRACTAL_NZ信息会被保留ReplayDevice 基于给定输入执行已捕获的图并得到输出。二、适用场景npugraph_ex 不是万能的它的设计目标决定了它的最佳战场在线推理场景追求简单快速适配希望以最小代码改动获得图模式收益熟悉 CUDAGraph 模式使用习惯与torch.cuda.CUDAGraph类似上手成本低LLM Decode 阶段固定 shape 的单 token 输入天然满足图静态化要求。与之对应LLM 的 Prefill 阶段通常不适合图模式输入序列长度动态变化、首 token 生成逻辑复杂、shape 不固定强制图化容易导致重编译因此 cann-recipes-infer 的执行框架中图模式主要用于 DecodePrefill 默认走 eager 路径参见 executor/core/model_worker/model_worker.py 中not is_prefill才使用 compiled model 的判断。三、使用约束接入 npugraph_ex 前务必先确认以下约束是否满足约束项说明PyTorch 版本需要 2.6.0 及以上版本支持场景在线推理场景不支持反向流程 capture随机数算子不支持 capturerandn、dropout 等动态控制流不支持需保证图静态Stream 同步不支持 stream sync 操作成熟度试验特性暂不支持商用产品此外forward内部不能出现.item()调用——将 Tensor 转换为 Python 标量会强制 Graph Break。若模型包含随机数算子、基于 Tensor 值的if/while或.item()需要先改造模型代码见下文LLM 模型改造要点。四、快速上手import torch import torch_npu model YourModel().npu() opt_model torch.compile(model, backendnpugraph_ex, fullgraphTrue, dynamicFalse) output opt_model(input_tensor)三步即可完成基本接入导入torch_npu模型必须已运行在 NPU 上且torch_npu已导入NPU 不支持aot_eager、inductor、cudagraphs等后端torch.compile指定后端backendnpugraph_ex并设置fullgraphTrue要求整图编译、dynamicFalse固定静态若模型存在list[int]形态的动态长度入参需dynamicTrue详见下文调用编译后的模型opt_model与原始model调用方式一致用相同的关键字入参即可触发图执行 / replay。在 cann-recipes-infer 框架中上述过程被封装为统一的compile_model_forward()见 executor/utils/graph_utils.py并自动处理集合通信入图与 Dynamo 配置tng.patch_for_hcom() # 集合通信入图PyTorch 2.6 之后通常可省略 torch._dynamo.config.inline_inbuilt_nn_modules False # 避免内建模块过度内联五、options 配置速查npugraph_ex 的编译配置通过torch.compile的optionskwargs 传入。以下配置覆盖了调试、FX 图优化、内存优化、性能优化与捕获控制五个维度opt_model torch.compile( model, backendnpugraph_ex, fullgraphTrue, options{ # 调试 force_eager: False, # 强制 eager 模式调试 # FX图优化 inplace_pass: True, # 原地操作优化 input_inplace_pass: True, # 输入原地优化 pattern_fusion_pass: True, # 算子融合 # 内存优化 reuse_graph_pool_in_same_fx: True, # 图池复用 clone_input: True, # 克隆输入 clone_output: False, # 克隆输出 use_graph_pool: None, # 图池配置 # 性能优化 static_kernel_compile: False, # 静态Kernel编译 remove_noop_ops: True, # 移除空操作 frozen_parameter: False, # 冻结参数 # 捕获控制 capture_limit: 64, # 重捕获次数限制 } )各配置项的用途调试类force_eagerTrue时后端强制以 eager 方式执行用于快速隔离图编译问题与模型代码问题FX 图优化类inplace_pass/input_inplace_pass允许对中间结果与输入做原地写以省内存、提带宽pattern_fusion_pass开启图内算子融合扩大后端优化视野内存优化类reuse_graph_pool_in_same_fx在同一 FX 图内复用图池clone_inputTrue为输入做克隆以避免捕获后输入被外部修改破坏图语义use_graph_pool可显式指定图池性能优化类static_kernel_compile开启静态 Kernel 编译该能力当前仅用于 npugraph_ex 相关路径GE 静态图中算子默认为静态算子remove_noop_ops消除冗余空操作frozen_parameter冻结参数减少每次回放的参数处理开销捕获控制类capture_limit限制同一图的重捕获次数防止异常场景下无限重编译。框架中的 options 组装方式cann-recipes-infer 框架在 graph_utils.py 中组装了 npugraph_ex 的 options并自动挂接缓存编译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: [.*]} } if enable_cache_compile: compiled torch.npu.npugraph_ex.inference.cache_compile( model_forward, cache_dircache_dir, dynamicenable_dynamic_graph, optionscompile_options) else: compiled torch.compile(model_forward, dynamicenable_dynamic_graph, fullgraphTrue, backendnpugraph_ex, optionscompile_options)注意其中的关键差异框架默认frozen_parameterTrue与速查表中的False默认值不同——冻结参数对 LLM Decode 这类参数固定不变的推理场景收益明显static_kernel_compile由配置项enable_static_kernel驱动仅支持 npugraph_exinference_config.py 中会做校验enable_static_kernelTrue而exe_mode ! npugraph_ex时直接报错enable_cache_compileTrue时不再走torch.compile而是改用torch.npu.npugraph_ex.inference.cache_compile首次编译产物落盘到cache_dir后续启动命中缓存后跳过 Dynamo compile 与 Capture。六、核心 APIAPI用途说明compile_fx()自定义 backend通过torch.compile(backend...)机制注册自定义图编译后端register_replacement()自定义算子融合注册图内子图替换规则实现自定义融合模式cache_compile()编译缓存torch.npu.npugraph_ex.inference.cache_compile把编译产物持久化到磁盘跨进程复用limit_core_num()限核功能通过 scope 限定 AI Core / Vector Core 数量用于多流、负载均衡场景这些 API 的完整参数与行为说明可参考 TorchAir 的 npugraph_ex API 文档torch.npu.npugraph_ex、compile_fx、register_replacement、inference (cache_compile, readable_cache)、scope (limit_core_num)五个模块。编译缓存的使用时机cache_compile适合在图模式已跑通、输入 shape / guard / 通信域稳定后启用用于降低冷启动或多次拉起时的编译耗时它不用于修复 Graph Break 或非预期重编译。使用要点详见 model-infer-graph-mode 技能说明使用后原torch.compile编译流程不再需要被缓存的函数应是 module method、未被其他装饰器修饰、能形成 full graph且同一缓存函数只能触发一次 Dynamo tracePrefill / Decode 或 guard 不同的场景应拆分封装避免相互污染缓存若模型代码、输入规格、分布式 rank/world_size、CANN/torch_npu 版本发生变化需重新生成或清理缓存。七、问题定界流程接入图模式后若出现问题按以下流程逐级定界先排除用户脚本问题再怀疑图模式问题问题发生 │ ├─→ aot_eager 验证 ──失败──→ 修复用户脚本 │ ↓ 正常 │ ├─→ force_eagerTrue ──失败──→ 修复用户脚本 │ ↓ 正常 │ └─→ npugraph_ex/FX图问题 ├── 重编译问题 → 阅读 LLM 模型改造指南 npugraph_ex 指南 ├── Graph Break 问题 → 阅读 TorchAir 在线文档中的典型案例 └── 其他问题 → 阅读对应模式文档npugraph_ex-guide.md 或 ge-graph-guide.md定界要点aot_eager 验证aot_eager 是 PyTorch 的图编译前端绕过后端图编译直接执行。若 aot_eager 都失败说明是用户脚本模型代码问题force_eager 验证force_eagerTrue让 npugraph_ex 后端强制 eager 执行。若此时失败同样先修复用户脚本只有 eager 全部正常才需要排查 npugraph_ex 捕获、FX 图优化或重编译问题问题归属图模式后按重编译 / Graph Break / 其他分类参考下文常见问题逐项处理。八、在 cann-recipes-infer 框架中的接入方式8.1 通过 YAML 配置启用框架通过 inference_config.py 的ModelConfig统一管理图模式开关相关字段如下配置字段类型默认值说明exe_modestreager执行模式可选eager/ge_graph/npugraph_exenable_cache_compileboolFalse是否启用编译缓存落盘enable_static_kernelboolFalse是否开启静态 Kernel 编译仅支持 npugraph_exenable_dynamic_graphboolTrue是否使用动态图编译仓库真实样例 deepseek_v3.2_exp_rank_64_64ep_w8a8c8_decode_npugraphex_benchmark.yaml 展示了完整用法model_config: model_name: deepseek_v3_2_exp model_path: /data/models/DeepSeek-V3.2-Exp-W8A8C8 exe_mode: npugraph_ex # [eager, npugraph_ex, ge_graph], mode of decode with_ckpt: True # [False, True] next_n: 0 force_eplb: False # [False, True] enable_weight_nz: True enable_profiler: False # [False, True] enable_cache_compile: False # [False, True] enable_static_kernel: True # [False, True] only support npugraph_ex custom_params: enable_multi_streams: True # [False, True] enable_offload: False其中exe_mode: npugraph_ex表示Decode 阶段使用 npugraph_ex 图模式Prefill 阶段仍走 eager。从配置校验逻辑inference_config.py可以看到两条隐含约束enable_static_kernelTrue时exe_mode必须是npugraph_ex否则直接抛错exe_mode npugraph_ex时框架会自动设置环境变量TASK_QUEUE_ENABLE1npugraph_ex 仅支持 TASK_QUEUE_ENABLE 取 0 或 1eager 模式则默认设为 2。8.2 框架内部调用链框架中图模式的编译与调用链如下model_worker.pycompile_model()在 warm-up 阶段调用compile_model_forward()L470-L520并将编译结果绑定为self.model_compiled编译对象是模型的main_decode主模型或mtp_decodedraft 模型方法——框架要求模型必须实现对应方法否则报错。主模型与 draft 模型使用不同 compile 接口避免互相触发不必要的重编译正式推理时L445-L460仅在exe_mode in [ge_graph, npugraph_ex]且非 prefill时走compiled_model(**model_inputs)其余情况走 eager 的self.model(**model_inputs)使用缓存编译且非 warm-up 时调用被包在torch.compiler.set_stance(skip_guard_eval_unsafeTrue)中跳过 guard 求值以降低每次回放的 Host 开销。这种设计保证了warm-up 阶段完成 Dynamo compile 与 aclgraph Capture 并缓存图正式推理直接复用图、仅走 Guards 校验、Input 处理和 Replay图编译开销不落在正式推理关键路径上。九、LLM 模型改造要点图模式适配的本质是将动态变化的东西提取为模型输入模型内部尽量保证静态。对于 LLM 推理模型必须严格区分 prefill 和 decode 阶段详见 LLM 模型改造指南动态因素问题表现解决思路内存地址变化Guard 失败、重编译预分配固定大小原地更新Shape 变化图中断、多次编译固定 shape 或通过参数控制Python 控制流Graph Break使用 Tensor 操作或模式参数.item()调用强制 Graph Break保持 Tensor 或外部传入9.1 Prefill 与 Decode 阶段限制阶段是否支持图模式原因Prefill禁止使用输入长度动态变化、首 token 生成逻辑复杂、shape 不固定Decode推荐使用输入长度固定通常为 1、shape 稳定、适合图捕获实现建议将 prefill 和 decode 的 forward 逻辑分离成不同方法仅对 decode 方法应用torch.compile图模式prefill 保持 eagerclass YourModel: def prefill(self, input_ids, ...): Prefill 阶段使用 eager 模式 # 输入长度动态变化不适合图模式 return self._forward(input_ids, ...) def decode(self, input_ids, ...): Decode 阶段可使用图模式 # 输入长度固定通常为 1适合图捕获 return self._forward(input_ids, ...) # 仅对 decode 方法应用图模式model.prefill 保持 eager model.decode torch.compile(model.decode, backendnpugraph_ex, ...)9.2 KV Cache 模块改造改造目标消除 KV Cache 动态扩展导致的 shape 变化实现固定大小 cache 的原地更新。# 问题模式动态扩展 KV cache key torch.cat([past_key, new_key], dim1) # shape 变化 # 改造模式固定大小预分配 原地更新 def _init_kv_cache(self, batch_size, max_seq_len, device): cache_shape (batch_size, 1, max_seq_len, head_dim) self.kv_cache torch.zeros(cache_shape, dtypedtype, devicedevice) def forward(self, ..., kv_len, past_key_value): torch_npu.scatter_update_(past_key_cache, kv_len, new_key_states, dim-2)常见问题对照问题现象根因解决方案每次 decode 触发重编译torch.cat扩展 KV cache预分配固定大小原地更新内存占用过大预分配浪费结合 PagedAttention 按 block 管理返回 KV cache 开销大图模式下返回大量 tensor已原地更新无需返回9.3 动态信息外部化设计将动态变化的信息从模型内部移到输入参数是图模式改造的核心手法动态信息内部计算应避免外部传入推荐位置索引position_ids torch.arange(seq_len)作为参数传入序列长度seq_len hidden_states.size(1)actual_seq_lengths参数写入位置内部计算kv_lenkv_len参数模式切换内部判断is_prefill参数一个图模式友好的 forward 签名参考完整示例见 llm-model-guide.mddef forward( self, input_ids: torch.LongTensor, position_ids: Optional[torch.LongTensor] None, # 位置相关Tensor 形式支持图追踪 kv_len: Optional[torch.IntTensor] None, # KV 写入位置 actual_seq_lengths_kv: Optional[List[int]] None, # 序列长度List[int] 传给 NPU 算子 actual_seq_lengths_q: Optional[List[int]] None, is_prefill: bool False, # 模式控制 past_key_values: Optional[Tuple[torch.Tensor]] None, cos: Optional[torch.Tensor] None, # 预计算的 cos/sin sin: Optional[torch.Tensor] None, ... ): pass9.4 重编译问题定位与解决如果图模式性能劣化必须先定位是否发生了重编译# 开启重编译日志 torch._logging.set_logs(recompilesTrue) # 运行模型如果发生重编译会打印类似 # [recompiles] Recompiling function func_name for reason: reason解决思路dynamicFalse: 检测到重编译 │ └─→ 分析重编译原因 ├── 固定 shape 但仍重编译 → dynamicFalse skip_guard_eval_unsafeTrue └── 输入 shape 变化 → dynamicTrue框架中正式推理时的skip_guard_eval_unsafeTruemodel_worker.py正是这一思路的实现——确认 guard 稳定后跳过 guard 求值进一步压低 Host 开销。十、npugraph_ex 与 FIA 融合算子的配合图模式与 Flash Attention 融合算子结合时actual_seq_lengths参数的处理是最常见的出错点。npugraph_ex 模式下的推荐配置详见 model-infer-graph-mode 技能说明图模式FA 接口来源actual_seq_lengths 类型dynamic 设置npugraph_extorch_npu FA 接口list[int]dynamicTruenpugraph_extorch_npu FA 接口如有 Tensor 接口TensordynamicFalse需确认接口支持GE推荐torchair FA 接口TensordynamicFalseGE不推荐torch_npu FA 接口list[int]dynamicTruemark_staticnpugraph_ex 常用配置为list[int]类型 dynamicTrueimport torch import torch_npu # 使用 torch_npu 的 FA 接口actual_seq_lengths 为 list[int] 类型 attn_output torch.ops.npu.npu_fused_infer_attention_score( query, key, value, actual_seq_lengths[seq_len], # list[int] 类型 actual_seq_lengths_kv[kv_len], # ... 其他参数 ) # 必须配置 dynamicTrue opt_model torch.compile(model, backendnpugraph_ex, dynamicTrue)需要理解的是npugraph_ex 常保持dynamicTrue并不是因为图本身必须动态而是与当前推理场景中部分 FIA 算子接口有关——部分actual_seq_lengths入参仍以list[int]形式传入若强行静态化容易触发重编译。后续算子接口补齐 Tensor 输入后这类配置可以继续收敛参见 NPU 图模式优化原理。仓库中 deepseek_v4 等不存在这类 list 输入的模型建议选择静态图即enable_dynamic_graphFalse以获得更稳定的图复用和更低的下发开销。此外LLM Decode 场景下框架约定npugraph_ex 的 decode 路径使用 Host list 长度字段即从ForwardMetaData取actual_seq_lengths_list_kv/q/actual_seq_lengths_cu_list_kv/qList[int]形态传给静态图路径普通 eager / GE 路径仍用对应的 Tensor 字段——接入时务必对齐这一约定避免类型不匹配。十一、常见问题1. 如何判断是否应该使用 npugraph_ex适合LLM decode 阶段、固定 shape 推理、简单快速适配不适合需要动态 shape、生产环境稳定性优先、训练场景。2. 报错不支持 capture怎么办检查代码中是否包含随机数算子randn、dropout动态控制流基于 tensor 值的if/while.item()调用。这三类写法都会破坏图的静态性需要按动态信息外部化思路改写为 Tensor 逻辑或外部传入。3. 性能劣化怎么办开启重编译日志torch._logging.set_logs(recompilesTrue)检查是否发生重编译分析原因固定 shape 仍重编译 → 用skip_guard_eval_unsafeTrue输入 shape 变化 →dynamicTrue固定输入 shape、KV Cache 地址和常驻 buffer观察重编译是否消失参考 LLM 模型改造指南 中的改造指南逐项核对。4. 高频问题速查表现象常见根因处理建议编译前报错eager 路径本身不正确或输入 shape、dtype、长度组织方式不稳定先单独验证 eager再检查图模式输入和前向参数图捕获中断.item()、Tensor 驱动的 Python 分支、print、自定义算子未适配改写为 Tensor 逻辑或补齐入图适配Decode 性能没有提升甚至变差发生重编译或图执行没有命中预期优化打开重编译日志检查 guard 变化、shape、地址和缓存命中情况actual_seq_lengths类型报错图模式与 FIA 接口不匹配ge_graph优先 Tensor torchair.opsnpugraph_ex对齐list[int]方案enable_static_kernel报错开启模式不匹配该选项只在npugraph_ex模式下使用cache compile 不生效缓存目录、输入规格或编译配置发生变化固定 cache 目录保持模型代码、函数名称、Tensor shape、dtype 和配置稳定图模式下精度异常KV Cache、FA 接口或原地更新语义变化回退到 eager 对齐再逐项恢复图模式优化十二、相关文档LLM 模型改造指南llm-model-guide.mdLLM 适配优先阅读GE 图模式指南ge-graph-guide.md图模式适配技能总览model-infer-graph-mode/SKILL.md方案设计、图中断修复、FA 参数配置、验证测试全流程NPU 图模式优化原理docs/cann/zh/npu_graph_optimization.mdeager 与图模式对比、两类图模式关系、编译缓存原理框架图编译实现executor/utils/graph_utils.py图模式相关配置定义executor/core/config/inference_config.pynpugraph_ex 真实配置样例deepseek_v3.2_exp_rank_64_64ep_w8a8c8_decode_npugraphex_benchmark.yaml。图模式通常不是单独存在的优化开关实际项目中常与缓存编译、静态 Kernel、多流NPU 多流原理、预取NPU Prefetch 原理、superkernelNPU superkernel 原理等能力组合使用。推荐的推进顺序是先跑通 eager 并完成功能与精度验证 → 打开图模式消除 graph break 和 recompile → 对比 eager 与 graph 输出至少覆盖一轮 Prefill 和多轮 Decode→ 再按需开启缓存编译、静态 Kernel、多流、限核或模型自带的 superkernel 开关。【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考