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：大块 Tiling（Big Tile）

M_tile = 128, N_tile = 128, K_tile = 128
每次读入 A[128, 128] + B[128, 128] = 64KB
L1 放得下，但 Cube 利用率只有 60%

模式 2：长条 Tiling（Long Strip）

M_tile = 16, N_tile = 4096, K_tile = 128
A 只有 4KB，B 是 1MB
B 需要反复读，L1 放不下整个 B，需要 double buffer
Cube 利用率 85%，但 HBM 读取次数多

模式 3：双缓冲 Tiling（Double Buffer）

同时准备两组分块：当前组在计算时，下一组在 DMA 搬入
计算和搬运完全重叠
Cube 利用率 95%+

ops-blas 默认用模式 3。但 AOE（自动调优引擎）可以搜索更优的 tile size 组合。

AOE 自动调优

AOE 是 CANN 的自动调优工具，遍历不同的 Tiling 参数组合，找到最优配置：

# 对特定矩阵尺寸做调优
aoe --job_type=2 --model_path=model.onnx --config=aoe_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"：

1. 按 batch 维度分配到不同的 Cube 单元（Atlas 800I A2 有多个 AI Core）
2. 每个 AI Core 内部按 M/N/K 切分

这意味着 batch 数最好是 AI Core 数量的整数倍。Atlas 800I A2 有 8 个 AI Core，QKV 三路 MergedMatMul 只有 3 个 batch，不是 8 的倍数。ops-blas 会把 3 个 batch 分配到 3 个 AI Core，剩下 5 个空闲。利用率只有 37.5%。

这就是为什么 QKV MergedMatMul 在昇腾NPU上的加速比没有理论上那么高——3 个 batch 打不满 8 个 AI Core。

性能数据

Atlas 800I A2，fp16 GEMM：

矩阵尺寸	默认 Tiling	AOE 调优后	理论峰值
[4096, 4096] × [4096, 4096]	210 TFLOPS	280 TFLOPS	310 TFLOPS
[1, 4096] × [4096, 14336]	85 TFLOPS	95 TFLOPS	310 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