昇腾CANN ops-blas：GEMM 在 NPU 上为什么可以快到极致

矩阵乘是所有深度学习计算的根。Attention、全连接、卷积展开——归根到底都是矩阵乘。ops-blas 是 CANN 里专门做高性能 GEMMGeneral Matrix Multiply的算子库核心目标是把昇腾 NPU 的 Cube 单元利用率拉到 90% 以上。ops-blas 和 ops-nn 的 MatMul 有重叠但分工不同ops-nn 侧重神经网络场景的融合算子MatMulBiasGELUops-blas 侧重纯计算的 GEMM 极致性能。如果 ops-nn 是一辆带自动挡的轿车ops-blas 就是赛道上拆了空调减重到极致的 F1。矩阵乘的硬件分块昇腾 NPU 的 Cube 单元一次算 16×16 的矩阵乘。对 M×K 的矩阵 A 和 K×N 的矩阵 B理论上需要 (M/16) × (N/16) × (K/16) 次 Mmad 指令。分块的粒度直接影响计算效率// ops-blas/gemm/gemm_tiling_strategy.cpp简化structGemmTiling{intM_block;// A 的行分块intN_block;// B 的列分块intK_block;// 公共维分块};GemmTilingselect_tiling(intM,intN,intK,DataType dtype){// L1 缓存可容纳的元素数考虑双缓冲需要 3 个 bufferconstexprintL1_CAPACITY192*1024/sizeof(float);// A 块: M_block × K_block// B 块: K_block × N_block// C 块: M_block × N_block// 总共必须 L1_CAPACITY// 经验公式: M_block ≈ N_block ≈ sqrt(L1_CAPACITY / 3)intoptimal(int)sqrt(L1_CAPACITY/3.0);// 对齐到 16 的倍数Cube 限制intMbmin((M15)/16*16,optimal/16*16);intNbmin((N15)/16*16,optimal/16*16);intKbmin((K15)/16*16,optimal/16*16);// 如果 M 很小batch1 的 FC 层调整策略if(M64){// M 小、N 大增加 M 块大小减少 load 的次数Mbmin(M,optimal/4*16);}if(N64){// N 小、M 大增加 N 块大小Nbmin(N,optimal/4*16);}return{Mb,Nb,Kb};}分块策略直接影响 Cube 利用率。分块太小Mmad 指令开销占比高分块太大L1 缓存溢出HBM 开始频繁换入换出。select_tiling的核心约束就是 L1 容量——三个 bufferA 块、B 块、C 块加起来不能超过 L1 的容量否则 Cube 会停下来等数据。GEMM 的大中小三条优化路径ops-blas 根据矩阵形状走三条不同的优化路径// ops-blas/gemm/gemm_dispatch.cppenumGemmPath{GEMM_LARGE,// M,N 1024: 大矩阵GEMM_MEDIUM,// M,N 在 64-1023: 中等矩阵GEMM_TINY// M,N 64: 小矩阵比如 batch 中的单样本};GemmPathclassify(intM,intN,intK){if(M1024N1024)returnGEMM_LARGE;if(M64||N64)returnGEMM_TINY;returnGEMM_MEDIUM;}Statusgemm(constTensorA,constTensorB,TensorC){auto[M,K]A.shape();auto[K2,N]B.shape();autopathclassify(M,N,K);switch(path){caseGEMM_LARGE:// 大矩阵多级 tiling packing// 对 A 和 B 做内存排布重排packing// 让 Cube 单元始终按连续地址读取returngemm_large(A,B,C);caseGEMM_MEDIUM:// 中等矩阵标准双缓冲流水returngemm_medium(A,B,C);caseGEMM_TINY:// 小矩阵batch 合并策略// 把多个小矩阵的 GEMM 合并成一个大矩阵乘// 比如把 32 个 [64, 768]×[768, 768] 合并成 [2048, 768]×[768, 768]returngemm_batched(A,B,C);}}大矩阵路径packing对矩阵 A 和 B 做内存重排把不连续的列主序或行主序重新排列成 Cube 单元需要的连续片。这个操作叫 packing需要一次额外的 HBM 拷贝但只要矩阵足够大packing 的时间会被 Cube 的峰值吞吐摊平。中等矩阵路径双缓冲不额外 packing用双缓冲流水让搬运和计算并行。搬运下一块数据的同时 Cube 在算当前块。小矩阵路径batch 合并单个小矩阵乘的开销全在启动延迟kernel launch HBM 初始化实际计算时间微乎其微。ops-blas 把多个小矩阵的 GEMM 调用合并成一次 batch GEMM——比如 32 个 [64,768]×[768,768] 合并成 [2048,768]×[768,768]延迟从 32 次降到 1 次。Batched GEMM 的实战LLaMA 的 attention 层里每个头独立做 Q×K^T32 个头就是 32 个独立的矩阵乘。ops-blas 的 batched GEMM 在这种场景下有天然优势// ops-blas/batched_gemm/batched_gemm_kernel.cpp__aicore__voidbatched_gemm_kernel(GlobalTensorfloatC,// [batch, M, N]GlobalTensorfloatA,// [batch, M, K]GlobalTensorfloatB,// [batch, K, N]intbatch,intM,intN,intK){// 把所有 batch 的矩阵拼接成大矩阵// A_merged: [batch*M, K] 等价于把所有 batch 沿 M 维拼接// B_merged: [K, N] 所有 batch 共享同一个 B// C_merged: [batch*M, N] 输出同样沿 M 维拼接intMergedMbatch*M;autotilingselect_tiling(MergedM,N,K);for(inti0;iMergedM;itiling.M_block){for(intj0;jN;jtiling.N_block){// 标准 GEMM tiling 逻辑// batch 维被合并到了 M 维// Cube 单元一次 Mmad 处理 16 个元素不再受 batch 拆分的影响LocalTensorfloata_tile(tiling.M_block,tiling.K_block);LocalTensorfloatb_tile(tiling.K_block,tiling.N_block);LocalTensorfloatc_tile(tiling.M_block,tiling.N_block);DataCopy(a_tile,A[i*K],{tiling.M_block,tiling.K_block});DataCopy(b_tile,B[j],{tiling.K_block,tiling.N_block});Mmad(c_tile,a_tile,b_tile,{tiling.M_block,tiling.N_block,tiling.K_block});DataCopy(C[i*Nj],c_tile,{tiling.M_block,tiling.N_block});}}}合并后的效果32 个独立的小 GEMM 变成 1 个大 GEMMCube 单元始终保持满负荷不用等 kernel launch。踩坑一GEMM 维度的顺序C A × B 的维度约定有三种常见记法M×K × K×N M×N标准、N×K × K×M 和 K×M × M×N。Pytorch 侧torch.nn.functional.linear用的是 M×K × K×N但torch.matmul可以接受任意维度组合。混淆维度顺序是高频错误。错误写法importtorch,torch_npu# 错误维度传反了# 想要计算 [B, M, K] × [B, K, N] [B, M, N]Atorch.randn(8,1024,768).npu().half()# [B, M, K]Btorch.randn(8,1024,768).npu().half()# [B, N, K] -- 错了! K 应该在最后# torch.matmul 会做广播但维度对应不对# 实际计算的是 [8, 1024, 768] × [8, 768, 1024]# K768, N1024 被当成内部维相乘输出形状 [8, 1024, 1024] 而非预期的 [8, 1024, 768]Ctorch.matmul(A,B)正确写法# 正确明确 K 维度位置Atorch.randn(8,1024,768).npu().half()# [B, M, K]Btorch.randn(8,768,1024).npu().half()# [B, K, N]Ctorch.matmul(A,B)# [B, 1024, 1024]# 或者用 transpose 显式控制B_transposedB.transpose(-1,-2)# [B, 1024, 768] - [B, 768, 1024]Ctorch.matmul(A,B_transposed)# [B, 1024, 1024]# 验证shape 是否正确assertC.shape(8,1024,1024),f期望 (8, 1024, 1024), 得到{C.shape}踩坑二Batch GEMM 的 stride 对齐Batched GEMM 合并多个矩阵拼成大矩阵时A 的各 batch 在 HBM 上是否连续存储决定了合并后是否需要额外的内存拷贝。错误写法# 错误A 的每个 batch 之间 stride 不对齐A_list[torch.randn(i1,768).npu().half()foriinrange(32)]# A_list[0] shape: [1, 768], A_list[1] shape: [2, 768]# 不同 batch 的 M 维不等长——不能直接拼接# 强行拼接会引入 padding 开销A_paddedtorch.nn.utils.rnn.pad_sequence(A_list,batch_firstTrue)# padding 填了额外的 0 行参与无效计算浪费算力正确写法# 正确保证所有 batch 的 M 维等长拼接后 stride 连续Atorch.stack([torch.randn(64,768)for_inrange(32)]).npu().half()# [32, 64, 768] -- 所有 batch M 恒定Btorch.randn(32,768,768).npu().half()# 方式一直接 batched matmulCtorch.bmm(A,B)# [32, 64, 768], ops-blas 内部自动合并# 方式二手动合并显式控制 strideA_mergedA.reshape(32*64,768)# [2048, 768]B_mergedB.reshape(32*768,768)# [24576, 768]C_mergedtorch.mm(A_merged,B_merged[:768])# 如果用相同的 B# 注意B 的 batch 不同时需要不同的拼接方式C 侧原理ops-blas 的 batched GEMM 内部用ldaleading dimension of A参数来表示 A 在 HBM 上的实际列步长。如果各 batch 的 M 不等长lda 每 batch 不同合并策略失效。踩坑三Packing 和转置的隐式开销Packing 需要对 A 或 B 做额外的内存拷贝和重排。如果矩阵太小packing 的开销占了总时间的大部分反而比不 packing 慢。错误写法# 错误小矩阵也强制 packingAtorch.randn(64,128).npu().half()Btorch.randn(128,64).npu().half()# 64x64 的 GEMM 计算量很小packing 的 HBM 拷贝开销占了大头# 1. packing: 拷贝 A 和 B 到连续的 Cube-friendly 内存布局: ~3us# 2. GEMM 计算: ~0.5usCube 几乎是瞬间算完的# packing 占比 85% -- 得不偿失# 这种小矩阵应该走 gemm_batched 合并而不是单独调用正确写法# 正确收集多个小 GEMM 一起跑A_batch[]B_batch[]for_inrange(256):A_batch.append(torch.randn(64,128).npu().half())B_batch.append(torch.randn(128,64).npu().half())# 合并后一次 batch GEMMA_stackedtorch.stack(A_batch)# [256, 64, 128]B_stackedtorch.stack(B_batch)# [256, 128, 64]Ctorch.bmm(A_stacked,B_stacked)# [256, 64, 64]# 合并后总计算量是大矩阵规模packing 占比降到 5% 以下性能实测Ascend 910 上FP16 GEMM 表现矩阵尺寸路径耗时Cube 利用率4096×4096×4096GEMM_LARGE1.82 ms92%1024×768×768GEMM_MEDIUM0.28 ms78%64×768×768GEMM_TINY(单独)0.12 ms35%512×64×768×768GEMM_TINY(batch合并)2.10 ms91%小矩阵单独调 Cube 利用率只有 35%——kernel launch 和 HBM 初始化占了大头。512 个合并成 batch 后利用率回到 91%证明合并策略对小矩阵场景是刚需。GEMM 是个数学上极其简单的操作——只有乘法和加法。复杂度全在工程上多大的块、要不要 packing、怎么处理不规则的 batch。ops-blas 的三路径分法大矩阵 packing 中等矩阵双缓冲 小矩阵 batch 合并覆盖了从 attention 到 FC 层的全部矩阵乘场景Cube 利用率 90% 以上是这套策略体系的目标不是单个 kernel 的能力。