用途:为昇腾 NPU 上的 PyTorch 模型添加 torch_npu.npu_prefetch 权重预取,优化推理性能,适用于 profiling 显示 MatMul/QBMM/GMM 存在 memory-bound 热点、需为模型添加权重预取或迁移 prefetch 模式的场景。前置条件:模型须在 NPU 上运行且已导入 torch_npu;必须经 profiling 确认存在 memory-bound 瓶颈,且前序算子不能是 memory-bound,否则会争抢带宽。
基本流程:先做方案设计,确定预取位置与 max_size,并输出 Markdown 设计方案;经用户确认后实施开发,贯通 enable_prefetch 开关、插入条件保护的 torch_npu.npu_prefetch 调用并配置 max_size 与 dependency;最后做编译、功能与性能验证,输出性能优化报告。目标算子优先级为 MatMul、BatchMatMul、QBMM、GMM 及大型 MLP/MoE 线性层;LayerNorm、RMSNorm、ROPE、SwiGLU 等通常作为依赖窗口而非预取目标。若所有前驱都是重搬运或重通信窗口,应说明无安全窗口并放弃预取。
关键原则:所有 prefetch 必须由 enable_prefetch 开关保护,默认 False;max_size 从保守值(理论值约 50%)开始,按 profiling 结果逐步调整;遵循单点验证,先验证单个目标算子再扩展;图模式下依赖节点须与目标算子在逻辑时序上匹配,不能跨分支复用。若没有合适前驱,不要硬插 prefetch。
使用方式:推荐通过 wrapper 函数调用。
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
max_size 计算按权重维度乘积乘数据类型字节数(bf16/fp16 为 2、int8 为 1、fp32 为 4)再乘保守系数确定。
限制:缺少模型权重时无法运行实际测试,只能提供测试脚本、报告模板与预期改善分析。具体性能提升幅度、收费、价格及支持平台等资料未提供。