External Cannbot Model Model Infer Graph Mode

图模式适配优化技能

提供 npugraph_ex 和 GE 两种图模式的适配指南，覆盖方案设计、图中断修复、FA 参数配置和验证测试。

重要原则

• 前置条件：模型必须在 NPU 上运行且已导入 torch_npu
• 图模式仅适用于 Decode 阶段：Prefill 阶段输入长度动态变化，不适合图模式，保持 eager 执行
• 保持模型完整性：不为适配图模式而简化模型逻辑
• NPU 不支持的后端：aot_eager、inductor、cudagraphs 等不可用
• 固定 tensor 图外预创建：atten_mask、KV cache 等推理全程不变的 tensor 在图外预创建，用 torch._dynamo.mark_static() 标记
• LLM Decode 重编译注意：kv_len/actual_seq_lengths_kv 每步变化，npugraph_ex 可能触发重编译，遇到 hit recompile_limit 参考重编译解决方案排查
• 精度问题调试：遇到精度问题可调用 model-infer-precision-debug skill 进行排查

工作流程

第一步：方案设计

1. 分析模型结构：识别模型中可能阻碍图模式的代码模式
2. 识别问题点：找出 Graph Break、重编译、动态 shape 等问题
3. 设计改造方案：
   • 需要修改哪些文件
   • 修改的具体内容
   • 对现有功能的影响评估
4. 输出设计方案文档：以 Markdown 格式呈现

第二步：方案确认

等待用户确认方案后再进行开发。有修改意见则返回第一步。

第三步：实施开发

按照确认的设计方案逐步实施代码修改。

第四步：验证测试

由 Agent 实际执行以下测试，记录真实结果：

• 编译验证：运行 torch.compile，检查是否成功，记录编译日志
• 功能验证：运行模型，对比图模式前后输出
• 性能验证：记录 Prefill/Decode 阶段耗时
• 测试报告：整理测试环境、测试用例、对比数据

图模式选择

在开始适配前先确定使用哪种模式。如果用户未指定，默认采用 npugraph_ex。

npugraph_ex 快速示例

import torch
import torch_npu

model = YourModel().to("npu")
opt_model = torch.compile(model, backend="npugraph_ex", fullgraph=True, dynamic=False)
# 注：LLM Decode 场景 actual_seq_lengths 每步变化时需 dynamic=True
output = opt_model(input_tensor)

关键约束：PyTorch 2.6.0+、仅支持在线推理、不支持随机数算子和动态控制流、forward 中不可使用 .item()

GE 图模式快速示例

import torch
import torch_npu
import torchair
from torchair import patch_for_hcom

patch_for_hcom()  # 集合通信入图（有 TP/EP 并行时需调用）

model = YourModel().to("npu")
config = torchair.CompilerConfig()
npu_backend = torchair.get_npu_backend(compiler_config=config)
opt_model = torch.compile(model, backend=npu_backend)
output = opt_model(input_tensor)

详细文档：npugraph_ex 见 references/npugraph_ex-guide.md，GE 见 references/ge-graph-guide.md

LLM 模型适配要点

对于 LLM 推理模型，必须严格区分 prefill 和 decode 阶段。

Prefill 与 Decode 阶段限制

实现建议：

• 将 prefill 和 decode 的 forward 逻辑分离成不同方法
• 仅对 decode 方法应用 torch.compile 图模式
• prefill 阶段使用 eager 模式执行

class 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, backend="npugraph_ex", ...)

核心改造原则

核心原则：将动态变化的东西提取为模型输入，模型内部尽量保证静态。

详细改造指南：references/llm-model-guide.md

问题定界流程

问题定界优先基于本 skill 内置的知识进行独立分析。

问题发生
    │
    ├─→ aot_eager 验证 ──失败──→ 修复用户脚本
    │       ↓ 正常
    │
    ├─→ force_eager/run-eagerly ──失败──→ 修复用户脚本
    │       ↓ 正常
    │
    └─→ 图模式问题
            ├── 重编译问题 → 阅读 LLM 指南 + npugraph_ex 指南
            ├── Graph Break 问题 → 阅读 TorchAir 在线文档中的典型案例
            └── 其他问题 → 阅读对应模式文档（npugraph_ex-guide.md 或 ge-graph-guide.md）

调试知识来源（按优先级）：

1. 本 SKILL.md 中的方法、原则和检查清单
2. references/npugraph_ex-guide.md - npugraph_ex 详细指南
3. references/llm-model-guide.md - LLM 模型改造指南
4. references/ge-graph-guide.md - GE 图模式指南
5. TorchAir 官方文档

注意：避免盲目复制其他模型的图模式配置，应基于当前模型结构独立分析。

图模式 + FA 融合算子快速 Debug

图模式与 FA 融合算子结合时，actual_seq_lengths 参数的处理是最常见的出错点。 本节提供快速诊断和解决方案。

问题现象

图模式 + FA 场景下的典型问题：

• 编译报错：actual_seq_lengths 类型不匹配
• 运行时报错：重编译（recompile）触发
• 性能问题：动态 shape 导致无法充分优化

关键参数：actual_seq_lengths

FA 算子的 actual_seq_lengths / actual_seq_qlen / actual_seq_kvlen 参数在不同图模式下有不同的要求：

GE 模式配置指南

方案一：torchair FA 接口（推荐）

import torch
import torch_npu
import torchair
from torchair.ge_concrete_graph.ge_graph import mark_static

# 使用 torchair 提供的 FA 接口
# actual_seq_lengths 为 Tensor 类型
attn_output = torchair.ops.npu_fused_infer_attention_score(
    query, key, value,
    actual_seq_qlen=actual_seq_qlen_tensor,   # Tensor 类型
    actual_seq_kvlen=actual_seq_kvlen_tensor, # Tensor 类型
    # ... 其他参数
)

# 编译配置
opt_model = torch.compile(model, backend=npu_backend, dynamic=False)

优点：

• dynamic=False 可获得更好的静态图优化
• 无需额外的 mark_static 配置
• 图编译更稳定

约束：

• TorchAir FA 接口仅支持 GE 图模式，不支持 Eager 模式和 npugraph_ex 模式调用
• 必须在 GE 图模式（torchair.get_npu_backend()）下使用，不可在 Eager 或 npugraph_ex 模式下调用

方案二：torch_npu FA 接口（不推荐）

import torch
import torch_npu
import torchair
from torchair.ge_concrete_graph.ge_graph import mark_static

# 使用 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],
    # ... 其他参数
)

# 必须配置 dynamic=True
# 并使用 mark_static 标记除 actual_seq_lengths 外的静态输入
mark_static(input_ids)          # 静态输入
mark_static(position_ids)       # 静态输入
mark_static(attention_mask)     # 静态输入
# actual_seq_lengths 保持动态

# 编译配置
opt_model = torch.compile(model, backend=npu_backend, dynamic=True)

缺点：

• 需要配置 dynamic=True，性能略逊于静态图
• 需要手动调用 mark_static 标记所有静态输入
• 配置繁琐，易遗漏导致问题

npugraph_ex 模式配置指南

方案一：list[int] 类型 + dynamic=True

import 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],
    # ... 其他参数
)

# 必须配置 dynamic=True
opt_model = torch.compile(model, backend="npugraph_ex", dynamic=True)

方案二：Tensor 类型 + dynamic=False（如有接口支持）

import torch
import torch_npu

# 查询是否有 Tensor 类型的 actual_seq_lengths 接口
# 通过 subagent 调用 /model-infer-fusion 查询

# 如果有支持的接口：
attn_output = torch.ops.npu.npu_fused_infer_attention_score(
    query, key, value,
    actual_seq_lengths=actual_seq_lengths_tensor,  # Tensor 类型
    actual_seq_lengths_kv=actual_seq_kvlen_tensor,
    # ... 其他参数
)

# 可配置 dynamic=False
opt_model = torch.compile(model, backend="npugraph_ex", dynamic=False)

FA 接口查询

如需查询具体的 FA 接口参数和版本支持，使用 subagent 调用 /model-infer-fusion 技能：

启动 subagent（类型：general-purpose），提示词包含：
1. 明确指示调用 /model-infer-fusion 技能
2. 询问具体问题，例如：
   - "查询 npu_fused_infer_attention_score 的 actual_seq_lengths 参数是否支持 Tensor 类型"
   - "查询 npu_fused_infer_attention_score_v2 在图模式下的推荐配置"
   - "查询 torchair 提供的 FA 接口及其参数要求"
3. 提供当前使用的图模式（GE / npugraph_ex）

常见错误与修复

Debug 检查清单

在图模式 + FA 场景下，按以下清单逐一排查：

• 1. 确认使用的图模式：GE 还是 npugraph_ex
• 2. 确认 FA 接口来源：torch_npu 还是 torchair
• 3. 若使用 torchair FA 接口，确认当前为 GE 图模式（不支持 Eager 和 npugraph_ex）
• 4. 检查 actual_seq_lengths 类型：GE+torchair→Tensor；GE+torch_npu→list[int]+dynamic=True；npugraph_ex→list[int]+dynamic=True
• 5. 检查 dynamic 配置是否与 actual_seq_lengths 类型匹配
• 6. 若使用 GE + torch_npu FA + list[int]，检查是否已 mark_static 标记所有静态输入
• 7. 如有疑问，调用 model-infer-fusion skill 查询接口详情

TorchAir 官方文档索引

在线文档：gitcode.com/Ascend/torchair/docs/zh

入门必读

GE 图模式

自定义算子入图

常见案例和定位方法

性能分析与精度比对