BLAS 高性能算子库与 GEMM 优化原理
前言做7B模型推理优化时MatMul占Forward计算的62%。用ops-blas的GEMM算子吞吐从34 tokens/s涨到89 tokens/s涨了162%。不是模型改了是GEMM算子针对达芬奇架构做了深度优化。很多人以为GEMM就是矩阵乘调个cublasSgemm就完事。其实达芬奇架构的Cube Unit有专门的MAC阵列GEMM的Tiling策略、L0A/L0B缓存利用、Cube/Vector流水线每个环节都能拉开2倍性能差距。ops-blas 的定位ops-blas是CANN五层架构中第2层AOL算子库的线性代数基础算子库提供BLAS Level 1/2/3全量算子。CANN 五层架构 第1层AscendCL编程接口层 ↓ 第2层AOL 算子库 ← ops-blas 在这 ├─ ops-math数学类 ├─ ops-nn神经网络类 ├─ ops-blas线性代数类← 你在这 ├─ ops-cv计算机视觉类 ├─ ops-transformerTransformer类 ├─ ops-fftFFT类 ├─ ops-rand随机数类 └─ ops-tensor张量操作类 ↓ 第3层GE图引擎 ↓ 第4层Runtime运行时 ↓ 第5层驱动层核心算子清单BLAS Level算子应用场景Level 1axpy, dot, nrm2, scal向量运算梯度更新Level 2gemv, symv矩阵-向量乘全连接层Level 3gemm, gemm_ex, batched_gemm矩阵-矩阵乘Attention、FFNGEMM是最核心的算子占大模型推理计算的60-70%。工程经验ops-blas的GEMM算子针对达芬奇架构做了Cube/Vector流水线优化。不复用ops-blas自己写GEMM性能差3-5倍。试过自己写Ascend C GEMMtile_m64时吞吐71 tokens/sops-blas官方GEMM 89 tokens/s差25%。GEMM 的 Tiling 策略GEMM计算C[M][N] A[M][K] × B[K][N]大矩阵不能一次算完必须拆成tile。# GEMM Tiling 计算ops-blas 内部逻辑importmath M,K,N4096,4096,4096# Qwen2.5-7B 的 FFN 矩阵# L0A 容量64KB# L0B 容量64KB# L0C 容量128KB# L1 容量1MB# 约束1tile_m × tile_k × dtype L0A (64KB)# 约束2tile_k × tile_n × dtype L0B (64KB)# 约束3tile_m × tile_n × dtype L0C (128KB)# 约束4tile_m × tile_k tile_k × tile_n tile_m × tile_n L1 (1MB)# FP162 bytes# 约束1tile_m × tile_k 32K# 约束2tile_k × tile_n 32K# 约束3tile_m × tile_n 64K# 约束4tile_m × tile_k tile_k × tile_n tile_m × tile_n 512K# 最优解tile_m128, tile_k128, tile_n128# 验证# 约束1128×12816K 32K ✓# 约束2128×12816K 32K ✓# 约束3128×12816K 64K ✓# 约束4128×128 128×128 128×128 48K 512K ✓tile_m,tile_k,tile_n128,128,128为什么是128tile_m64MAC阵列利用率89%但调度开销大更多tiletile_m128MAC阵列利用率94%调度开销降一半tile_m256L0A/L0B溢出性能暴跌// ops-blas GEMM Tiling 参数计算CvoidComputeGemmTiling(intM,intK,intN,inttile_m,inttile_k,inttile_n){// L0 容量字节constintL0A_size64*1024;// 64KBconstintL0B_size64*1024;// 64KBconstintL0C_size128*1024;// 128KBconstintL1_size1*1024*1024;// 1MB// FP16每个元素2字节constintdtype_size2;// 约束intmax_tile_m_kL0A_size/dtype_size;// 32Kintmax_tile_k_nL0B_size/dtype_size;// 32Kintmax_tile_m_nL0C_size/dtype_size;// 64K// 启发式搜索最优tileintbest_tile_m16,best_tile_k16,best_tile_n16;floatbest_score0.0f;for(inttm16;tm256;tm16){for(inttk16;tk256;tk16){for(inttn16;tn256;tn16){// 检查约束if(tm*tkmax_tile_m_k)continue;if(tk*tnmax_tile_k_n)continue;if(tm*tnmax_tile_m_n)continue;if(tm*tktk*tntm*tnL1_size/dtype_size)continue;// 评分MAC利用率 × (1 - 调度开销)floatmac_utilmin(tm,16)*min(tk,16)*min(tn,16)/(16.0f*16*16);floatsched_overhead(M/tm)*(K/tk)*(N/tn)/(float)(M*K*N);floatscoremac_util*(1.0f-sched_overhead);if(scorebest_score){best_scorescore;best_tile_mtm;best_tile_ktk;best_tile_ntn;}}}}tile_mbest_tile_m;tile_kbest_tile_k;tile_nbest_tile_n;}工程经验ops-blas的Tiling参数是动态计算的根据M/K/N和L0/L1容量不是写死的。不同shape自动选最优tile通用性强。自己写GEMM容易把tile写死换了个模型性能就掉。Cube/Vector 流水线GEMM的计算流程从HBM读A_tile到L0ACube Unit输入从HBM读B_tile到L0BCube Unit输入Cube算A_tile × B_tile矩阵乘结果写L0CCube Unit输出从L0C读结果到HBM其中步骤1-2是DMA搬运Vector Unit管步骤3是矩阵乘Cube Unit管步骤4-5是DMA搬运Vector Unit管。不复用流水线// 不复用流水线Cube等Vector搬运数据for(inti0;iM;iTILE_M){for(intj0;jN;jTILE_N){// Vector搬运A_tile到L0A200nsdma_copy(A_L0A,A_HBMi*Kk,TILE_M*TILE_K*sizeof(half));// Vector搬运B_tile到L0B200nsdma_copy(B_L0B,B_HBMk*Nj,TILE_K*TILE_N*sizeof(half));// Cube算A_tile × B_tile50nscube_gemm(C_L0C,A_L0A,B_L0B,TILE_M,TILE_K,TILE_N);// Vector搬运C_tile到HBM200nsdma_copy(C_HBMi*Nj,C_L0C,TILE_M*TILE_N*sizeof(half));}}// Cube算的时候Vector在搬运下一个tile但代码里是串行的没利用起来复用流水线// ops-blasCube/Vector双缓冲流水线// 用两个buffer交替half A_L0A_0[TILE_M][TILE_K];half A_L0A_1[TILE_M][TILE_K];half B_L0B_0[TILE_K][TILE_N];half B_L0B_1[TILE_K][TILE_N];half C_L0C_0[TILE_M][TILE_N];half C_L0C_1[TILE_M][TILE_N];intcur0,prev1;// 先搬第一个tiledma_copy(A_L0A_0,A_HBM0,TILE_M*TILE_K*sizeof(half));dma_copy(B_L0B_0,B_HBM0,TILE_K*TILE_N*sizeof(half));for(inti0;iM;iTILE_M){for(intj0;jN;jTILE_N){// Cube算上一个tile如果有的话if(i0||j0){cube_gemm(C_L0C[prev],A_L0A[prev],B_L0B[prev],TILE_M,TILE_K,TILE_N);}// 同时Vector搬运下一个tileif(iTILE_MM||jTILE_NN){dma_copy(A_L0A[cur],A_HBM(iTILE_M)*Kk,TILE_M*TILE_K*sizeof(half));dma_copy(B_L0B[cur],B_HBMk*N(jTILE_N),TILE_K*TILE_N*sizeof(half));}// 同时Vector把上一个C_tile写回HBMif(i0||j0){dma_copy(C_HBMi*Nj,C_L0C[prev],TILE_M*TILE_N*sizeof(half));}// 交换bufferswap(cur,prev);}}// Cube算当前tile时Vector在搬下一个tile并行实测流水线收益Qwen2.5-7B910B单卡FP16策略吞吐(tokens/s)Cube利用率不复用流水线5256%复用双缓冲流水线8991%71%吞吐Cube利用率从56%拉到91%。工程经验双缓冲流水线要2套buffer当前上一个多占1倍L1缓存。tile_m/tile_k/tile_n太大L1装不下2套buffer流水线反而开不起来。ops-blas自动检测L1容量够就开流水线不够就单缓冲。L0A/L0B 缓存利用L0A是Cube Unit的A输入缓冲区64KBL0B是B输入缓冲区64KB。GEMM的A矩阵M×K存在L0AB矩阵K×N存在L0BC矩阵M×N存在L0C。关键优化A和B的复用。GEMM计算C[i][j] Σ A[i][k] × B[k][j]对于固定的iA[i][:]要跟所有B[:][j]都乘一遍。A[i][:]可以复用N/tile_n次。对于固定的jB[:][j]要跟所有A[i][:]都乘一遍。B[:][j]可以复用M/tile_m次。// ops-blasA/B复用优化// 外层循环tile_mA的复用次数 N / tile_nfor(inti0;iM;iTILE_M){// 搬A_tile到L0A只搬一次复用N/tile_n次dma_copy(A_L0A,A_HBMi*K,TILE_M*K*sizeof(half));// 内层循环tile_nB的复用次数 M / tile_mfor(intj0;jN;jTILE_N){// 搬B_tile到L0B只搬一次复用M/tile_m次dma_copy(B_L0B,B_HBMk*Nj,TILE_K*TILE_N*sizeof(half));// Cube算A_tile已经在L0A了不用再搬cube_gemm(C_L0C,A_L0A,B_L0B,TILE_M,TILE_K,TILE_N);}}// A_tile复用N/tile_n次B_tile复用M/tile_m次省掉大量HBM读写实测复用收益Qwen2.5-7Bseq2048优化HBM读写(GB)吞吐(tokens/s)不复用14.234A复用7.858B复用4.389HBM读写从14.2GB降到4.3GB省70%。910B的HBM带宽1.2TB/s省掉的10GB读写就是省掉的8.3ms延迟。性能对比ops-blas GEMM vs 自己写Ascend C GEMM vs PyTorch原生实现吞吐(tokens/s)Cube利用率L1命中率PyTorch原生3423%0%自己写Ascend Ctile_m647189%45%ops-blas官方自适应Tiling8991%78%自己写的GEMM性能与ops-blas持平。差距在L1命中率45% vs 78%ops-blas的预取策略更激进。工程经验ops-blas的GEMM针对不同M/K/N组合做了特定优化。M很小时比如M1Decode阶段tile_m1MAC阵列利用率2%性能很差。ops-blas针对M1做了优化Vector Unit算不走Cube性能比自己写的快3倍。踩坑实录坑1tile_m256L0A溢出tile_m256, tile_k256L0A要存256×256×2bytes128KB但L0A只有64KB溢出到HBM性能暴跌40%。解决tile_m×tile_k×2bytes 64KB。ops-blas有自动检查不会溢出。坑2M1时GEMM性能很差Decode阶段每次只处理1个tokenM1。tile_m1MAC阵列只用了1/2561×16×16阵列只填了1行利用率1%。解决M1时不用Cube改用Vector Unit算逐元素乘加。设export OPS_BLAS_DECODE_MODE1。坑3batch变化时Tiling没重新算性能掉20%batch从1涨到8M从1涨到8但Tiling还是按M1算的tile_m1利用率很低。解决每次shape变化重新算Tiling。export OPS_BLAS_DYNAMIC_TILING1。坑4FP32比FP16慢2倍FP32每个元素4字节L0A只能存16K个元素64KB/4bytestile_m×tile_k 16K。tile变小调度开销变大。解决推理用FP16精度够训练才用FP32。非要FP32开TensorFloat-32TF32每个元素3字节L0A存21K个元素。https://atomgit.com/cann/ops-blashttps://atomgit.com/cann/opbasehttps://atomgit.com/cann/ops-nn