CANN-ops-blas-昇腾NPU上MatMul的Tiling为什么决定一切
CANN-ops-blas-昇腾NPU上MatMul的Tiling为什么决定一切大模型推理 70% 的时间在跑矩阵乘法。ops-blas 仓库管的 GEMM通用矩阵乘法是所有 MatMul 的底层实现。在昇腾NPU上GEMM 的性能不取决于算术多快——取决于数据怎么搬进 Cube 单元。这就是 Tiling。Cube 单元的工作方式昇腾NPU的 Cube 单元一次能算 [16×16, 16×16] 的矩阵乘法fp16。一个 [4096, 4096] × [4096, 4096] 的 GEMM 需要 4096/16 256 个分块在 M 维、256 个分块在 N 维、256 个分块在 K 维。问题来了Cube 单元的片上缓存L1只有 1MB。一个 [16, 4096] 的分块在 fp16 下占 128KB。如果 M 和 N 维的分块各读一次、K 维分块反复读取L1 放不下。Tiling 就是解决怎么把大矩阵切小让 Cube 单元算得最快的问题。Tiling 策略的三种模式模式 1大块 TilingBig TileM_tile 128, N_tile 128, K_tile 128 每次读入 A[128, 128] B[128, 128] 64KB L1 放得下但 Cube 利用率只有 60%模式 2长条 TilingLong StripM_tile 16, N_tile 4096, K_tile 128 A 只有 4KBB 是 1MB B 需要反复读L1 放不下整个 B需要 double buffer Cube 利用率 85%但 HBM 读取次数多模式 3双缓冲 TilingDouble Buffer同时准备两组分块当前组在计算时下一组在 DMA 搬入 计算和搬运完全重叠 Cube 利用率 95%ops-blas 默认用模式 3。但 AOE自动调优引擎可以搜索更优的 tile size 组合。AOE 自动调优AOE 是 CANN 的自动调优工具遍历不同的 Tiling 参数组合找到最优配置# 对特定矩阵尺寸做调优aoe--job_type2--model_pathmodel.onnx--configaoe_config.json# 调优结果保存在 op_tiling 目录# 后续编译自动使用调优后的参数AOE 的搜索空间很大M_tile × N_tile × K_tile 的组合有上千种一次调优可能跑 2-4 小时。但调优后的 GEMM 性能可以比默认 Tiling 好 20-40%。什么场景需要 AOE固定 shape 的推理服务shape 不变调优一次永久生效。动态 shape 的推理服务不适合——每次 shape 变化都要重新调优。和 ops-transformer MergedMatMul 的关系MergedMatMul 把多个共享输入的 MatMul 合成一次 Batch GEMM。Batch GEMM 的 Tiling 跟普通 GEMM 不同——多了一个 batch 维度。ops-blas 对 Batch GEMM 的 Tiling 策略是先切 batch再切 M/N/K按 batch 维度分配到不同的 Cube 单元Atlas 800I A2 有多个 AI Core每个 AI Core 内部按 M/N/K 切分这意味着 batch 数最好是 AI Core 数量的整数倍。Atlas 800I A2 有 8 个 AI CoreQKV 三路 MergedMatMul 只有 3 个 batch不是 8 的倍数。ops-blas 会把 3 个 batch 分配到 3 个 AI Core剩下 5 个空闲。利用率只有 37.5%。这就是为什么 QKV MergedMatMul 在昇腾NPU上的加速比没有理论上那么高——3 个 batch 打不满 8 个 AI Core。性能数据Atlas 800I A2fp16 GEMM矩阵尺寸默认 TilingAOE 调优后理论峰值[4096, 4096] × [4096, 4096]210 TFLOPS280 TFLOPS310 TFLOPS[1, 4096] × [4096, 14336]85 TFLOPS95 TFLOPS310 TFLOPS小矩阵decode 阶段的 [1, 4096] × [4096, 14336]无论如何都打不满 Cube 单元——M 维只有 1算力浪费 70%。这是 decode 阶段 NPU 利用率低的根本原因Tiling 解决不了。GEMM 的 Tiling 是昇腾NPU性能调优的核心。默认 Tiling 覆盖了大部分场景但固定 shape 的推理服务用 AOE 调优可以额外提升 20-40%。小 batch decode 场景的 GEMM 利用率低是硬件限制不是软件问题。仓库在这里https://atomgit.com/cann/ops-blas