external-cannbot-model-model-infer-graph-mode

Graph Mode Adaptation and Optimization Skill

Provides adaptation guides for npugraph_ex and GE graph modes, covering solution design, graph break fixes, FA parameter configuration, and validation testing.

---

Important Principles

• Prerequisites: The model must run on NPU and must have imported torch_npu
• Graph mode applies only to the Decode phase: In the Prefill phase, the input length changes dynamically, so it is not suitable for graph mode; keep eager execution
• Preserve model integrity: Do not simplify model logic just to adapt to graph mode
• Backends not supported by NPU：aot_eager、inductor、cudagraphs etc. are unavailable
• Pre-create fixed tensors outside the graph: Tensors that remain unchanged throughout inference, such as atten_mask and KV cache, should be pre-created outside the graph, using torch._dynamo.mark_static() markers
• Notes on LLM Decode recompilation：kv_len/actual_seq_lengths_kv changes at each step, and npugraph_ex may trigger recompilation; when encountering hit recompile_limit refer to the recompilation solution for troubleshooting
• Precision issue debugging: When encountering precision issues, you can invoke the model-infer-precision-debug skill for troubleshooting

---

Workflow

Step 1: Solution Design

1. Analyze the model structure: Identify code patterns in the model that may hinder graph mode
2. Identify problem points: Find issues such as Graph Breaks, recompilation, and dynamic shapes
3. Design the modification plan：
   • Which files need to be modified
   • The specific content of the modifications
   • Assessment of the impact on existing functionality
4. Output the design document: Present it in Markdown format

Step 2: Solution Confirmation

Wait for the user to confirm the solution before proceeding with development. If there are revision comments, return to Step 1.

Step 3: Implementation and Development

Implement the code modifications step by step according to the confirmed design plan.

Step 4: Verification Testing

The Agent actually executes the following tests and records the real results:

• Compilation verification: Run torch.compile, check whether it succeeds, and record the compilation logs
• Functional verification: Run the model and compare the outputs before and after graph mode
• Performance verification: Record the time consumed in the Prefill/Decode phases
• Test report: Organize the test environment, test cases, and comparison data

---

Graph Mode Selection

Before starting adaptation, first determine which mode to use. If the user does not specify, use npugraph_ex by default.

Scenario	Recommended Mode	Detailed Documentation
LLM Large Language Model	npugraph_ex (read the LLM guide first)	references/llm-model-guide.md
General models	npugraph_ex	references/npugraph_ex-guide.md
Requires GE graph mode	GE（Ascend IR）	references/ge-graph-guide.md

Feature	npugraph_ex backend (aclgraph)	GE graph mode (Ascend IR)
Enablement method	backend="npugraph_ex"	torchair.get_npu_backend()
Implementation principle	Capture & Replay	FX graph converted to Ascend IR, compiled and executed by GE engine
Maturity	Experimental feature, not yet commercially supported	More mature and stable
PyTorch version	Requires 2.6.0+	No special requirements
Supported scenarios	Online inference	General scenarios
Similar technology	torch.cuda.CUDAGraph	Traditional graph compilation

npugraph_ex quick example

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)

Key constraints: PyTorch 2.6.0+, only supports online inference, does not support random number operators and dynamic control flow, cannot be used in forward .item()

GE graph mode quick example

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)

Detailed documentation: npugraph_ex see references/npugraph_ex-guide.md, GE see references/ge-graph-guide.md

---

LLM model adaptation key points

For LLM inference models,the prefill and decode phases must be strictly distinguished。

Prefill and Decode phase restrictions

Phase	Whether graph mode is supported	Reason
Prefill	Prohibited	Input length changes dynamically, first token generation logic is complex, shape is not fixed
Decode	Recommended	Input length is fixed (usually 1), shape is stable, suitable for graph capture

Implementation suggestions：

• Separate the forward logic of prefill and decode into different methods
• Apply torch.compile graph mode only to the decode method
• The prefill phase uses eager mode execution

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", ...)

Core refactoring principles

Core principle: Extract dynamically changing things as model inputs, and keep the model internals as static as possible.

Dynamic factor	Problem manifestation	Solution approach
Memory address change	Guard failure, recompilation	Pre-allocate a fixed size, update in place
Shape change	Graph break, multiple compilations	Fix the shape or control it through parameters
Python control flow	Graph Break	Use Tensor operations or pattern parameters
.item() call	Forces Graph Break	Keep Tensor or pass in externally

Detailed refactoring guide:references/llm-model-guide.md

---

Problem scoping process

Problem scoping should prioritize independent analysis based on the built-in knowledge of this skill.

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

Debugging knowledge sources (by priority)：

1. Methods, principles, and checklists in this SKILL.md
2. references/npugraph_ex-guide.md - npugraph_ex detailed guide
3. references/llm-model-guide.md - LLM model refactoring guide
4. references/ge-graph-guide.md - GE graph mode guide
5. TorchAir official documentation

Note: Avoid blindly copying the graph mode configurations of other models; analyze independently based on the current model structure.

---

Quick Debug for graph mode + FA fused operator

When graph mode is combined with the FA fused operator,actual_seq_lengths parameter handling is the most common source of errors. This section provides quick diagnosis and solutions.

Problem symptoms

Typical problems in graph mode + FA scenarios:

• Compilation error:actual_seq_lengths type mismatch
• Runtime error: recompile triggered
• Performance issue: dynamic shape prevents full optimization

Key parameter: actual_seq_lengths

The actual_seq_lengths / actual_seq_qlen / actual_seq_kvlen parameter of the FA operator has different requirements under different graph modes:

Graph mode	FA interface source	actual_seq_lengths type	dynamic setting	Execution mode constraint	Description
GE mode (recommended)	torchair FA interface	Tensor	dynamic=False	Only supports GE graph mode	Best solution, static graph
GE mode (not recommended)	torch_npu FA interface	list[int]	dynamic=True + mark_static	No restriction	Requires extra configuration, error-prone
npugraph_ex mode	torch_npu FA interface	list[int]	dynamic=True	No restriction	Dynamic capture mode
npugraph_ex mode	torch_npu FA interface	Tensor (if available)	dynamic=False	No restriction	Need to confirm whether the interface supports it

GE mode configuration guide

Solution 1: torchair FA interface (recommended)

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)

Advantages：

• dynamic=False Can achieve better static graph optimization
• No need for extra mark_static configuration
• More stable graph compilation

Constraints：

• The TorchAir FA interfaceonly supports GE graph mode, and does not support calls in Eager mode or npugraph_ex mode
• It must be used in GE graph mode (torchair.get_npu_backend()), and cannot be called in Eager or npugraph_ex mode

Solution 2: torch_npu FA interface (not recommended)

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)

Disadvantages：

• Requires configuration of dynamic=True, performance slightly inferior to static graph
• Need to manually call mark_static to mark all static inputs
• Configuration is cumbersome and prone to omissions that cause problems

npugraph_ex mode configuration guide

Solution 1: list[int] type + 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)

Solution 2: Tensor type + dynamic=False (if the interface supports it)

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 interface query

To query specific FA interface parameters and version support, use a subagent to invoke the /model-infer-fusion skill:

启动 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）

Common errors and fixes

Error symptom	Root cause	Fix
Compilation error: actual_seq_lengths type error	Tensor passed to torch_npu FA interface in GE mode	Switch to the torchair FA interface, or switch to list[int] + dynamic=True
Frequent recompilation at runtime	dynamic=False but actual_seq_lengths is list[int]	Switch to Tensor type + torchair interface, or set dynamic=True
Performance below expectations	dynamic=True prevents full optimization	Use the torchair FA interface + dynamic=False as much as possible
npugraph_ex mode error	actual_seq_lengths is a Tensor but the interface does not support it	Confirm the interface support situation, or switch to list[int] + dynamic=True
Calling TorchAir FA in Eager or npugraph_ex mode reports an error	The TorchAir FA API only supports GE graph mode	Switch to the torch_npu FA API, or switch to GE graph mode

Debug checklist

In the graph mode + FA scenario, troubleshoot one by one according to the following checklist:

• 1. Confirm the graph mode in use: GE or npugraph_ex
• 2. Confirm the source of the FA API: torch_npu or torchair
• 3. If the torchair FA API is used, confirm that the current mode is GE graph mode (Eager and npugraph_ex are not supported)
• 4. Check the actual_seq_lengths type: GE+torchair→Tensor; GE+torch_npu→list[int]+dynamic=True; npugraph_ex→list[int]+dynamic=True
• 5. Check whether the dynamic configuration matches the actual_seq_lengths type
• 6. If GE + torch_npu FA + list[int] is used, check whether mark_static has been called to mark all static inputs
• 7. If in doubt, call the model-infer-fusion skill to query API details

---

TorchAir official documentation index

Online documentation:gitcode.com/Ascend/torchair/docs/zh

Must-read for beginners

Topic	Online link
Introduction	overview.md
List of supported ATen APIs	aten_api.md

GE graph mode

Topic	Online link
GE graph mode overview	ascend_ir.md
GE graph mode quick start	quick_start.md
GE graph mode features (basic)	basic/
GE graph mode features (advanced)	advanced/
GE graph mode API reference	api/

Custom operator graph capture

Topic	Online link
Overview	overview.md
In-place operator development and graph capture example	in_place_op_cases.md
Non-in-place operator development and graph capture example	non_in_place_op_cases.md
Operator plugin adaptation	op_plugin_adapt_torchair.md

Common cases and troubleshooting methods

Topic	Online link
Large model graph mode inference case	infer_cases.md
Dynamic/static graph concepts and typical problem troubleshooting	dynamic_static_graph/
Graph capture failure delimitation and troubleshooting	graph_failed_cases.md

Performance analysis and accuracy comparison

Topic	Online link
Graph mode performance analysis	perfermance_cases.md
Graph mode accuracy comparison	accuracy_cases.md
FAQ	faq.md