CANNBot技能：模型推理预取优化

【免费下载链接】cannbot-skillsCANNBot 是面向 CANN 开发的用于提升开发效率的系列智能体本仓库为其提供可复用的 Skills 模块。项目地址: https://gitcode.com/cann/cannbot-skillsname: model-infer-prefetch description: 基于 PyTorch 框架的昇腾 NPU 模型推理权重预取优化技能。为模型添加 torch_npu.npu_prefetch 特性以优化推理性能。触发场景profiling 显示 MatMul/QBMM/GMM 算子存在 memory-bound 热点、需要为模型添加权重预取、将 prefetch 模式迁移到新模型。 user-invocable: truePrefetch 预取优化技能提供手动指定和自动化分析两种方式确定预取位置,覆盖方案设计、大小计算、代码实现和性能验证。重要原则前置条件:模型必须在 NPU 上运行且已导入torch_npu仅在 memory-bound 热点时使用:必须通过 profiling 确认存在 memory-bound 瓶颈必须有安全的依赖窗口:前序算子不能是 memory-bound,否则会争抢带宽始终提供开关保护:所有 prefetch 必须用enable_prefetch开关保护,默认False先保守后激进:max_size从保守值开始,根据 profiling 结果逐步调整单点验证原则:先验证单个目标算子,成功后再扩展到多个位置图模式兼容性:在图模式下,依赖节点必须与目标算子在逻辑时序上匹配工作流程第一步:方案设计确定预取位置:选择手动指定或自动化分析方式手动指定:用户提供目标算子和依赖窗口自动化分析:调用npu-roofline-analysis从 profiling 数据中识别候选位置确定预取大小:根据权重张量维度和数据类型计算max_size设计改造方案:需要修改哪些文件如何传递enable_prefetch开关在哪里插入torch_npu.npu_prefetch调用max_size和dependency的具体取值输出设计方案文档:以 Markdown 格式呈现第二步:方案确认等待用户确认方案后再进行开发。有修改意见则返回第一步。第三步:实施开发按照确认的设计方案逐步实施代码修改:新增或贯通enable_prefetch开关插入条件保护的torch_npu.npu_prefetch调用配置max_size和dependency参数第四步:验证测试由 Agent 实际执行以下测试,记录真实结果:编译验证:检查代码语法和运行时错误功能验证:对比 prefetch 前后输出一致性性能验证:运行 profiling,对比目标算子耗时、关键路径、依赖窗口测试报告:整理前后对比数据,评估优化效果重要:必须实际运行测试并生成性能优化报告,包括:Baseline 性能数据 (enable_prefetchFalse)优化后性能数据 (enable_prefetchTrue)性能对比和改善百分比验证检查清单结果后续调优建议预取位置确定在开始实施前先确定使用哪种方式。如果用户未指定,优先采用自动化分析。方式输入要求输出结果适用场景手动指定位置用户提供目标算子和依赖窗口直接进入大小确定阶段已有明确优化目标,熟悉模型结构自动化分析profiling 数据或模型代码候选位置列表 理论预取大小需要从数据中发现机会,不确定优化点手动指定方式用户直接提供:目标算子:需要优化的 memory-bound 算子(如 MatMul, QBMM, GMM)依赖窗口:目标算子的前序非 memory-bound 算子权重信息:权重张量的维度和数据类型直接跳转到预取大小确定章节。自动化分析方式通过调用npu-roofline-analysisskill 和分析模型代码自动识别候选位置。自动化分析流程开始 │ ├─→ 步骤 1:调用 npu-roofline-analysis │ ↓ │ 获取 memory-bound top 10 算子列表 │ ↓ ├─→ 步骤 2:对每个 memory-bound 算子 │ ├── 分析前序算子 │ ├── 判断前序是否访存密集 │ │ │ ├─→ 前序是访存密集 ──→ 无预取空间 ──→ 跳过该算子 │ │ │ └─→ 前序非访存密集 ──→ 有预取空间 ──→ 标记为候选 │ ↓ ├─→ 步骤 3:计算理论预取大小 │ ├── 获取权重张量维度 │ ├── 获取数据类型(bf16/int8) │ └── 计算 max_size 维度 × 字节数 │ ↓ └─→ 输出:候选位置列表 理论预取大小自动化分析决策表步骤检查项结果下一步1调用npu-roofline-analysis获取 memory-bound top 10 算子进入步骤 22调用npu-roofline-analysis分析每个算子的前序算子前2-3个判断前序是否访存密集进入步骤 33a前序算子是访存密集(大 MatMul/通信)但其输出可用有预取空间(使用前序输出作为依赖窗口)计算理论预取大小3b前序算子非访存密集(LayerNorm/ROPE/SwiGLU/小算子)有预取空间(使用前序输出作为依赖窗口)计算理论预取大小3c前序算子是通信或无合适窗口无预取空间跳过该算子4计算权重大小理论 max_size输出候选列表前序算子访存密集判断标准访存密集定义:指 memory-bound 算子,即内存带宽成为瓶颈的算子。算子类型是否访存密集(memory-bound)说明大型 MatMul, BatchMatMul, QBMM, GMM是权重搬运密集,memory-bound,但其输出可作为依赖窗口通信算子(send, recv, all_reduce)是带宽密集,不适合作为依赖窗口大型数据搬运(大 transpose, 大 concat)是数据搬运量大,会争抢带宽LayerNorm, RMSNorm否虽是访存类算子,但不是 memory-bound,可作为依赖窗口ROPE, cast, reshape, 小 transpose否计算或搬运量小,可作为依赖窗口SwiGLU, GELU, Silu, npu_dequant_swiglu_quant否计算密集,可作为依赖窗口router topk, dispatch, combine否计算密集,可作为依赖窗口router.classifier(小型 MatMul)否权重较小,输出可作为依赖窗口自动化分析示例场景:LongCat-Flash 模型优化调用 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, ...分析 gate_up_proj 前序算子:前序:o_proj(大型 MatMul, memory-bound)判断:访存密集 ✗结论:但实际 o_proj 输出可作为依赖窗口(o_proj 执行完成后带宽空闲)实际使用:在 o_proj 完成时预取 gate_up_proj 权重 ✓分析 down_proj 前序算子:前序:SwiGLU(计算密集,非 memory-bound)判断:非访存密集 ✓结论:可预取,依赖窗口 SwiGLU 输入 x实际使用:在 gate_up_proj 输入时预取 down_proj 权重 ✓分析 q_a_proj 前序算子:前序:RMSNorm(虽是访存类算子,但非 memory-bound)判断:非访存密集 ✓结论:可预取,依赖窗口 down_proj 输出实际使用:在 down_proj 完成时预取下一层 q_a_proj 权重 ✓分析 router.classifier 前序算子:前序:router 输入(来自 recv 或 o_proj)判断:非访存密集 ✓结论:可预取,依赖窗口 o_proj 输出实际使用:在 o_proj 完成时预取 router.classifier 权重 ✓计算理论预取大小:gate_up_proj:hidden_size × intermediate_size × 2 × 2字节(bf16)保守取值:理论值的 50%输出候选列表:候选 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-analysisskill预取大小确定根据权重张量维度和数据类型计算max_size,遵循先保守后激进原则。计算公式基本公式:max_size 权重维度乘积 × 数据类型字节数 × 保守系数数据类型字节数:bf16/fp16: 2 字节int8: 1 字节fp32: 4 字节保守系数:首次尝试:0.5(理论值的 50%)验证成功后:逐步提升到 0.7 → 1.0不同算子的计算示例算子类型权重维度数据类型理论 max_size保守 max_size(50%)MatMul(q_proj)hidden_size × hidden_sizebf16H × H × 2H × HMatMul(o_proj)hidden_size × hidden_sizebf16H × H × 2H × HMatMul(gate_up_proj)hidden_size × intermediate_size × 2bf16H × I × 4H × I × 2MatMul(down_proj)intermediate_size × hidden_sizebf16I × H × 2I × HQBMMhidden_size × hidden_sizeint8H × HH × H / 2GMM(MoE)hidden_size × intermediate_size × 2 / moe_tp_size × experts_per_rankbf16/int8见下方见下方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调整策略场景调整方向说明目标算子加速明显,依赖窗口稳定增大 max_size从 50% → 70% → 100%目标算子加速不明显保持或减小 max_size可能预取收益有限依赖窗口明显退化减小 max_size从 50% → 30% → 10%依赖窗口严重退化调整位置或放弃前移 prefetch 或选择其他依赖窗口LongCat-Flash 实际取值参考位置目标算子依赖窗口max_size说明Line 1244router classifiero_proj18 MB固定值Line 1259q_a_projdense MLP output18 MB保守值Line 1260q_b_projdense MLP output36 MB中等值Line 1261kv_a_proj_with_mqadense MLP output7 MB保守值Line 1286next layer attentiondown_proj计算值动态计算参考文件:models/longcat-flash/models/modeling_longcat_flash.py,models/longcat-flash/models/ffn.py目标算子与依赖窗口选择目标算子优先级优先级算子类型说明高MatMul, BatchMatMul, QBMM, GMM主要预取目标,权重搬运密集高大型 MLP/FFN 线性层q_proj, k_proj, v_proj, o_proj, up_proj, down_proj, gate_up_proj高MoE 相关 MatMulrouter.classifier, experts.w13_weight, experts.w2_weight高融合前置大搬运算子如 MLAProlog(需 roofline 支持)低LayerNorm, RMSNorm, ROPE通常作为依赖窗口,不是预取目标低cast, reshape, transpose, concat同上低router topk, dispatch, combine同上低SwiGLU, GELU, Silu同上依赖窗口选择流程开始选择依赖窗口 │ ├─→ 找到目标算子的紧邻前驱 │ ↓ │ 检查:是否 memory-bound? │ ├─→ 是 ──→ 前移到更早的前驱 ──→ 重新检查 │ │ │ └─→ 否 ──→ 检查:是否有足够执行时间? │ ├─→ 是 ──→ 选择为依赖窗口 ✓ │ │ │ └─→ 否 ──→ 前移到更早的前驱 ──→ 重新检查 │ └─→ 所有前驱都不合适 ──→ 报告无安全窗口 ──→ 放弃该目标依赖窗口选择标准从目标算子向前找同一关键路径上的前序算子,优先选择:离目标算子足够近(在同一 layer 或相邻 layer)自身有一定执行时间(不是瞬时完成的小算子)被npu-roofline-analysis判断为非 memory-bound不是另一个大矩阵/大通信/大搬运算子重要:即使前序算子本身是大型 MatMul(如 o_proj),其输出仍可作为依赖窗口,因为该算子执行完成后带宽空闲,可以开始预取下一个算子的权重。LongCat-Flash 实际使用模式目标算子依赖窗口代码位置说明router.classifiero_proj 输出modeling_longcat_flash.py:1244o_proj 完成后预取 router 权重gate_up_projo_proj 输出modeling_longcat_flash.py:210o_proj 完成后预取 MLP 权重down_projgate_up_proj 输入 xmodeling_longcat_flash.py:239在 gate_up_proj 执行时预取 down_proj 权重q_a_proj(下一层)down_proj 输出modeling_longcat_flash.py:1259down_proj 完成后预取下一层 attention 权重q_b_proj(下一层)down_proj 输出modeling_longcat_flash.py:1261同上kv_a_proj(下一层)down_proj 输出modeling_longcat_flash.py:1263同上experts.w2_weightrouter_logits 输出ffn.py:76router 完成后预取 GMM2 权重router.classifier(下一层)gmm2_out 输出ffn.py:137GMM2 完成后预取下一层 router 权重experts.w13_weight(下一层)gmm2_out 输出ffn.py:139GMM2 完成后预取下一层 GMM1 权重关键观察:大型 MatMul 的输出可作为依赖窗口:o_proj, down_proj, gmm2_out 都是大型 MatMul,但其输出用于预取下一个算子跨层预取:在当前层的 down_proj 完成时预取下一层的 attention 权重MoE 链式预取:router → GMM1 → GMM2 → 下一层 router,形成预取链何时放弃预取如果所有合理前驱都表现为重搬运或重通信窗口,则明确说明当前阶段没有安全窗口,不要为了有改动而硬插 prefetch。代码实现指南1. 开关传递路径检查开关传递链路:命令行/配置文件 → ModelRunner → 模型类__init__→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) 02. Prefetch 调用模式推荐:使用 wrapper 函数(executor/utils/common_utils.py):def npu_prefetch(switch_flag, weight, depend, size, offset0): 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, weightself.next_layer.weight, dependdependency_output, sizeself.prefetch_size, offset0 ) # 执行目标算子 output self.next_layer(dependency_output) return output3. 图模式下的依赖节点图模式下,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. 仓库内的先验经验场景经验LongCat-Flash优先看 QBMM、MLAProlog、Matmul;对更重的后续算子,可提前预取更早权重AFD/FFN 分离Attention 提速后 FFN 成瓶颈时,优先看 GMM/Matmul;仅在通信间隙和非重搬运窗口存在时尝试MoE专家计算常见候选是 grouped_matmul 和大专家线性层;dispatch/combine/router 通常不是 prefetch 目标Dense MLPgate_up_proj/down_proj/o_proj 常是候选;但其前序窗口若已被另一个 MatMul 占满,不要硬加性能测试验证验证检查清单在 prefetch 前后运行 profiling,按以下清单逐一检查:目标算子执行时间是否下降:对比 prefetch 前后目标算子的耗时关键路径是否缩短:对比主流上的总耗时或空洞/等待时间依赖窗口是否稳定:前序算子耗时不应明显增加是否出现新的性能瓶颈:检查其他流是否出现新的长拖尾功能正确性是否保持:对比 prefetch 前后输出一致性成功标准至少满足其中两项:目标算子耗时下降关键路径缩短前序窗口没有明显退化功能结果无差异失败信号与调整策略失败信号根因分析调整策略目标算子没快,前序窗口反而变慢max_size过大,争抢带宽减小max_size:50% → 30% → 10%关键路径没有缩短,只是等待位置挪动依赖窗口选择不当前移或后移 prefetch 位置多流上出现新的长拖尾prefetch 引入新的同步点调整依赖节点或放弃该位置图模式下依赖关系不合法依赖节点与目标算子不匹配修正依赖节点或改为 eager 模式测试调优顺序先固定位置,只调max_size:从 50% → 70% → 100% 或 50% → 30% → 10%再固定max_size,微调 dependency 位置:前移或后移依赖窗口最后才考虑多目标叠加:单个目标验证成功后再扩展常见问题与修复错误现象根因修复方案不先跑 roofline 就选目标缺少 profiling 数据支持必须先用npu-roofline-analysis确认 memory-bound选了 LayerNorm/ROPE 作为预取目标误判目标算子这些通常是依赖窗口,不是预取目标dependency 窗口也是 memory-bound依赖窗口选择不当选非 memory-bound 的前驱,或前移 prefetch一次性预取全量权重max_size过大先保守max_size(50%),后续再调优多处同时加 prefetch改动范围过大先单个目标,验证后再扩展没有开关保护缺少回退机制所有 prefetch 必须有enable_prefetch开关前序窗口退化但不调整未根据 profiling 调整缩小max_size或调整位置,必要时回退图模式下编译失败依赖节点不合法确保依赖节点与目标算子在同一路径输出要求完成后必须汇报:选了哪些目标算子,为什么:基于 roofline 分析结果或用户手动指定roofline skill 的关键结论:memory-bound 算子列表、关键路径(如使用自动化分析)为什么选这个 dependency 窗口:非 memory-bound、有足够执行时间加了哪些开关与代码路径:配置文件 → 模型初始化 → forward 调用max_size取值依据:权重维度、数据类型、保守系数或用户指定值做了哪些验证:profiling 对比、功能验证还有哪些风险或后续调优建议:潜在问题、下一步优化方向必须交付的文件:修改后的代码文件:包含 prefetch 实现的模型代码测试脚本:用于对比 baseline 和优化后性能的自动化测试脚本性能优化报告:Markdown 格式,包含:优化配置和参数代码修改说明性能测试结果(如有模型权重)验证检查清单后续调优建议如何运行测试的说明如果无法运行实际测试(如缺少模型权重):创建完整的测试脚本和报告模板说明如何运行测试提供预期的性能改善分析参考资料索引主题文件路径npu_prefetch wrapper 函数executor/utils/common_utils.pyLongCat-Flash prefetch 实现models/longcat-flash/models/modeling_longcat_flash.pyFFN prefetch 大小计算models/longcat-flash/models/ffn.py配置验证逻辑models/longcat-flash/models/model_setting.py配置示例models/longcat-flash/config/README.mdRoofline 分析调用npu-roofline-analysisskill图模式适配调用model-infer-graph-modeskill【免费下载链接】cannbot-skillsCANNBot 是面向 CANN 开发的用于提升开发效率的系列智能体本仓库为其提供可复用的 Skills 模块。项目地址: https://gitcode.com/cann/cannbot-skills创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考