external-cannbot-model-model-infer-fusion

torch_npu fused operator optimization skill

Analyze computation patterns in PyTorch model code and match them with torch_npu fused operators for replacement optimization. Based on the torch_npu local docstring query script (scripts/torch_npu_query.py) and the fused operator usage experience of existing models in the repository, analyze, match, replace, and verify module by module.

Workflow

Step 1: Analyze the model code and break it down into modules

Analyze the model code and output a module list.

Break down the network structure: Identify the top-level structure (Embedding → Transformer Blocks → LM Head), and break down each Block into independent modules:

Attention layer: Norm → QKV projection → RoPE → KV Cache → Flash Attention → O projection (MLA includes V absorb)
MoE / FFN layer: Norm → Gate → routing dispatch → expert computation → aggregation
Other modules: Embedding, LM Head, cross-layer residual, etc.

Record runtime scenario characteristics：

Prefill / Decode branch differences: the same module may follow different operator paths in the two phases, and must be handled separately during replacement
Quantization requirements: BF16 / W8A8 / W8A8C8 / W4A16
Distributed configuration: TP / EP / DP, affecting MoE routing operator selection
Fixed constraints: layout (NZ/BSND/TND), cache format, metadata

Expand key sub-chains: For complex modules such as Attention and MoE, expand to the replaceable chain level:

Attention: RoPE, KV Cache write/read, Attention Core (FA / FA v2 / Sparse FA)
MoE: Gate / TopK, Routing Init / Dispatch, Expert computation, Finalize / Combine, Shared Expert (if any)
Other key chains (Residual + Norm, QKV Projection, O Projection, MC2 / AllToAll, etc.) are also included

Step 2: Independently match each module against repository reference implementations

For each module broken down in Step 1,independentlymatch the closest operator chain in the repository reference implementations. Once a match is found, use that path as acandidate blueprint, and supplement the analysis by comparing it with the key sub-chain list from Step 1 for parts not covered by that path; for unmatched modules, skip to Step 3 and search in the operator master table.

Attention layer

Note: If the FA fused operator has already been replaced in a previous stage, skip only the FA call itself. The other sub-chains of Attention (RoPE, KV Cache write, Residual+Norm, QKV/O Projection, etc.) still need to be checked one by one for available fused operators. Also note multi-step overall fused operators (such as npu_kv_rmsnorm_rope_cache fused RMSNorm+RoPE+Cache write,npu_mla_prolog_v3 fused Q/KV projection+RMSNorm+RoPE+Cache write); do not simply classify a single submodule as "standard linear layer with no fused operator" and skip evaluation.

When analyzing Attention,KV Cache needs to be discussed together with Attention Core ; it is not recommended to separate them completely; in addition to the architecture, the following should also be confirmed:

cache organization method: contiguous non-PA / PA
write index semantics:kv_len / start_pos, or slot_mapping / cache_index
subsequent Attention Core: non-fused attention,npu_fused_infer_attention_score、npu_fused_infer_attention_score_v2、npu_sparse_flash_attention

Common implementation combinations include:

Continuous non-PA + continuous write (e.g., scatter_update_ + kv_len / start_pos) → then evaluate non-fused attention or fused FA
PA + slot-based write (e.g., npu_scatter_nd_update_ / npu_scatter_pa_kv_cache / cache_index) → then evaluate the FA
MLA absorb specialized path:npu_mla_prolog_v3(cache_index / slot_mapping) + fused FA

First, based on the current implementation's architecture, cache form, write method, and Attention Core organization, determine which reference path it is closer to, then proceed to the corresponding detail document.

Attention 架构？
  │
  ├─ GQA / MHA（标准多头 / 分组查询注意力）
  │    │
  │    ├─ Prefill: 当前 batch 的 q/k/v → FA(sparse_mode=3 推荐，sparse_mode=2 不推荐)
  │    └─ Decode:  从 KV Cache 读取 → FA(actual_seq_lengths_kv)
  │    → 详情：references/module-attention-gqa.md
  │
  └─ MLA（Multi-head Latent Attention，低秩 KV 压缩）
       │
       ├─ 无 Indexer
       │    ├─ Prefill: 分步投影 → 展开 K/V → FA v1 或 v2
       │    └─ Decode:  absorb（手动或 npu_mla_prolog_v3）→ FA v1 或 v2 → V absorb
       │    → 详情：references/module-attention-mla-absorb.md
       │
       └─ 有 Indexer（稀疏 Top-K KV 选择）
            ├─ Prefill/Decode 共路径
            │  npu_mla_prolog_v3 → Indexer → 稀疏 FA → V absorb
            → 详情：references/module-attention-mla-indexer.md

MoE / FFN Layer

First confirm whether the current module belongs to Dense FFN or MoE; if it is MoE, then combine the gate form, parallel mode, and stage to determine which reference path it is closer to.

是否有 MoE？
  │
  ├─ 无（Dense FFN）→ 在算子总表中确认 Dense Linear / Activation / Norm 可用融合算子
  │
  └─ 有 MoE
       │
       ├─ Gate 算子
       │    ├─ softmax 打分 → npu_moe_gating_top_k_softmax（qwen3-moe）
       │    └─ sigmoid/noaux → npu_moe_gating_top_k（deepseek 系列）
       │
       └─ 路由 + 专家计算（按并行模式和阶段区分）
            ├─ Prefill（纯 TP）: init_routing_v2 → grouped_matmul → finalize_routing
            ├─ Prefill（EP）:    init_routing_v2 → AllToAll → re_routing → grouped_matmul → finalize_routing
            └─ Decode（EP+TP）:  MC2 dispatch_v2 → grouped_matmul → MC2 combine_v2
       → 详情：references/moe-fusion-reference.md

Unmatched Module

For other modules not covered by the above decision tree (Embedding, LM Head, cross-layer residual, Diffusion-specific modules, etc.), skip to Step 3 and search in the operator summary table.

Step 3: Consult the operator interface documentation to confirm availability and adaptability

For operators hit in Step 2 or modules that do not match any pattern, consult the torch_npu interface documentation in the following priority order:

Priority: local docstring

Built into the torch_npu wheel _op_plugin_docs.py, containing complete Chinese documentation for all operators. Query via scripts/torch_npu_query.py :

# 单 API 详情
python3 scripts/torch_npu_query.py show <api_name>

# 反向搜索
python3 scripts/torch_npu_query.py search "<keyword>"

# 枚举所有算子名
python3 scripts/torch_npu_query.py list [--prefix npu_]

Confirm availability：

For operators hit by a pattern: consult the detail documentation to confirm parameter constraints and applicable scenarios
For modules that do not hit any pattern: search in the operator summary table, and read the detail documentation to analyze functionality

Adaptation verification：

Check whether shape, dtype, layout, cache organization, and metadata meet the operator requirements
If the differences can be resolved through reasonable upfront modifications (such as format conversion, RoPE precomputation and value-fetch path reorganization, KV Cache staticization/PA modification, etc.), it should be marked as "candidate + requires upfront modification", and the required modification items should be stated; for some modifications with relatively complex processes, you may ask the user whether to adopt them
Only when the differences are hard constraints and cannot be resolved through reasonable upfront modifications may it be marked as "not adaptable". When marking, the specific hard constraint must be noted (such as operator error information, parameter limits explicitly stated in the documentation), and it must not be marked as not adaptable merely on the grounds of "requiring relatively large changes"

Step 4: Analysis-stage review

Before proceeding to code replacement, review the following items. Incomplete items must be returned to the corresponding step to be completed, and must not be skipped to directly proceed to implementation.

If the current task only requires analysis, then end at this step, output the analysis results, candidate solutions, and verification plan, and do not proceed to code implementation.

Analysis-stage review items：

Module decomposition is complete: it has been expanded to the replaceable link level (including Prefill/Decode branch differences), and the key Attention/MoE sub-links have been covered
Reference matching is complete: repository reference implementations have been matched by module (or it has been confirmed that there is no match), and sub-links not covered by the reference paths have been supplemented for analysis
Operator constraints have been confirmed: reference / API documentation has been consulted to confirm the adaptation constraints and quantization compatibility of candidate operators
Candidate list has been formed: each candidate module clearly specifies prerequisites, the minimal verification切入点, and blocking points

Step 5: Implement replacement module by module

After a comprehensive operator candidate analysis is available, iterate through and replace all modules in the candidate list one by one according to the replacement process, changing only one place at a time, and proceeding to the next only after verification passes; no module that has entered the candidate list may be skipped. If the current module cannot be further implemented, its failure evidence, blocking reasons, and current conclusion must still be recorded.

If the candidate module depends on the following prerequisite transformations and they have not yet been completed, the corresponding resources may be consulted as needed:

RoPE Precomputation, caching, and value retrieval path organization:references/rotary-embedding-pattern.md
KV Cache Staticization / PA transformation: refer to model-infer-kvcache skill

Fused operator replacement process：

Complete prerequisite transformations (if needed)
Replace the operator code of this module
Verification: accuracy comparison (output alignment before and after fusion)
Verification: performance comparison (confirm there is a gain)
If verification passes → keep, continue to the next module
If verification fails → roll back the changes, re-analyze this module:
- Are there other available operators or alternative solutions? → try the alternative solution
- Confirm that there is currently no suitable operator but there is a fusion gain → record as a new requirement (module location, computation semantics, reason for unsuitability, expected fusion form)
- Currently unable to continue implementation → after rolling back, record the failure evidence and blocking reasons, continue to the next one
Record the optimization report for this module: accuracy comparison results, performance comparison results, logs or error paths

Note: During replacement, refer to the closest model implementation in the repository and the references/ operator interface documentation usage instructions.

Implementation phase checklist：

All modules in the candidate list have been processed one by one (replaced or skipped reason recorded), with no omissions
Each implemented module has recorded accuracy and performance comparison results
Modules that failed implementation have recorded error messages, attempted handling methods, and final conclusions

Reference documentation index

The following documents are consulted as needed to avoid loading them all at once and consuming tokens

Topic	Path
torch_npu API documentation query script	`scripts/torch_npu_query.py`
Attention: GQA reference chain	`references/module-attention-gqa.md`
Attention: MLA Absorb reference chain	`references/module-attention-mla-absorb.md`
Attention: MLA+Indexer reference chain	`references/module-attention-mla-indexer.md`
MoE operator pattern details	`references/moe-fusion-reference.md`
RotaryEmbedding precomputation and invocation pattern	`references/rotary-embedding-pattern.md`
op-plugin repository (online)	gitcode.com/Ascend/op-plugin