昇腾CANN ops-blas 仓：GEMM 算子的高性能实现

前言矩阵乘法是深度学习里最核心的操作没有之一。Transformer 的 Attention 要做 QK.T 和 PVFFN 要做两 个 MatMul。GEMMGeneral Matrix Multiply就是专门优化矩阵乘的算子。ops-blas 仓是 CANN 的线性代数基础算子库GEMM 是它的核心产品。这篇文章拆开看它怎么把 Cube 单元跑满的。GEMM 在深度学习中的地位GEMM 的全称是 General Matrix Multiply做的是 C alpha * A * B beta * C 这种通用矩阵乘法。在深度学习里的典型用法是# PyTorch 里的 Linear 层本质就是一个 GEMM# Y X W^T b# 写成 GEMM 格式就是 Y 1.0 * X W 0.0 * Y# 其中 X 是 (batch, input_dim)W 是 (output_dim, input_dim)# 矩阵乘的结果是 (batch, output_dim)Transformer 的核心计算几乎全是矩阵乘Attention 里的 QK.T 是 (batch, head, seq, head_dim) (batch, head, head_dim, seq) - (batch, head, seq, seq)FFN 里的两个 MatMul 分别是 hidden - intermediate 和 intermediate - hidden。大模型训练和推理的计算量 70% 以上都花在矩阵乘上优化 GEMM 就是优化整个模型。昇腾达芬奇架构的 Cube 单元昇腾达芬奇架构的计算核心是两个单元Cube 单元和 Vector 单元。Cube 单元专门做矩阵乘Vector 单元做向量和标量运算。Cube 单元的名字来自 3D Cube——它一次能处理三维张量的矩阵乘。具体来说Cube 单元一次可以做 16×16×16 的矩阵乘累加。这个 16×16×16 来自硬件设计每个 cycle 能跑 4096 个乘累加运算MAC在 FP16 精度下峰值算力是 256 TFLOPS Ascend 910。理解 Cube 单元的关键是tiling分块。要把一个大矩阵乘拆到 Cube 单元上跑需要把矩阵切成小块。每个小块要能装进 L1 Buffer然后交给 Cube 单元处理。tile 的大小选择直接影响性能太大了 L1 Buffer 不够用需要频繁读写 HBM太小了 Cube 单元的并行度上不去。ops-blas 仓的核心工作就是设计 tiling 策略——怎么切块能让 Cube 单元的利用率最高同时让 HBM 访问最少。ops-blas GEMM 的分块策略ops-blas 的 GEMM 实现用了多层 tiling。简单说就是“大块套小块”第一层是Core 级别的 tiling。Ascend 910 有多个 AI Core每个 Core 负责矩阵的一部分。ops-blas 会把矩阵按 Core 数量做切分把 A 按行切、把 B 按列切、每个 Core 拿自己那份去算。第二层是Tile 级别的 tiling。每个 Core 内部再把任务分成多个 tile。每个 tile 要满足两个约束能装进 L1 Buffer、能让 Cube 单元跑满。典型配置下A 按 16×K 切B 按 K×16 切C 产 16×16 的结果块。第三层是指令级别的 tiling。Cube 单元内部还有一层微 tiling用指令流水来隐藏内存访问延迟。看一段简化版的伪代码理解 tiling 逻辑// GEMM 核心计算A B - C// 这里展示 tiling 的思路voidgemm_core(half*A,half*B,half*C,intM,intN,intK,intM_tile,intN_tile,intK_tile){// M 方向切成 M_tile 大小的块for(inti0;iM;iM_tile){// N 方向切成 N_tile 大小的块for(intj0;jN;jN_tile){// C(i:iM_tile, j:jN_tile) // A(i:iM_tile, :) B(:, j:jN_tile)// 每个 C 块内部按 K 方向切half C_tile[M_tile][N_tile]{0};for(intk0;kK;kK_tile){// 加载 A 块从 HBM 到 L1// 这个 block 大小要能装进 L1 Bufferload_block(A,i,k,M_tile,K_tile);// 加载 B 块load_block(B,k,j,K_tile,N_tile);// Cube 单元执行矩阵乘// 一次跑 16x16x16 的 MACcube_mm(A_block,B_block,C_tile);}// 把结果写回 Cstore_block(C,i,j,C_tile);}}}双缓冲隐藏 HBM 访问延迟GEMM 的瓶颈往往不在计算而在数据搬运。从 HBM 加载 A 块和 B 块到 L1 的时间远大于 Cube 单元计算的时间。ops-blas 用双缓冲double buffering来解决这个问题。双缓冲的核心是算第 j 块的同时搬运第 j1 块。这样计算和数据搬运并行进行HBM 带宽的延迟被藏在 Cube 计算的背后。// 双缓冲示例计算和搬运并行voidgemm_with_double_buffer(half*A,half*B,half*C,intM,intN,intK){// 准备两个 buffer 轮换用half A_buf0[BLOCK_A],A_buf1[BLOCK_A];half B_buf0[BLOCK_B],B_buf1[BLOCK_B];intbuf_idx0;// 预加载第一块load_async(A_buf0,A0);load_async(B_buf0,B0);for(intk0;kK;kBLOCK_K){// 等待当前块加载完成wait_load();// 启动下一块的异步加载if(kBLOCK_KK){load_async(A_buf[1-buf_idx],A(kBLOCK_K));load_async(B_buf[1-buf_idx],B(kBLOCK_K));}// Cube 单元计算当前块cube_mm(A_buf[buf_idx],B_buf[buf_idx],C_block);// 切换 bufferbuf_idx1-buf_idx;}}这段代码展示了双缓冲的思路不是等上一块算完才开始搬下一块而是算着当前块的同时搬下一块。硬件上DMA 引擎负责 HBM 搬运和 Cube 单元负责计算是独立运行的只要调度得当两者可以完美 overlap。L1 Cache 优化除了双缓冲还有一个关键优化是L1 Cache 的利用。Cube 单元每次计算的输入 A 和 B 可以复用同一个 A 块要和多个 B 块相乘同一个 B 块要和多个 A 块相乘。把常用的数据块保持在 L1 Cache 里能大幅减少 HBM 访问。ops-blas 的 tiling 策略专门考虑了缓存复用A 块在 K 方向复用A(i, k) 这个块会和 B(k, j) 的所有 j 相乘所以 A 块一次性加载后可以留在 L1 里很久B 块在 M 方向复用B(k, j) 这个块会和 A(i, k) 的所有 i 相乘但复用的机会比 A 少一些实际效果是HBM 访问量能降到理论最低值的 1/3 到 1/2。性能数据不同配置下的实测数据Ascend 910FP16配置TFLOPS利用率MNK102423090%MNK204824595%MNK409625097%可以看到当矩阵尺寸变大时利用率更高因为大矩阵的缓存命中率更高HBM 延迟能被更好地隐藏。跟其他实现对比实现TFLOPSops-blas GEMM250cuBLAS (NVIDIA A100)312理论峰值256昇腾的 Cube 单元利用率已经非常接近理论峰值了。跟 NVIDIA 的差距主要在峰值算力上A100 的 Tensor Core 峰值比 Ascend 910 高但软件层面的优化已经做到位了。如何调用PyTorch 调用 GEMM 最简单的方式是通过 Linear 层importtorchimporttorch_npu# Linear 内部就是 GEMMlineartorch.nn.Linear(4096,11008).npu()xtorch.randn(1,4096,dtypetorch.float16).npu()# forward 会调用 ops-blas.GEMMylinear(x)print(y.shape)# (1, 11008)如果想直接调用 GEMM用于自定义算子开发可以用 AscendCL 接口importacl# 初始化 ACLacl.init()# 创建 GEMM 算子gemm_opacl.op.create_gemm(transaFalse,# A 不转置transbFalse,# B 不转置m1024,n1024,k1024,alpha1.0,beta0.0,a_formatND,b_formatND)# 执行aacl.malloc(1024*1024*2)# FP16bacl.malloc(1024*1024*2)cacl.malloc(1024*1024*2)gemm_op(a,b,c)GEMM 是深度学习计算的基础设施。ops-blas 把昇腾的 Cube 单元压榨到了接近理论峰值。对于做模型优化的人来说理解 GEMM 的 tiling 策略和缓存优化是进一步提升性能的前提。仓库地址https://atomgit.com/cann/ops-blas