ops-blas：昇腾NPU上线性代数算子的性能天花板在哪？
前言
去年帮一个量化基金优化因子计算，每天要跑10^12次矩阵乘（GEMM），原来用cuBLAS在A100上跑，一天算不完。后来迁到ops-blas + Ascend 910，同样的计算量18小时跑完，省了30%的硬件成本。这个结果让我重新审视了线性代数算子的性能天花板——不是换个硬件就能解决的，关键在于你能不能把硬件的计算单元喂饱。
这篇文章不是ops-blas的API文档，是我实际项目中对GEMM优化的理解，以及ops-blas在Ascend 910上跑出理论峰值92%利用率的技术细节。
GEMM：深度学习的地基
先说一个很多人忽略的事实：深度学习90%的计算量都是GEMM。
- MatMul？就是GEMM（C = A × B）
- Conv2D？im2col展开后也是GEMM
- Attention（Q×K^T、softmax×V）？还是GEMM
- 全连接层？GEMM
- LSTM/GRU的隐藏状态计算？GEMM
所以GEMM的性能直接决定了模型训练和推理的速度。你调优模型，本质上是调优GEMM。
Ascend 910的Matrix单元（Cube Core）理论峰值：256 TOPS（FP16）。这个数字很漂亮，但实际能跑出多少？ops-blas给出了答案：236 TOPS，理论峰值利用率92%。
ops-blas的GEMM优化三板斧
92%的利用率不是白来的。ops-blas的GEMM实现用了三个核心优化策略：Tiling、双缓冲、L0 Cache优化。
第一斧：Tiling——把大矩阵切成小块
Ascend 910的Cube Core（矩阵计算单元）不能直接算一个4096×4096的矩阵乘——它的Local Memory（L1 Cache）只有192KB per Cube Core，装不下这么大的矩阵。
Tiling的做法：把大矩阵切成小块（tile），每次算一个小块，算完累加。
原始GEMM：C[M,N] = A[M,K] × B[K,N] Tiling后： for m in range(0, M, TILE_M): # 沿M轴切 for n in range(0, N, TILE_N): # 沿N轴切 for k in range(0, K, TILE_K): # 沿K轴切 C[m:m+TILE_M, n:n+TILE_N] += A[m:m+TILE_M, k:k+TILE_K] × B[k:k+TILE_K, n:n+TILE_K]tile大小怎么选？这是关键。ops-blas的自动tiling算法根据矩阵大小和Local Memory容量，自动选择最优的TILE_M、TILE_N、TILE_K。
# ops-blas GEMM调用示例（自动tiling）importtorchfromops_blasimportgemm A=torch.randn(4096,4096,dtype=torch.float16).npu()B=torch.randn(4096,4096,dtype=torch.float16).npu()# ops-blas自动选择最优tiling参数C=gemm(A,B)# 等价于 torch.matmul(A, B)，但走ops-blas的优化路径手动配置tiling参数（高级用法）：
fromops_blasimportgemm,TilingConfig# 手动指定tiling参数（比如你知道矩阵特征的场景）tiling=TilingConfig(tile_m=128,tile_n=128,tile_k=64)C=gemm(A,B,tiling=tiling)第二斧：双缓冲——计算和搬运并行
Tiling解决了"装不下"的问题，但引入了新问题：每次算完一个tile，要从HBM搬运下一个tile的数据到Local Memory，搬运期间Cube Core闲置。
双缓冲的做法：准备两块Local Memory（Buffer A和Buffer B），Cube Core算Buffer A的数据时，DMA控制器同步搬运下一个tile到Buffer B。算完后交换A和B，Cube Core算Buffer B，DMA搬运下一个tile到Buffer A。
时间线： Cube Core: [算tile 1][算tile 2][算tile 3][算tile 4] DMA搬运: [搬tile 2][搬tile 3][搬tile 4][搬tile 5] ↑ tile 1在初始化时已经搬好了结果：Cube Core几乎不等待数据搬运，利用率从60%提升到90%。
第三斧：L0 Cache优化——最大化数据复用
Cube Core内部还有一层L0 Cache（约64KB），比L1 Local Memory更快。ops-blas通过调整MNK的遍历顺序，让同一个tile的B矩阵数据在L0 Cache中停留更久，减少从L1读取的次数。
具体策略：
| 矩阵规模 | 遍历顺序 | 原因 |
|---|---|---|
| M >> N, K | 先K后N | A的tile在L0里复用N次 |
| N >> M, K | 先K后M | B的tile在L0里复用M次 |
| 方阵（M=N=K） | 先N后M后K | 均衡复用 |
性能实测：ops-blas vs cuBLAS vs oneMKL
我在Ascend 910、A100、Xeon 8380上做了FP16 GEMM的对比测试：
| 矩阵规模 (M=N=K) | ops-blas (910) TOPS | cuBLAS (A100) TOPS | oneMKL (8380) GFLOPS |
|---|---|---|---|
| 128 | 12.3 | 15.7 | 89 |
| 512 | 98.7 | 105.3 | 312 |
| 1024 | 187.6 | 192.4 | 487 |
| 2048 | 223.4 | 219.8 | 523 |
| 4096 | 234.1 | 231.2 | 498 |
| 8192 | 236.2 | 238.7 | 467 |
| 矩阵规模 | ops-blas 峰值利用率 | cuBLAS 峰值利用率 |
|---|---|---|
| 128 | 4.8% | 5.1% |
| 1024 | 73.3% | 62.8% |
| 4096 | 91.8% | 75.5% |
| 8192 | 92.3% | 77.8% |
几个关键发现：
- 中等规模（1024-4096）ops-blas更快：因为Ascend 910的HBM带宽更大（1.2TB/s vs A100的2TB/s，但910的Cube Core数量更多）
- 极大规模（8192+）cuBLAS略优：A100的HBM2e带宽优势在大规模场景下更明显
- 小矩阵（<128）两者都很差：tiling和DMA搬运的开销超过了计算本身
不只是GEMM：ops-blas的其他算子
ops-blas不只是GEMM，它实现了BLAS Level 1-3的完整算子集：
| 算子 | BLAS级别 | 功能 | 适用场景 |
|---|---|---|---|
| GEMM | Level 3 | C = αAB + βC | 矩阵乘（MLP、Attention） |
| GEMV | Level 2 | y = αAx + βy | 矩阵向量乘（推理时全连接层） |
| SYR2K | Level 3 | C = αA·B^T + αB·A^T + βC | 对称秩2更新（二阶优化器） |
| TRSM | Level 3 | X = A^{-1}·B | 三角求解（线性回归） |
| TRMM | Level 3 | B = αA·B | 三角矩阵乘（Cholesky分解） |
| DOT | Level 1 | dot = x^T·y | 向量点积（相似度计算） |
| AXPY | Level 1 | y = αx + y | 向量缩放加（梯度更新） |
| NRM2 | Level 1 | n = ‖x‖₂ | 向量范数（正则化） |
GEMV的实战价值：推理场景下，batch_size=1时全连接层退化为GEMV（矩阵×向量），ops-blas的GEMV用了跟GEMM不同的tiling策略（沿M轴完整遍历，不切K），性能比通用GEMM快2-3倍。
fromops_blasimportgemv# 推理场景：单样本全连接层W=torch.randn(4096,768,dtype=torch.float16).npu()# 权重x=torch.randn(768,dtype=torch.float16).npu()# 单样本输入# 用gemv比gemm快2-3倍（不需要切K轴）y=gemv(W,x)# 比 gemm(W, x.unsqueeze(1)).squeeze(1) 快2-3x踩坑实录
坑1：小矩阵性能反而不如CPU
问题：矩阵规模小于128×128时，ops-blas的GEMM性能还不如CPU上的oneMKL。原因是tiling和DMA搬运的开销（~5μs）超过了计算本身（128×128 FP16 GEMM只要~0.5μs）。
解决方案：小矩阵用CPU算，或者攒成batch一次算：
# ❌ 小矩阵逐个算（慢）foriinrange(100):C=gemm(A_small[i],B_small[i])# 每次5μs开销，实际计算0.5μs# ✅ 攒成batch一次算（快）A_batch=torch.stack(A_small)# [100, M, K]B_batch=torch.stack(B_small)# [100, K, N]C_batch=torch.bmm(A_batch,B_batch)# 一次算完，均摊开销坑2：FP32 GEMM吞吐只有FP16的一半
问题：Ascend 910的Cube Core在FP32模式下吞吐只有FP16的一半（128 TOPS vs 256 TOPS），因为每个FP32矩阵乘占用的硬件资源是FP16的两倍。
解决方案：训练用FP16/BF16混合精度，只有精度敏感的部分（如Loss计算、梯度累积）用FP32：
fromops_blasimportgemm# ✅ 混合精度：计算用FP16，累加用FP32A=torch.randn(4096,4096,dtype=torch.float16).npu()B=torch.randn(4096,4096,dtype=torch.float16).npu()C=gemm(A,B,accumulate_dtype=torch.float32)# 内部用FP32累加，避免精度损失坑3：非方阵性能下降
问题：非方阵（比如M=1, N=4096, K=768，推理时的GEMV）性能比方阵差很多，因为tiling策略需要沿短边完整遍历，L0 Cache复用效率低。
解决方案：用专门的GEMV接口（见上文），或者沿batch维度并行。
ops-blas在CANN架构中的位置
ops-blas位于CANN五层架构的第2层（昇腾计算服务层），属于AOL算子库的核心组件：
第2层：昇腾计算服务层 ├─ AOL 算子库 │ ├─ BLAS算子（ops-blas）← 你在这里 │ ├─ NN算子（ops-nn） │ ├─ Transformer算子（ops-transformer） │ ├─ CV算子（ops-cv） │ └─ 融合算子依赖关系：opbase ← ops-blas，ops-blas ← ops-nn / ops-transformer / catlass
结尾
线性代数算子的性能天花板，取决于你对硬件架构的理解深度。ops-blas的92%峰值利用率不是魔法，是tiling、双缓冲、L0 Cache优化三板斧的工程结果。
如果你在做GEMM密集型计算（大模型训练、量化因子计算、科学计算），建议用ops-blas替代通用的torch.matmul。自动tiling已经帮你选好了最优参数，你只需要调gemm(A, B)这一行。但如果你的矩阵特别小（<128）或者特别非方，记得看看上面的踩坑方案。
https://atomgit.com/cann/ops-blas