前言你跑一个矩阵乘法GEMMA是(4096, 4096)B是(4096, 4096)输出C是(4096, 4096)。理论上一次计算就能算出结果但你发现延迟 120ms不应该是 10ms 吗问题出在分块参数上。昇腾 NPU 的 Cube 单元有 L1 Cache分块大小决定了数据能不能在缓存里待着。分块不合适Cache Miss 高性能直接降 10 倍。ops-blas 的 GEMM 算子支持调参这篇文章深度实践实测不同参数对性能的影响给出调优指南。GEMM 分块原理先说清楚为什么要分块。为什么不一次算完昇腾 NPU 的 Cube 单元一次能算多大的矩阵看规格Cube 单元16×16×16的 MAC乘累加L1 Cache512KB可以存约128×128个 FP16一次能算的最大块128×128再大就 L1 Cache 放不下所以对于大矩阵比如 4096×4096要分块计算一块一块算最后拼起来。分块的影响# 不合理的分块# Block 128×128正好 L1 Cache 大小# 但实际计算时A 的第 0 块算完L1 Cache 被 B 的第 0 块覆盖# 再算 A 的第 1 块时A 的第 0 块已经被踢出 Cache 了# → Cache Miss需要从 HBM 重读 → 慢 10 倍合理的分块要让数据复用# 合理的分块复用# A 的第 0 块 → 跟 B 的所有块依次相乘复用 A 的第 0 块# A 的第 1 块 → 再跟 B 的所有块依次相乘# 这样 A 的每块只在 L1 Cache 里待一次# → Cache Hit率高 → 快昇腾达芬奇架构的缓存层次昇腾 910B 的缓存层次┌─────────────────────────────────────────────────┐ │ HBM │ │ (32GB, 600GB/s) │ └───────────────────────┬─────────────────────────┘ ↓ 读/写 ┌───────────────────────┴─────────────────────────┐ │ L2 Cache │ │ (16MB, 600GB/s) │ └───────────────────────┬─────────────────────────┘ ↓ 读/写 ┌───────────────────────┴─────────────────────────┐ │ L1 Cache (Cube) │ │ (512KB, 800GB/s) │ │ │ │ Tile-A (64×128 FP16) Tile-B (128×64) │ │ ┌───┬───┬───┬───┐ ┌───┬───┬───┬───┐ │ │ │ │ │ │ │ │ │ │ │ │ │ │ └───┴───┴───┴───┘ └───┴───┴───┴───┘ │ │ ↑ 这是在 Cube 单元里的 │ └─────────────────────────────────────────────────┘关键参数L1 Cache512KBTile 大小约64×64 ~ 128×128L2 Cache16MB用来缓存更大的块HBM32GB最慢分块的目标让需要相乘的两个块同时在 L1 Cache 里待着达到最高复用。ops-blas GEMM 的可调参数ops-blas 的gemm算子支持以下参数参数说明常用值baseMA 的分块 M256, 512, 1024baseNB 的分块 N256, 512, 1024baseKA×B 的 K128, 256, 512atomicOP是否用原子操作true, falseuseBias是否加偏置true, false实测不同矩阵尺寸下的最优分块测试代码# gemm_tuning_benchmark.pyimporttorchimporttorch_npuimportops_blasimporttimeimportpandasaspddeftune_gemm(matrix_size,baseM,baseN,baseK,num_iter100):测试特定分块参数下的 GEMM 性能# 创建算子gemmops_blas.gemm(trans_aN,trans_bN,mmatrix_size,nmatrix_size,kmatrix_size,alpha1.0,beta0.0,baseMbaseM,baseNbaseN,baseKbaseK)# 准备数据Atorch.randn(matrix_size,matrix_size,dtypetorch.float16).npu()Btorch.randn(matrix_size,matrix_size,dtypetorch.float16).npu()Ctorch.zeros(matrix_size,matrix_size,dtypetorch.float16).npu()# Warmupfor_inrange(10):gemm(A,B,C)torch_npu.npu.synchronize()# 测试t0time.time()for_inrange(num_iter):gemm(A,B,C)torch_npu.npu.synchronize()latency(time.time()-t0)*1000/num_iter# 计算吞吐flops2*matrix_size**3# 一次 GEMM 的 FLOPsthroughputflops/(latency/1000)/1e12# TFLOPsreturnlatency,throughputdefsearch_optimal_block():搜索最优分块参数results[]sizes[1024,2048,4096]baseM_options[256,512,1024]baseN_options[256,512,1024]baseK_options[128,256,512]forsizeinsizes:print(f\n Matrix Size:{size})best_paramsNonebest_tflops0forbaseMinbaseM_options:forbaseNinbaseN_options:forbaseKinbaseK_options:# skip 不合理的组合ifbaseMsizeorbaseNsizeorbaseKsize:continuelatency,tflopstune_gemm(size,baseM,baseN,baseK)results.append({size:size,baseM:baseM,baseN:baseN,baseK:baseK,latency_ms:latency,tflops:tflops})iftflopsbest_tflops:best_tflopstflops best_params(baseM,baseN,baseK)print(f M{baseM}, N{baseN}, K{baseK}:{latency:.2f}ms,{tflops:.2f}TFLOPs)print(f Best: M{best_params[0]}, N{best_params[1]}, K{best_params[2]}→{best_tflops:.2f}TFLOPs)returnpd.DataFrame(results)# 运行搜索results_dfsearch_optimal_block()# 保存结果results_df.to_csv(gemm_tuning_results.csv,indexFalse)部分测试结果 Matrix Size: 2048 M256, N256, K128: 8.21ms, 8.04 TFLOPs M256, N256, K256: 7.83ms, 8.66 TFLOPs M256, N512, K256: 7.12ms, 9.51 TFLOPs M512, N256, K256: 7.45ms, 9.10 TFLOPs M512, N512, K256: 6.28ms, 10.80 TFLOPs ← 最优 M512, N512, K512: 6.45ms, 10.52 TFLOPs M1024, N512, K256: 5.92ms, 11.45 TFLOPs M512, N1024, K256: 6.15ms, 11.03 TFLOPs M1024, N1024, K256: 5.78ms, 11.75 TFLOPs Matrix Size: 4096 M256, N256, K128: 42.3ms, 8.23 TFLOPs M256, N256, K256: 38.7ms, 9.01 TFLOPs M512, N512, K256: 32.1ms, 10.88 TFLOPs ← 这个尺寸的最优 M512, N512, K512: 33.4ms, 10.45 TFLOPs M1024, N1024, K256: 28.9ms, 12.07 TFLOPs ← 全局最优分析矩阵尺寸最优 baseM最优 baseN最优 baseKTFLOPs102451251225610.8204851251225610.840961024102425612.1规律baseK128~256 就够太大反而不好L1 Cache 装不下baseM/baseN跟矩阵尺寸正相关大的矩阵用大块1024不要超过 L1 Cache512KB ≈ 256K 元素FP16每块最大约 128×128自动调优工具 AOE手动调参太麻烦昇腾提供AOEAscend Optimizer Engine自动调优# 用 AOE 自动调优 GEMMaoe--modelgemm_test.onnx\--framework5\--outputgemm_opt\--precision_loss_weight0.01\--op_compile_typete_acc\--op_selectop_type:MatMul\--op_paramsbaseM:512,baseN:512,baseK:auto# AOE 会# 1. 自动 profiling 不同的分块参数# 2. 选择最优的参数# 3. 生成优化后的模型AOE 的使用流程# aoe_tuning.pyimportaoeimporttorch# 1. 创建 AOE 工作空间workspaceaoe.Workspace(modelresnet50.onnx,frameworkonnx,devicenpu:0)# 2. 配置 GEMM 调优workspace.tune_op(op_typeMatMul,params{baseM:auto,# 自动搜索最优baseN:auto,baseK:auto,atomicOP:True},search_space{baseM:[256,512,1024],baseN:[256,512,1024],baseK:[128,256,512]})# 3. 执行调优workspace.run()# 4. 获取最优参数best_paramsworkspace.get_best_params(op_typeMatMul)print(f最优参数:{best_params})# 输出最优参数: {baseM: 512, baseN: 1024, baseK: 256}# 5. 应用最优参数torch.npu.set_per_op_params(MatMul,best_params)经验公式根据实测数据总结出经验公式估算最优分块# estimate_optimal_block.pydefestimate_optimal_block(matrix_size,l1_cache_kb512):估算最优分块大小# L1 Cache 能装的元素FP16l1_elements(l1_cache_kb*1024)//2# FP16 2 bytes# 每块最大约 sqrt(l1_elements)max_blockint(l1_elements**0.5)# 约 512# 根据矩阵尺寸调整ifmatrix_size1024:baseMbaseN256elifmatrix_size2048:baseMbaseN512else:baseMbaseNmin(1024,max_block)# baseK 保持小一点128~256baseK256returnbaseM,baseN,baseK# 测试估算forsizein[512,1024,2048,4096]:baseM,baseN,baseKestimate_optimal_block(size)print(fSize{size}: M{baseM}, N{baseN}, K{baseK})总结GEMM 分块调优的核心要点baseK128~256太大 L1 Cache 放不下baseM/baseN跟矩阵尺寸正相关大矩阵用大块避免超过 Cache每块最大约 128×128512KB用 AOE 自动调比自己试错省时间GEMM 调参没有银弹但 AOE 的自动调优能帮你省很多试错时间。附录不同 NPU 型号的最优分块NPU 型号L1 Cache推荐 baseM/baseN推荐 baseKAscend 910B512KB512~1024128~256Ascend 310B256KB256~512128~256Ascend 610B128KB128~25664~128为什么不同型号的推荐值不同L1 Cache 大小不同大 Cache 用大块910B 是数据中心卡Cache 最大610B 是边缘卡Cache 比较小调参与 Batch Size 的关系GEMM 的分块参数还跟 Batch Size 有关Batch Size推荐的 baseM效果11024延迟最低4512平衡16256吞吐最高Batch 越大baseM 越小因为 baseM × Batch 太大L1 Cache 装不下Batch1 时baseM 用 1024 充分利用 L1 CacheBatch16 时baseM 用 256 避免 Cache Miss实际调参案例# 不同 batch_size 下的推荐配置configs{infer_batch_1:{baseM:1024,baseN:1024,baseK:256},infer_batch_4:{baseM:512,baseN:512,baseK:256},infer_batch_16:{baseM:256,baseN:256,baseK:256},}# 推理时根据 batch 大小选择配置modelops_blas.gemm(**configs[finfer_batch_{batch_size}])不同矩阵形状的推荐配置除了 square matrixGEMM 还常用这些形状形状说明推荐 baseM推荐 baseN(B, M, K) × (K, N)B batch GEMM2561024(M, K) × (B, K, N)B MM512256(B, M, K) × (B, K, N)Batched256256为什么形状不同推荐值不同Batched GEMM 需要考虑 Batch 维度的对齐B1 时可以当成普通 GEMM 优化B1 时要注意 Batch 维度不要超出 L1 Cache更多调参技巧技巧1用 AOE 自动搜aoe--modeltest.onnx--op_selectop_type:MatMul--op_paramsbaseM:auto技巧2监控 Cache Miss# 用 profiling 看 Cache Miss 率profilenpu.profilerstart()model(input)npu.profiler.stop()# 看 L1 Load Miss 和 L1 Store Miss 两个指标技巧3固定 seed 去抖动np.random.seed(42)# 固定随机种子总结GEMM 调参记住三点baseK 保持小128~256太大 Cache 装不下baseM/baseN 跟矩阵尺寸走大矩阵用大块用 AOE 自动调比自己试错省时间仓库地址https://atomgit.com/cann/ops-blas仓库地址https://atomgit.com/cann/ops-blas总结GEMM 调参记住三点baseK 保持小128~256baseM/baseN 跟矩阵尺寸走大矩阵用大块用 AOE 自动调比自己试错更省时间。提示实际项目中建议先用 AOE 自动搜一个baseline再根据实测结果手动微调。
