mcgatom/codex-cannbot/model-infer-prefetch

Prefetch 预取优化技能

提供手动指定和自动化分析两种方式确定预取位置,覆盖方案设计、大小计算、代码实现和性能验证。

重要原则

• 前置条件:模型必须在 NPU 上运行且已导入 torch_npu
• 仅在 memory-bound 热点时使用:必须通过 profiling 确认存在 memory-bound 瓶颈
• 必须有安全的依赖窗口:前序算子不能是 memory-bound,否则会争抢带宽
• 始终提供开关保护:所有 prefetch 必须用 enable_prefetch 开关保护,默认 False
• 先保守后激进:max_size 从保守值开始,根据 profiling 结果逐步调整
• 单点验证原则:先验证单个目标算子,成功后再扩展到多个位置
• 图模式兼容性:在图模式下,依赖节点必须与目标算子在逻辑时序上匹配

工作流程

第一步:方案设计

1. 确定预取位置:选择手动指定或自动化分析方式

• 手动指定:用户提供目标算子和依赖窗口
• 自动化分析:调用 npu-roofline-analysis 从 profiling 数据中识别候选位置

1. 确定预取大小:根据权重张量维度和数据类型计算 max_size
2. 设计改造方案:

• 需要修改哪些文件
• 如何传递 enable_prefetch 开关
• 在哪里插入 torch_npu.npu_prefetch 调用
• max_size 和 dependency 的具体取值

1. 输出设计方案文档:以 Markdown 格式呈现

第二步:方案确认

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

第三步:实施开发

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

• 新增或贯通 enable_prefetch 开关
• 插入条件保护的 torch_npu.npu_prefetch 调用
• 配置 max_size 和 dependency 参数

第四步:验证测试

由 Agent 实际执行以下测试,记录真实结果:

• 编译验证:检查代码语法和运行时错误
• 功能验证:对比 prefetch 前后输出一致性
• 性能验证:运行 profiling,对比目标算子耗时、关键路径、依赖窗口
• 测试报告:整理前后对比数据,评估优化效果

重要:必须实际运行测试并生成性能优化报告,包括:

1. Baseline 性能数据 (enable_prefetch=False)
2. 优化后性能数据 (enable_prefetch=True)
3. 性能对比和改善百分比
4. 验证检查清单结果
5. 后续调优建议

预取位置确定

在开始实施前先确定使用哪种方式。如果用户未指定,优先采用自动化分析。

手动指定方式

用户直接提供:

• 目标算子:需要优化的 memory-bound 算子(如 MatMul, QBMM, GMM)
• 依赖窗口:目标算子的前序非 memory-bound 算子
• 权重信息:权重张量的维度和数据类型

直接跳转到"预取大小确定"章节。

自动化分析方式

通过调用 npu-roofline-analysis skill 和分析模型代码自动识别候选位置。

自动化分析流程

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

自动化分析决策表

前序算子访存密集判断标准

访存密集定义:指 memory-bound 算子,即内存带宽成为瓶颈的算子。

自动化分析示例

场景:LongCat-Flash 模型优化

1. 调用 npu-roofline-analysis:

• 识别 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, ...

1. 分析 gate_up_proj 前序算子:

• 前序:o_proj(大型 MatMul, memory-bound)
• 判断:访存密集 ✗
• 结论:但实际 o_proj 输出可作为依赖窗口(o_proj 执行完成后带宽空闲)
• 实际使用:在 o_proj 完成时预取 gate_up_proj 权重 ✓

1. 分析 down_proj 前序算子:

• 前序:SwiGLU(计算密集,非 memory-bound)
• 判断:非访存密集 ✓
• 结论:可预取,依赖窗口 = SwiGLU 输入 x
• 实际使用:在 gate_up_proj 输入时预取 down_proj 权重 ✓

1. 分析 q_a_proj 前序算子:

• 前序:RMSNorm(虽是访存类算子,但非 memory-bound)
• 判断:非访存密集 ✓
• 结论:可预取,依赖窗口 = down_proj 输出
• 实际使用:在 down_proj 完成时预取下一层 q_a_proj 权重 ✓

1. 分析 router.classifier 前序算子:

• 前序:router 输入(来自 recv 或 o_proj)
• 判断:非访存密集 ✓
• 结论:可预取,依赖窗口 = o_proj 输出
• 实际使用:在 o_proj 完成时预取 router.classifier 权重 ✓

1. 计算理论预取大小:

• gate_up_proj:hidden_size × intermediate_size × 2 × 2 字节(bf16)
• 保守取值:理论值的 50%

1. 输出候选列表:

• 候选 1:gate_up_proj, 依赖窗口 = o_proj 输出, max_size = 计算值
• 候选 2:down_proj, 依赖窗口 = gate_up_proj 输入, max_size = 计算值
• 候选 3:q_a_proj(下一层), 依赖窗口 = down_proj 输出, max_size = 18 MB
• 候选 4:router.classifier, 依赖窗口 = o_proj 输出, max_size = 18 MB
• ...

详细的 roofline 分析方法参见 npu-roofline-analysis skill

预取大小确定

根据权重张量维度和数据类型计算 max_size,遵循"先保守后激进"原则。

计算公式

基本公式:

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

数据类型字节数:

• bf16 / fp16: 2 字节
• int8: 1 字节
• fp32: 4 字节

保守系数:

• 首次尝试:0.5(理论值的 50%)
• 验证成功后:逐步提升到 0.7 → 1.0

不同算子的计算示例

MoE GMM 预取大小计算

# 数据类型字节数
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

调整策略

LongCat-Flash 实际取值参考

参考文件:cann-recipes-infer/models/longcat-flash/models/modeling_longcat_flash.py, cann-recipes-infer/models/longcat-flash/models/ffn.py

目标算子与依赖窗口选择

目标算子优先级

依赖窗口选择流程

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

依赖窗口选择标准

从目标算子向前找同一关键路径上的前序算子,优先选择:

• 离目标算子足够近(在同一 layer 或相邻 layer)
• 自身有一定执行时间(不是瞬时完成的小算子)
• 被 npu-roofline-analysis 判断为非 memory-bound
• 不是另一个大矩阵/大通信/大搬运算子

重要:即使前序算子本身是大型 MatMul(如 o_proj),其输出仍可作为依赖窗口,因为该算子执行完成后带宽空闲,可以开始预取下一个算子的权重。

LongCat-Flash 实际使用模式

关键观察:

1. 大型 MatMul 的输出可作为依赖窗口:o_proj, down_proj, gmm2_out 都是大型 MatMul,但其输出用于预取下一个算子
2. 跨层预取:在当前层的 down_proj 完成时预取下一层的 attention 权重
3. MoE 链式预取:router → GMM1 → GMM2 → 下一层 router,形成预取链

何时放弃预取

如果所有合理前驱都表现为重搬运或重通信窗口,则明确说明"当前阶段没有安全窗口",不要为了"有改动"而硬插 prefetch。

代码实现指南

1. 开关传递路径

检查开关传递链路:YAML → InferenceConfig → 模型类 __init__（通过 infer_config.model_config.custom_params 或专用字段）→ forward/decode/prefill

配置文件示例(config.yaml):

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

模型初始化示例:

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 调用模式

推荐:使用 wrapper 函数(cann-recipes-infer/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

在模型中调用:

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 必须是明确的 tensor 节点
• 依赖节点必须与目标算子在逻辑时序上匹配
• 不要跨分支、跨阶段乱复用依赖节点

正确示例:

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

错误示例:

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

4. 最小改动原则

优先:

• 单个目标算子
• 单个稳定窗口
• 单个保守 max_size
• 单文件或单模块的最小侵入式改动

避免:

• 多处同时预取
• 多个重量级目标一起上
• 全量权重一次性预取
• 改动多个并行策略和 prefetch 策略混在一个 patch 里

5. 仓库内的先验经验

性能测试验证

验证检查清单

在 prefetch 前后运行 profiling,按以下清单逐一检查:

• [ ] 目标算子执行时间是否下降:对比 prefetch 前后目标算子的耗时
• [ ] 关键路径是否缩短:对比主流上的总耗时或空洞/等待时间
• [ ] 依赖窗口是否稳定:前序算子耗时不应明显增加
• [ ] 是否出现新的性能瓶颈:检查其他流是否出现新的长拖尾
• [ ] 功能正确性是否保持:对比 prefetch 前后输出一致性

成功标准

至少满足其中两项:

1. 目标算子耗时下降
2. 关键路径缩短
3. 前序窗口没有明显退化
4. 功能结果无差异

失败信号与调整策略