external-cannbot-model-model-infer-prefetch

Prefetch optimization skill

Provides two methods, manual specification and automated analysis, to determine prefetch locations, covering solution design, size calculation, code implementation, and performance verification.

---

Important Principles

• Prerequisites:The model must run on the NPU and must have imported torch_npu
• Use only for memory-bound hotspots:Profiling must confirm the existence of a memory-bound bottleneck
• Must have a safe dependency window:The preceding operator must not be memory-bound, otherwise it will contend for bandwidth
• Always provide switch protection:All prefetch must be protected by the enable_prefetch switch, default False
• Conservative first, then aggressive:max_size Start from a conservative value and gradually adjust based on profiling results
• Single-point verification principle:First verify a single target operator, and only after success expand to multiple locations
• Graph mode compatibility:In graph mode, the dependency node must match the target operator in logical timing

---

Workflow

Step 1: Solution Design

1. Determine prefetch locations:Choose manual specification or automated analysis
   • Manual specification: the user provides the target operator and dependency window
   • Automated analysis: call npu-roofline-analysis to identify candidate locations from profiling data
2. Determine prefetch size:Calculate max_size
3. Design the modification plan:
   • Which files need to be modified
   • How to pass the enable_prefetch switch
   • Where to insert the torch_npu.npu_prefetch call
   • max_size and dependency specific values
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 modification suggestions, return to Step 1.

Step 3: Implementation and Development

Implement the code modifications step by step according to the confirmed design solution:

• Add or connect through the enable_prefetch switch
• Insert the conditionally protected torch_npu.npu_prefetch call
• Configure the max_size and dependency parameters

Step 4: Verification Testing

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

• Compilation verification:Check code syntax and runtime errors
• Functional verification:Compare output consistency before and after prefetch
• Performance verification:Run profiling and compare the target operator latency, critical path, and dependency window
• Test report:Organize before-and-after comparison data and evaluate the optimization effect

Important:Tests must actually be run and a performance optimization report must be generated, including:

1. Baseline performance data (enable_prefetch=False)
2. Optimized performance data (enable_prefetch=True)
3. Performance comparison and improvement percentage
4. Verify Checklist Results
5. Follow-up Tuning Suggestions

---

Prefetch Position Determination

Determine which method to use before starting implementation. If the user does not specify, prefer automated analysis.

Method	Input Requirements	Output Result	Applicable Scenarios
Manually Specify Position	User provides target operator and dependency window	Proceed directly to the size determination stage	Already have a clear optimization target and are familiar with the model structure
Automated Analysis	profiling data or model code	Candidate position list + theoretical prefetch size	Need to discover opportunities from data and are unsure of the optimization points

Manual Specification Method

The user directly provides:

• Target Operator: the memory-bound operator to optimize (such as MatMul, QBMM, GMM)
• Dependency Window: the preceding non-memory-bound operator of the target operator
• Weight Information: the dimensions and data type of the weight tensor

Jump directly to the "Prefetch Size Determination" section.

Automated Analysis Method

By calling the npu-roofline-analysis skill and analyzing the model code to automatically identify candidate positions.

Automated Analysis Process

开始
  │
  ├─→ 步骤 1:调用 npu-roofline-analysis
  │       ↓
  │   获取 memory-bound top 10 算子列表
  │       ↓
  ├─→ 步骤 2:对每个 memory-bound 算子
  │       ├── 分析前序算子
  │       ├── 判断前序是否访存密集
  │       │
  │       ├─→ 前序是访存密集 ──→ 无预取空间 ──→ 跳过该算子
  │       │
  │       └─→ 前序非访存密集 ──→ 有预取空间 ──→ 标记为候选
  │               ↓
  ├─→ 步骤 3:计算理论预取大小
  │       ├── 获取权重张量维度
  │       ├── 获取数据类型(bf16/int8)
  │       └── 计算 max_size = 维度 × 字节数
  │           ↓
  └─→ 输出:候选位置列表 + 理论预取大小

Automated Analysis Decision Table

Step	Check Item	Result	Next Step
1	Call npu-roofline-analysis	Get the top 10 memory-bound operators	Proceed to step 2
2	Call npu-roofline-analysisAnalyze each operator's preceding operators (the previous 2-3)	Determine whether the preceding operator is memory-intensive	Proceed to step 3
3a	The preceding operator is memory-intensive (large MatMul/communication) but its output is usable	There is prefetch space (use the preceding output as the dependency window)	Calculate the theoretical prefetch size
3b	The preceding operator is not memory-intensive (LayerNorm/ROPE/SwiGLU/small operators)	There is prefetch space (use the preceding output as the dependency window)	Calculate the theoretical prefetch size
3c	The preceding operator is communication or there is no suitable window	No prefetch space	Skip this operator
4	Calculate the weight size	Theoretical max_size	Output the candidate list

Criteria for Determining Whether the Preceding Operator Is Memory-Intensive

Definition of Memory-Intensive: refers to memory-bound operators, that is, operators where memory bandwidth becomes the bottleneck.

Operator Type	Whether Memory-Intensive (memory-bound)	Description
Large MatMul, BatchMatMul, QBMM, GMM	Yes	Weight movement is intensive, memory-bound, but itsoutput can be used as a dependency window
Communication operators (send, recv, all_reduce)	Yes	Bandwidth-intensive, not suitable as a dependency window
Large data movement (large transpose, large concat)	Yes	Large data movement volume, will compete for bandwidth
LayerNorm, RMSNorm	No	Although it is a memory access operator, it is not memory-bound,can serve as a dependency window
ROPE, cast, reshape, small transpose	No	Computation or transfer volume is small,can serve as a dependency window
SwiGLU, GELU, Silu, npu_dequant_swiglu_quant	No	Compute-intensive,can serve as a dependency window
router topk, dispatch, combine	No	Compute-intensive,can serve as a dependency window
router.classifier(small MatMul)	No	Weights are relatively small,the output can serve as a dependency window

Automated analysis example

Scenario:LongCat-Flash model optimization

1. Call npu-roofline-analysis:

   • Identify memory-bound top 10:gate_up_proj(MatMul), down_proj(MatMul), q_a_proj(MatMul), q_b_proj(MatMul), kv_a_proj_with_mqa(MatMul), router.classifier(MatMul), GMM1, GMM2, ...

2. Analyze the preceding operators of gate_up_proj:

   • Preceding:o_proj(large MatMul, memory-bound)
   • Judgment:memory access intensive ✗
   • Conclusion:but in practice the o_proj output can serve as a dependency window (bandwidth is idle after o_proj finishes execution)
   • Actual use:prefetch gate_up_proj weights when o_proj completes ✓

3. Analyze the preceding operators of down_proj:

   • Preceding:SwiGLU(compute-intensive, not memory-bound)
   • Judgment:not memory access intensive ✓
   • Conclusion:can prefetch, dependency window = SwiGLU input x
   • Actual use:prefetch down_proj weights when gate_up_proj is input ✓

4. Analyze the preceding operators of q_a_proj:

   • Preceding:RMSNorm(although it is a memory access operator, it is not memory-bound)
   • Judgment:not memory access intensive ✓
   • Conclusion:can prefetch, dependency window = down_proj output
   • Actual use:prefetch the next layer's q_a_proj weights when down_proj completes ✓

5. Analyze the preceding operators of router.classifier:

   • Preceding:router input(from recv or o_proj)
   • Judgment:not memory access intensive ✓
   • Conclusion:can prefetch, dependency window = o_proj output
   • Actual use:prefetch router.classifier weights when o_proj completes ✓

6. Calculate the theoretical prefetch size:

   • gate_up_proj:hidden_size × intermediate_size × 2 × 2 bytes (bf16)
   • Conservative value:50% of the theoretical value

7. Output the candidate list:

   • Candidate 1:gate_up_proj, dependency window = o_proj output, max_size = calculated value
   • Candidate 2:down_proj, dependency window = gate_up_proj input, max_size = calculated value
   • Candidate 3:q_a_proj(next layer), dependency window = down_proj output, max_size = 18 MB
   • Candidate 4:router.classifier, dependency window = o_proj output, max_size = 18 MB
   • ...

For the detailed roofline analysis method, see npu-roofline-analysis skill

---

Determining the prefetch size

Calculate based on the weight tensor dimensions and data type max_size, following the principle of "be conservative first, then aggressive".

Calculation formula

Basic formula:

max_size = 权重维度乘积 × 数据类型字节数 × 保守系数

Data type byte count:

• bf16 / fp16: 2 bytes
• int8: 1 byte
• fp32: 4 bytes

Conservative coefficient:

• First attempt: 0.5 (50% of theoretical value)
• After successful verification: gradually increase to 0.7 → 1.0

Calculation examples for different operators

Operator type	Weight dimension	Data type	Theoretical max_size	Conservative max_size (50%)
MatMul(q_proj)	hidden_size × hidden_size	bf16	H × H × 2	H × H
MatMul(o_proj)	hidden_size × hidden_size	bf16	H × H × 2	H × H
MatMul(gate_up_proj)	hidden_size × intermediate_size × 2	bf16	H × I × 4	H × I × 2
MatMul(down_proj)	intermediate_size × hidden_size	bf16	I × H × 2	I × H
QBMM	hidden_size × hidden_size	int8	H × H	H × H / 2
GMM(MoE)	hidden_size × intermediate_size × 2 / moe_tp_size × experts_per_rank	bf16/int8	See below	See below

MoE GMM prefetch size calculation

# 数据类型字节数
dtype_bit = 1 if quant_mode == "w8a8" else 2  # int8: 1 字节, bf16: 2 字节

# GMM1(gate_up_proj)预取大小
gmm1_prefetch_size = hidden_size * intermediate_size * 2 * dtype_bit // moe_tp_size * experts_per_rank // 2

# GMM2(down_proj)预取大小
gmm2_prefetch_size = hidden_size * intermediate_size * dtype_bit // moe_tp_size * experts_per_rank

Adjustment strategy

Scenario	Adjustment direction	Description
Target operator acceleration is significant, dependency window is stable	Increase max_size	From 50% → 70% → 100%
Target operator acceleration is not significant	Keep or decrease max_size	Prefetch benefit may be limited
Dependency window degrades significantly	Decrease max_size	From 50% → 30% → 10%
Dependency window degrades severely	Adjust position or abandon	Move prefetch earlier or choose a different dependency window

LongCat-Flash actual value reference

Position	Target operator	Dependency window	max_size	Description
Line 1244	router classifier	o_proj	18 MB	Fixed value
Line 1259	q_a_proj	dense MLP output	18 MB	Conservative value
Line 1260	q_b_proj	dense MLP output	36 MB	Moderate value
Line 1261	kv_a_proj_with_mqa	dense MLP output	7 MB	Conservative value
Line 1286	next layer attention	down_proj	Calculated value	Dynamically calculated

Reference files:models/longcat-flash/models/modeling_longcat_flash.py, models/longcat-flash/models/ffn.py

---

Target operator and dependency window selection

Target operator priority

Priority	Operator type	Description
High	MatMul, BatchMatMul, QBMM, GMM	Main prefetch target, weight transfer intensive
High	Large MLP/FFN linear layer	q_proj, k_proj, v_proj, o_proj, up_proj, down_proj, gate_up_proj
High	MoE-related MatMul	router.classifier, experts.w13_weight, experts.w2_weight
High	Fused preceding large transfer operator	Such as MLAProlog (requires roofline support)
Low	LayerNorm, RMSNorm, ROPE	Usually serves as a dependency window, not a prefetch target
Low	cast, reshape, transpose, concat	Same as above
Low	router topk, dispatch, combine	Same as above
Low	SwiGLU, GELU, Silu	Same as above

Dependency window selection process

开始选择依赖窗口
  │
  ├─→ 找到目标算子的紧邻前驱
  │       ↓
  │   检查:是否 memory-bound?
  │       ├─→ 是 ──→ 前移到更早的前驱 ──→ 重新检查
  │       │
  │       └─→ 否 ──→ 检查:是否有足够执行时间?
  │               ├─→ 是 ──→ 选择为依赖窗口 ✓
  │               │
  │               └─→ 否 ──→ 前移到更早的前驱 ──→ 重新检查
  │
  └─→ 所有前驱都不合适 ──→ 报告"无安全窗口" ──→ 放弃该目标

Dependency window selection criteria

From the target operator, search backward for preceding operators on the same critical path, prioritizing:

• Close enough to the target operator (in the same layer or adjacent layer)
• Has some execution time of its own (not an instantaneous small operator)
• Judged by npu-roofline-analysis asnon memory-bound
• Is not another large matrix/large communication/large transfer operator

Important: Even if the preceding operator itself is a large MatMul (such as o_proj), itsoutputcan still serve as a dependency window, because after that operator finishes executing, bandwidth is idle and prefetching of the next operator's weights can begin.

LongCat-Flash actual usage pattern

Target operator	Dependency window	Code location	Description
router.classifier	o_proj output	modeling_longcat_flash.py:1244	Prefetch router weights after o_proj completes
gate_up_proj	o_proj output	modeling_longcat_flash.py:210	Prefetch MLP weights after o_proj completes
down_proj	gate_up_proj input x	modeling_longcat_flash.py:239	Prefetch down_proj weights while gate_up_proj executes
q_a_proj (next layer)	down_proj output	modeling_longcat_flash.py:1259	Prefetch next layer attention weights after down_proj completes
q_b_proj (next layer)	down_proj output	modeling_longcat_flash.py:1261	Same as above
kv_a_proj (next layer)	down_proj output	modeling_longcat_flash.py:1263	Same as above
experts.w2_weight	router_logits output	ffn.py:76	Prefetch GMM2 weights after router completes
router.classifier (next layer)	gmm2_out output	ffn.py:137	Prefetch next layer router weights after GMM2 completes
experts.w13_weight (next layer)	gmm2_out output	ffn.py:139	Prefetch next layer GMM1 weights after GMM2 completes

Key Observations:

1. Outputs of large MatMuls can serve as dependency windows: o_proj, down_proj, and gmm2_out are all large MatMuls, but their outputs are used to prefetch the next operator
2. Cross-layer prefetching: prefetch the next layer's attention weights when the current layer's down_proj completes
3. MoE chain prefetching: router → GMM1 → GMM2 → next layer router, forming a prefetch chain

When to abandon prefetching

If all reasonable predecessors exhibit heavy-transfer or heavy-communication windows, explicitly state that "there is no safe window at the current stage"; do not force-insert prefetch just to "have a change".

---

Code Implementation Guide

1. Switch propagation path

Check the switch propagation chain: command line/config file → ModelRunner → model class __init__ → forward/decode/prefill

Config file example(config.yaml):

model_config:
  enable_prefetch: true  # 默认 false
  enable_multi_stream: 2  # prefetch 需要多流支持

Model initialization example:

class YourModel:
    def __init__(self, config):
        self.enable_prefetch = config.get("enable_prefetch", False)
        # prefetch 需要图模式 + 多流
        assert not self.enable_prefetch or config.get("enable_multi_stream", 0) > 0

2. Prefetch call pattern

Recommended: use a wrapper function(executor/utils/common_utils.py):

def npu_prefetch(switch_flag, weight, depend, size, offset=0):
    if switch_flag:
        return torch_npu.npu_prefetch(weight, depend, size, offset)
    else:
        return None

Call in the model:

from executor.utils.common_utils import npu_prefetch

class YourModel:
    def forward(self, x):
        # 计算依赖窗口
        dependency_output = self.some_layer(x)

        # 预取下一个算子的权重
        npu_prefetch(
            self.enable_prefetch,
            weight=self.next_layer.weight,
            depend=dependency_output,
            size=self.prefetch_size,
            offset=0
        )

        # 执行目标算子
        output = self.next_layer(dependency_output)
        return output

3. Dependency nodes in graph mode

• In graph mode,dependency must be an explicit tensor node
• Dependency nodes must match the target operator in logical timing
• Do not arbitrarily reuse dependency nodes across branches or stages

Correct example:

# 依赖节点是目标算子的直接前驱
attn_output = self.attention(x)  # 依赖窗口
npu_prefetch(self.enable_prefetch, self.mlp.weight, attn_output, size)
mlp_output = self.mlp(attn_output)  # 目标算子

Incorrect example:

# 依赖节点与目标算子不在同一路径
attn_output = self.attention(x)
npu_prefetch(self.enable_prefetch, self.mlp.weight, x, size)  # 错误:x 不是 mlp 的直接前驱
mlp_output = self.mlp(attn_output)

4. Principle of minimal changes

Prefer:

• a single target operator
• a single stable window
• a single conservative max_size
• minimally invasive changes to a single file or module

Avoid:

• prefetching in multiple places simultaneously
• deploying multiple heavyweight targets together
• one-time prefetching of all weights
• mixing changes to multiple parallel strategies and prefetch strategies in one patch

5. Prior experience in the repository

Scenario	Experience
LongCat-Flash	Prioritize QBMM, MLAProlog, and Matmul; for heavier subsequent operators, earlier weights can be prefetched in advance
AFD/FFN separation	When FFN becomes the bottleneck after Attention speedup, prioritize GMM/Matmul; only attempt when communication gaps and non-heavy-transfer windows exist
MoE	Expert computation common candidates are grouped_matmul and large expert linear layers; dispatch/combine/router are usually not prefetch targets
Dense MLP	gate_up_proj/down_proj/o_proj are often candidates; but if their preceding window is already filled by another MatMul, do not force adding

---

Performance test verification

Verification checklist

Run profiling before and after prefetch, and check one by one according to the following checklist:

• Whether the target operator execution time decreases:Compare the duration of the target operator before and after prefetch
• Whether the critical path is shortened:Compare the total duration or gaps/waiting time on the main stream
• Whether the dependency window is stable:The duration of preceding operators should not increase significantly
• Whether new performance bottlenecks appear:Check whether other streams show new long tails
• Whether functional correctness is maintained:Compare output consistency before and after prefetch

Success criteria

At least satisfytwo:

1. Target operator duration decreases
2. Critical path is shortened
3. Preceding window does not degrade significantly
4. No difference in functional results

Failure signals and adjustment strategies

Failure signal	Root cause analysis	Adjustment strategy
The target operator does not get faster, and the preceding window instead slows down	max_size too large, contending for bandwidth	Reduce max_size:50% → 30% → 10%
The critical path is not shortened, only the waiting position shifts	Improper dependency window selection	Move the prefetch position earlier or later
New long tails appear on multiple streams	Prefetch introduces new synchronization points	Adjust the dependency node or abandon this position
Dependency relationships are invalid in graph mode	Dependency node does not match the target operator	Fix the dependency node or switch to eager mode for testing

Tuning order

1. First fix the position, only tune max_size:from 50% → 70% → 100% or 50% → 30% → 10%
2. Then fix max_size, and fine-tune the dependency position:Move the dependency window earlier or later
3. Only finally consider stacking multiple targets:Expand after a single target is verified successfully

---

Common issues and fixes

Error symptom	Root cause	Fix
Choosing targets without running roofline first	Lack of profiling data support	Must first use npu-roofline-analysis to confirm memory-bound
Selected LayerNorm/ROPE as prefetch targets	Misjudged target operators	These are usually dependency windows, not prefetch targets
The dependency window is also memory-bound	Improper dependency window selection	Choose a predecessor that is not memory-bound, or move prefetch earlier
Prefetch all weights at once	max_size too large	Start conservatively with max_size(50%), and tune later
Add prefetch in multiple places at the same time	Scope of changes is too large	Start with a single target, then expand after verification
No switch protection	Lack of fallback mechanism	All prefetch must have enable_prefetch switch
Preceding window degrades but is not adjusted	Not adjusted based on profiling	Reduce max_size or adjust position, fall back when necessary
Compilation fails in graph mode	Dependency node is invalid	Ensure dependency nodes and target operators are on the same path

---

Output requirements

Upon completion, must report:

1. Which target operators were selected, and why:Based on roofline analysis results or manually specified by the user
2. Key conclusions of the roofline skill:Memory-bound operator list, critical path (if automated analysis is used)
3. Why this dependency window was chosen:Not memory-bound, with sufficient execution time
4. Which switches and code paths were added:Configuration file → model initialization → forward call
5. max_size basis for the value:Weight dimensions, data type, conservative coefficient, or user-specified value
6. What validation was performed:Profiling comparison, functional validation
7. What risks or follow-up tuning suggestions remain:Potential issues, next optimization directions

Files that must be delivered:

1. Modified code files:Model code containing the prefetch implementation
2. Test scripts:Automated test scripts used to compare baseline and optimized performance
3. Performance optimization report:Markdown format, including:
   • Optimization configuration and parameters
   • Code modification description
   • Performance test results (if model weights are available)
   • Validation checklist
   • Follow-up tuning suggestions
   • Instructions on how to run the tests

If actual tests cannot be run(such as when model weights are missing):

• Create complete test scripts and report templates
• Explain how to run the tests
• Provide expected performance improvement analysis

---

Reference index

Topic	File path
npu_prefetch wrapper function	executor/utils/common_utils.py
LongCat-Flash prefetch implementation	models/longcat-flash/models/modeling_longcat_flash.py
FFN prefetch size calculation	models/longcat-flash/models/ffn.py
Configuration validation logic	models/longcat-flash/models/model_setting.py
Configuration example	models/longcat-flash/config/README.md
Roofline analysis	Call npu-roofline-analysis skill
Graph mode adaptation	Call model-infer-graph-mode skill