cuBLAS矩阵乘法性能对比:手写kernel与官方库优化解析

📅 发布时间:2026/9/11 11:26:21
cuBLAS矩阵乘法性能对比:手写kernel与官方库优化解析
矩阵乘法大概是整个GPU计算世界里最“值钱”的一个运算。全连接层、注意力机制里的QK^T、卷积展开后的im2col、各种数值求解器的核心步骤最后几乎都能落到GEMM上。这也是我在这套深入篇里专门开一节讲cuBLAS的原因你迟早要在工程里和它打交道与其每次稀里糊涂地照抄API不如把背后的调用规则和性能逻辑一次搞明白。这篇不灌水。我会先写一个朴素的CUDA矩阵乘法和一个共享内存tile版本然后用它们跟cuBLAS的SGEMM做一组公平对比再把数字背后的优化手段拆开讲清楚。适合谁看已经写过CUDA kernel但没系统用过cuBLAS的朋友以及一直好奇“官方库为什么比我手搓的快这么多”的人。看完你至少能直接上手调用也能跟别人解释清楚性能差距到底是怎么来的。1. 为什么矩阵乘法值得单独讲一节1.1 从全连接到注意力机制底层都是GEMM先说一个直觉判断你在GPU上写的绝大多数“正经”算法拆到最底层基本都绕不开矩阵乘法。拿深度学习举例全连接层的本质就是 y x·W b这是典型的GEMM卷积层如果不做im2col那也逃不掉各种类似矩阵乘的滑窗累加Transformer里的自注意力Q和K^T做点积K和V做加权和两个核心操作全是矩阵乘法。所以NVIDIA专门搞了一个BLAS库来做这件事cuBLAS就是GPU版的BLAS实现。它提供了一大堆经过手工调优的线性代数函数从最基础的向量加法、矩阵乘法到LU分解、SVD、特征值分解都有。而我们日常碰到最多的就是矩阵乘法也就是BLAS里的GEMMGEneral Matrix Multiply。我见过不少同事算法写得很熟练一听说要调cuBLAS就有点发怵觉得又是句柄又是列主序又是各种参数太麻烦不如自己写个三重循环算了。但实际情况是你手写的kernel在常规规模下跟cuBLAS差着几十倍到上百倍的性能。这个差距不是靠多写几行#pragma unroll能追上的。1.2 cuBLAS到底解决了什么问题cuBLAS解决的核心问题可以归纳成三个正确性、性能、可维护性。正确性方面GPU浮点运算有很多隐蔽的坑比如累加顺序不同导致的结果差异、不同精度模式的取舍、边界处理等。官方库把这些细节都封装好了你传对参数就行。性能方面cuBLAS在底层用了非常复杂的分块策略、寄存器级优化、内存访问重排乃至针对不同GPU架构的自动调优这些在普通工程里根本不可能自己从头写。可维护性方面如果你今天要换GPU型号手写kernel可能要重新调优而cuBLAS的API设计基本不变它内部会自己适配。不过这里也要说句公道话cuBLAS不是万能的。当矩阵规模特别小比如32×32以下时kernel启动开销占比变大手写专用内核反而可能更快另外cuBLAS默认只追求通用吞吐不一定针对你的特定场景做形状定制。所以“手写”和“调库”并不是非此即彼的关系而是不同场景下的不同选择。理解了这一点再看后面的性能数据你会更清醒。1.3 本节内容导览与前置要求这一节的安排是这样的先自己动手写两个GPU矩阵乘法版本一个朴素版、一个共享内存tile版这是后面性能对比的“对照组”然后正式引入cuBLAS把句柄、API参数、行列主序这些关键概念一次说清接着跑一组实测把三个版本的性能数据放在一张表里看最后聊一聊常见坑和排查思路。前置要求不高你要能写基本的CUDA kernel知道block、thread、shared memory是什么概念会用nvcc编译。如果你是在这方面零基础建议先看看本系列前面几节入门篇再回来读这篇效果会好很多。如果你已经有基础直接往下走代码和踩坑部分应该能给你一些实用收获。2. 先写一个能打的“手写版”矩阵乘法2.1 朴素版本人人都会的写法先看最直接的思路计算C(m×n) A(m×k) × B(k×n)让每个线程负责C矩阵中的一个元素。这个线程需要遍历p 0到k-1累加A[row][p] * B[p][col]。代码如下__global__ void sgemm_naive(const float* A, const float* B, float* C, int m, int n, int k) { int row blockIdx.y * blockDim.y threadIdx.y; int col blockIdx.x * blockDim.x threadIdx.x; if (row m col n) { float sum 0.0f; for (int p 0; p k; p) { sum A[row * k p] * B[p * n col]; } C[row * n col] sum; } }启动配置上用二维block大小(16, 16)网格按(m/16, n/16)取整ceil计算。这个kernel逻辑没有任何问题但它跑起来会很慢慢到让你怀疑GPU是不是坏了。为什么慢关键在于全局内存的访问模式。内层循环每迭代一次线程都要从全局内存读两个float。更致命的是同一个A[row][p]会被同一行所有col方向的线程重复读取同一个B[p][col]会被同一列所有row方向的线程重复读取。以16×16的block为例一个A元素要被16个线程重复读一个B元素也要被16个线程重复读全局内存带宽直接被浪费掉一个数量级。我实测过一个具体场景在RTX 3090上这个内核算1024×1024的矩阵乘法性能大约只有50 GFLOPS左右而这张卡的FP32理论峰值有35 TFLOPS左右。换句话说我们只用了显卡0.1%的算力这就是入门教材通常不告诉你的事。2.2 共享内存tile版本第一次优化要解决全局内存重复读取的问题常规思路是利用共享内存做分块。把A和B各切成8×8或16×16的小块每次把一个tile加载到shared memory然后让block内的所有线程复用这块数据。这样每个全局内存数据只需要从显存读一次却被使用了TILE次全局内存流量降为原来的1/TILE。代码如下我用16×16的tile#define TILE 16 __global__ void sgemm_tiled(const float* A, const float* B, float* C, int m, int n, int k) { __shared__ float As[TILE][TILE]; __shared__ float Bs[TILE][TILE]; int row blockIdx.y * TILE threadIdx.y; int col blockIdx.x * TILE threadIdx.x; float sum 0.0f; for (int tileIdx 0; tileIdx (k TILE - 1) / TILE; tileIdx) { // 协作加载A的tile int aRow row; int aCol tileIdx * TILE threadIdx.x; if (aRow m aCol k) As[threadIdx.y][threadIdx.x] A[aRow * k aCol]; else As[threadIdx.y][threadIdx.x] 0.0f; // 协作加载B的tile int bRow tileIdx * TILE threadIdx.y; int bCol col; if (bRow k bCol n) Bs[threadIdx.y][threadIdx.x] B[bRow * n bCol]; else Bs[threadIdx.y][threadIdx.x] 0.0f; __syncthreads(); // 从shared memory累加 #pragma unroll for (int p 0; p TILE; p) { sum As[threadIdx.y][p] * Bs[p][threadIdx.x]; } __syncthreads(); } if (row m col n) C[row * n col] sum; }这个版本的逻辑也不复杂但性能提升非常明显。核心变化有两点一是全局内存访问变成了分段式、可复用的访问二是内层计算直接从shared memory读数据延迟比全局内存低了一个数量级。实测下来同样的1024×1024矩阵这个kernel能做到400 GFLOPS左右比朴素版提升了大约8倍。不过这个版本仍然远没到硬件极限。它的问题在于shared memory的读取次数还是太多。每个线程在p循环里要读TILE次As、TILE次Bs而这些数据其实可以在寄存器里复用另外我们没有做任何向量化加载也没有考虑bank conflict更别说双缓冲这些进阶手段了。2.3 手写到这个程度还能往哪里走共享内存tile版本其实是很多教材里“优化demo”的终点但在真实的高性能计算里这只是起跑线。后面还可以做寄存器级tile让每个线程一次计算多个输出元素把数据在寄存器里反复复用可以引入float4向量化加载让每次内存事务传输4个float可以用双缓冲或者软件流水线让shared memory的加载和计算重叠还可以针对不同的GPU架构调整tile大小和unroll因子。我自己早年也走过这条路从朴素版优化到寄存器tile 向量化 double buffering大概能让性能再翻好几倍。但越往后越会发现每一步都需要对硬件细节有极深的理解而且要反复实验。这也是为什么cuBLAS一上来就能跑出比你折腾几天之后还好的性能——因为它背后是NVIDIA工程师多年积累的自动调优成果。所以结论很明确如果你想深入了解GPU优化手写kernel是极好的学习路径但如果你要完成实际项目直接用cuBLAS才是性价比最高的选择。下面我们就正式进入cuBLAS的调用环节。3. cuBLAS调用代码比想象中简单3.1 句柄与库初始化cuBLAS的API设计风格和其他CUDA库类似第一步都是创建一个库句柄handle。句柄相当于一个上下文对象保存了库的内部状态、设备信息、算法配置等。创建和销毁的代码如下#include cublas_v2.h cublasHandle_t handle; cublasStatus_t status cublasCreate(handle); if (status ! CUBLAS_STATUS_SUCCESS) { // 处理错误 } // ... 执行计算 ... cublasDestroy(handle);有一点要注意句柄绑定创建它的线程所在的设备和CUDA context。在多GPU场景下如果你要在device 0和device 1上分别计算就需要各创建一个句柄或者用cublasSetStream/cudaSetDevice切换上下文。另外如果工程里有多个线程并发调用cuBLAS建议每个线程持有自己的handle避免句柄内部的流和状态被多个线程互相踩踏。从实践角度讲我习惯把handle的创建放在一次性的初始化函数里放到全局或者类的成员变量中而不是每次调用都重新创建。因为cublasCreate底层会初始化一些内部缓冲反复创建销毁会有不必要的开销。3.2 cublasSgemm参数逐项拆解我们这节以单精度矩阵乘法为例子对应的函数是cublasSgemm。完整函数签名如下cublasStatus_t cublasSgemm( cublasHandle_t handle, cublasOperation_t transa, // A是否转置 cublasOperation_t transb, // B是否转置 int m, // C矩阵的行数 int n, // C矩阵的列数 int k, // 累加维度 const float* alpha, // 标量alpha const float* A, int lda, // A矩阵及其leading dimension const float* B, int ldb, // B矩阵及其leading dimension const float* beta, // 标量beta float* C, int ldc); // C矩阵及其leading dimension这个API计算的是C alpha * op(A) * op(B) beta * C。其中op(A)取决于transa参数如果传CUBLAS_OP_N就不转置传CUBLAS_OP_T就转置。这里有个容易搞混的点当我们说“A是m×k”时默认是在“不转置”的视角下如果transa CUBLAS_OP_T那么op(A)是k×m所以C的维度会变成k×n。再看lda、ldb、ldc这三个参数。ld是leading dimension的缩写意思是“主维步长”。在列主序存储下它等于矩阵的行数在行主序存储下它等于矩阵的列数。很多新手在这里翻车把lda直接填成矩阵的“逻辑行数”m实际上得看你的数据在内存里是怎么排的以及你当前用的是行主序还是列主序视角。还有alpha和beta。默认情况下这两个参数是指向主机内存的指针也就是说你可以直接传alpha。cuBLAS内部在kernel启动时会自动拷贝到设备端。如果你要用设备端指针控制alpha/beta需要先调用cublasSetPointerMode(handle, CUBLAS_POINTER_MODE_DEVICE)然后传入设备内存地址。这个特性在动态修改标量时有用但不是我们这节的重点。3.3 行主序和列主序的恩怨lda与转置这是cuBLAS调用里最容易踩的坑值得单独说一说。cuBLAS的接口沿用了Fortran的列主序约定。也就是说一个m×k的矩阵在内存里是“先存第0列再存第1列……”同一列的元素在地址上是连续的。而C/C里的原生二维数组是行主序同一行的元素在地址上连续。当你直接用C的行主序数据去喂cuBLAS它会把你的数据误解成另一个矩阵行主序的A(m×k)在cuBLAS眼里其实是A^T(k×m)。怎么办最简单的办法是显式转置在主机端把A转置成A_T后再拷到设备端然后正常把A_T当作m×k的列主序矩阵传入。这种方法的优点是理解简单代码不容易绕晕缺点是每次调用要多一次转置的开销。更优雅的做法是利用transa/transb参数不显式转置直接换一下传参顺序。推演一遍你就明白了行主序的A(m×k)在内存中对应列主序的A^T(k×m)行主序的B(k×n)对应列主序的B^T(n×k)。你想算行主序的R(m×n) A·B等价于算列主序的R^T(n×m) B^T·A^T。于是调用就变成// 行主序数据直接使用不手动转置 // 目标R(m×n) A(m×k) * B(k×n) // cuBLAS视角R^T(n×m) B^T · A^T cublasSgemm(handle, CUBLAS_OP_N, CUBLAS_OP_N, n, m, k, // 注意顺序 alpha, d_B, n, // 第一个矩阵是B^T列主序下lda n d_A, k, // 第二个矩阵是A^Tldb k beta, d_R, n); // C^T的leading dimension是n简单记就是B和A的位置互换输出维度从m×n换成n×mlda用第一个矩阵的行数ldb用第二个矩阵的行数。我早期用cuBLAS的时候在这上面折腾了很久。后来我给自己定了个规矩管线内部统一用列主序这样传参清晰只有在和外部数据交互的边界才做转置处理。如果你刚开始接触建议先写一个小规模的测试矩阵打印结果跟CPU计算结果对比确认无误再集成到正式代码里这会帮你省下大量的排查时间。4. 性能对比用数字说话4.1 测试方法与计时工具性能对比最怕不公平。手写kernel和cuBLAS如果不在同等条件下测试得到的数据就是自欺欺人。我这边统一采用以下方法。首先是矩阵规模。分别测试512、1024、2048、4096这四档矩阵类型是方阵即m n k。数据随机初始化确保没有稀疏性带来的偶然加速。接下来是计时方式GPU kernel是异步执行的用CPU的clock_gettime直接包住kernel调用测到的时间基本是错的。正确做法是使用CUDA eventcudaEvent_t start, stop; cudaEventCreate(start); cudaEventCreate(stop); cudaEventRecord(start); // 执行kernel或cublasSgemm调用 cudaEventRecord(stop); cudaEventSynchronize(stop); float ms 0.0f; cudaEventElapsedTime(ms, start, stop);然后根据矩阵规模换算GFLOPS。矩阵乘法的浮点运算次数是2×m×n×k一次乘和一次加各算一次浮点操作所以GFLOPS 2×m×n×k / (ms × 1e6)。我在每次测试前会执行一次“热身”调用再连续跑5次取中位数。原因很简单第一次调用可能涉及缓存未命中、驱动编译shader等额外开销实际成绩偏低而GPU频率呢随着温度升高会波动取中位数能更稳定地反映典型性能。4.2 实测数据与解读测试环境是RTX 3090、CUDA 12.4、编译选项-O3。矩阵规模从512到4096三种实现分别是朴素kernel、共享内存tile kernel、cuBLAS SGEMM。单位是GFLOPS数值越高越好矩阵规模朴素kernel共享内存tilecuBLAS SGEMM512293128230102451425158002048635102270040967553425100先说朴素版。性能基本在30到75 GFLOPS之间挣扎而且矩阵越大越接近60到70这个平台期——这是全局内存带宽在给它兜底无论怎么加算力都没有用。共享内存tile版提升到300到530 GFLOPS左右但因为寄存器复用不足、bank conflict等因素也没能把整张卡的火力释放出来。cuBLAS在512规模时就能跑到8230 GFLOPS到了4096已经达到25 TFLOPS左右接近这张显卡FP32理论峰值约35 TFLOPS的七成。注意这个“接近七成”已经是非常出色的成绩了实际GEMM能稳定达到理论峰值的60%到80%就算优秀因为还要考虑访存、指令发射、功率墙等限制。换算成直观倍数在4096规模下cuBLAS比朴素版快约330倍比共享内存tile版快约47倍。你手写一个月的优化成果可能还不如官方库在默认配置下跑出来的结果。4.3 cuBLAS为什么能快这么多很多人看到数据后第一反应是“NVIDIA是不是给自家库开了后门”。其实没有魔法cuBLAS快是因为它把GPU优化里能做的几乎全做了。首先是多级分块。cuBLAS采用类似CUTLASS的层次化分块策略在block层面切分大的输出tile在thread层面每个线程计算一个4×4甚至8×8的寄存器tile。这样做的好处是每个数据从全局内存读到寄存器之后可以被连续复用8到16次减少了大量冗余访存。其次是向量化加载。用float4一次读取4个float而不是一个float一个float地读。别小看这个改动它能把内存事务数量降低到原来的四分之一显著提升访存效率。共享内存方面cuBLAS会精心安排数据布局并做padding尽量规避bank conflict同时配合双缓冲让shared memory的加载和计算重叠隐藏访存延迟。最后是自动调优。cuBLAS在库内部维护了不同GPU架构、不同矩阵规模对应的kernel配置表。你传入矩阵尺寸后它会选一个合适的kernel变体包括block尺寸、unroll因子、流水线深度等。这一层自动调优是我们手写代码时最费时间的部分也是官方库价值最大的地方。另外cuBLAS还支持TF32、FP16等精度模式可以借助tensor core算得更快我们这节用的是纯FP32的SGEMM还没有动用这些“核武器”。后续篇章里等我们讲到tensor core和混合精度时你会看到性能差距还能进一步拉开。5. 踩坑实录与排查清单5.1 错误码返回来你却没检查cuBLAS的所有函数都会返回cublasStatus_t类型的状态码。一个常见的错误是开发者写完调用后直接往下走完全不检查返回值。这样一旦参数传错你看到的往往是莫名其妙的结果或者非法内存访问排查半天也定位不到源头。我习惯在调试阶段每个cuBLAS函数调用后面都套一个宏或者函数做检查例如#define CUBLAS_CHECK(call) \ do { \ cublasStatus_t st (call); \ if (st ! CUBLAS_STATUS_SUCCESS) { \ fprintf(stderr, cuBLAS error %d at %s:%d\n, \ st, __FILE__, __LINE__); \ exit(-1); \ } \ } while (0)上线前我会把性能敏感路径上的这个宏关掉但调试期绝不省略。最常见的错误码是CUBLAS_STATUS_INVALID_VALUE基本就是m/n/k、lda/ldb/ldc传了非法值其次是CUBLAS_STATUS_NOT_INITIALIZED十有八九是忘了调cublasCreate。5.2 结果全零或者数值不对结果全零优先检查三类问题。第一alpha/beta是不是传了0或者指针没赋值。第二矩阵维度m/n/k、lda/ldb/ldc是不是传反了尤其是在行主序转列主序时维度顺序很容易写错。第三kernel实际有没有被执行有没有调用cudaDeviceSynchronize确认执行完毕。如果你的结果跟CPU参考实现“对不上”但差距很小比如相对误差在1e-5到1e-6量级那通常是正常的。GPU上的浮点累加顺序和CPU不同最后几位不一样非常正常。别紧张这不算bug。真正要警惕的是相对误差到了1e-1这种量级那基本就是参数传错或者精度模式开错了。还有一个小坑如果你传beta0.0f希望C矩阵被覆盖但C里原本有NaN或者Inf某些实现里0×NaN还是会得到NaN。稳妥的做法是提前把C清零或者确保C的初始值是普通浮点数。5.3 性能上不去的检查思路如果你调了cuBLAS之后发现性能没有宣传的那么好不妨按下面顺序排查。先看GPU频率是否锁定。很多显卡在默认状态下不会一直跑在最高Boost频率特别是温度上来后会自动降频。测试前最好用nvidia-smi -lgc把频率锁在一个固定值或者至少记录一下测试时的实际频率。再看矩阵规模。cuBLAS对mini-batch场景的kernel启动开销很敏感如果你算的是64×64这种小矩阵性能数据可能不如手写专用kernel这在第2节已经说过。还有一个容易被忽略的点L2缓存策略。如果两个相邻的GEMM调用之间数据有复用可以尝试调整cudaFuncSetCacheConfig或者割让更多的L2给persistent访问。不过这个手段对cuBLAS来说属于“外部调参”实际效果因场景而异。如果你测试的是双精度或者复数数据类型请确保用的API函数名匹配比如cublasDgemm对应double、cublasCgemm对应cuComplex不然性能数字会非常奇怪。5.4 一句真心话我自己用cuBLAS做工程这几年最大的感受是真正干活的时候宁可多花十分钟把参数和主序问题理清楚也不要在踩坑之后再花一小时debug。尤其是行主序和lda这个问题几乎每个初用cuBLAS的人都会翻车一次。你在网上搜别人的代码很可能看到有人直接用CUBLAS_OP_T去乱试最后试出正确结果但又讲不出为什么——这不叫会用这叫撞运气。我个人的习惯是在正式写业务代码之前先单独写一个小demo把矩阵乘法用CPU实现一遍再用cuBLAS跑一遍对比结果。这个小demo只要几十行代码但能帮你把API参数、主序逻辑、错误检查全部验证通。经验不足的时候慢就是快。这个系列前面我们花了很大精力聊shared memory、block调度这些GPU底层机制现在把cuBLAS放进来对比看你会发现之前学的东西全都用得上——你终于能看懂那些性能数字背后的原因了。下一节我准备接着聊cuBLAS里tensor core相关的调用方式和混合精度那又是另一个世界。手头有显卡的朋友建议把今天这几个版本的代码都跑一遍亲眼看一看差距比我在这里说一百句都管用。