ARTICLE · INTELLIGENCE

战地情报 · 详情页

来自尧图项目组的一线实战观察与深度解析

CUDA SGEMM手写内核优化:从800到10600 GFLOPS的调优实践

CUDA SGEMM手写内核优化:从800到10600 GFLOPS的调优实践

前阵子我把一个SGEMM kernel从八百 GFLOPS 调到了一万零六百 GFLOPS。终端里打印出10623.47 GFLOPS的时候,我盯着屏幕停了几秒。这个数字在深度学习大模型遍地走的今天不算炸裂,Tensor Core 随便一开就是百 T 级别,但它是靠纯手工 SIMT 风格优化、用 float4 向量化、共享内存分块一步步怼上去的。对我来说,这比直接调 cuBLAS 爽多了。

先交代一下背景:我用的是一张 FP32 理论峰值在三四十 TFLOPS 级别的消费级 GPU,矩阵规模是 4096×4096,数据类型是 FP32。全程没有用 Tensor Core,没有用 cuBLAS,所有内核都是自己写的 CUDA C++,一步步从千行 profile 数据里抠性能。这篇文章不是教科书,是我把矩阵乘法从 800 GFLOPS 干到 10600 GFLOPS 的完整记录,代码能抄,坑也给你标出来了。

如果你在学 CUDA、或者正在为什么 SGEMM 跑不满发愁,看完这篇应该能少走不少弯路。

1. 10000 GFLOPS 到底是个什么水平

1.1 先算明白 GFLOPS 这个数是怎么来的

FLOPS 全称是 Floating-point Operations Per Second,GFLOPS 就是每秒十亿次浮点运算。10000 GFLOPS 等于 10 TFLOPS,也就是一秒内完成一万亿次浮点运算。这个单位在显卡参数页上特别常见,但很多人没搞清楚它到底怎么测出来的。

矩阵乘法的浮点运算量有固定公式:一次规模为 M×N×K 的矩阵乘,需要2 × M × N × K次浮点运算。为什么是 2?因为每算一个输出元素要执行一次乘法和一次加法。拿我的测试规模来说,M=N=K=4096,总运算量就是:

2 × 4096 × 4096 × 4096 = 137,438,953,472 ≈ 137.4 GFLOPS

注意单位换算,这个数是 GFLOP(总运算量),不是 GFLOPS(每秒速率)。速率需要用总运算量除以实际耗时。假如这 137.4 GFLOP 的任务在 1 秒内跑完,速率就是 137.4 GFLOPS;如果 13.7 毫秒跑完,速率就是 137.4 / 0.0137 ≈ 10030 GFLOPS,也就是我刚跨过的那条线。

我见过不少人直接在 kernel 结束后用clock()测时间,这个精度在毫秒级任务上不够看。正确做法是用 CUDA event,也就是cudaEventRecord包住 kernel 启动前后,再cudaEventElapsedTime拿到精确到微秒的耗时。另外还有个很多人忽略的点:测之前必须确认计算结果是对的。NVCC 在开启优化后有可能把没用的计算整段删掉,导致测出来的“性能”高得离谱。我自己的习惯是每版 kernel 跑完后做一次残差校验,最大误差超过 1e-4 就直接判定优化无效。

1.2 矩阵乘法为什么被当成性能标尺

这不只是因为矩阵乘法在数值计算里出现频率高,更重要的是它的算术强度(Arithmetic Intensity)极高。算术强度 = 总运算量 / 总访存量,它决定了一个计算任务在计算机里到底是被算力卡住还是被带宽卡住。

以 4096³ 的 FP32 矩阵乘为例,总数据量大约是三个矩阵加一起约 201 MB,总运算量 137.4 GFLOP,算下来算术强度在 680 FLOP/Byte 左右。就算 GPU 的全局内存带宽是 1 TB/s,要达到 10 TFLOPS 你也只需要大约 14.7 GB/s 的访存流量,带宽余量相当充裕。换句话说,这是个典型的“计算受限”问题——理论上只要把计算单元喂饱,性能就能拉满。

正因为它的计算密度极高,矩阵乘法成了检验一个并行体系结构“底线效率”的试金石。深度学习里的全连接层、卷积的矩阵展开、Transformer 的 QKV 投影和注意力打分,底层本质都是 GEMM。如果你能把 GEMM 写好,你在 AI 推理和训练的数据流优化上就成功了一大半。

2. 向量化不是玄学,是一笔经济账

2.1 FMA:一次乘加,两条流水线一起跑

CUDA 内核里最常见的一个计算操作是这种模式:

acc += a[k] * b[k];

这个表达式在 CPU 上会被拆成一条乘法和一条加法,但在 GPU 上,编译器会尽可能把它融合成一条 FFMA 指令。FFMA 的全称是 Fused Multiply-Add,它在一个指令周期内完成d = a * b + c的完整操作。

为什么这很重要?因为 GPU 的 FP32 吞吐量通常就是按 FMA 来标定的。一张声称“35 TFLOPS”的卡,真实含义是它每秒能执行 17.5 万亿次 FFMA,因为每次 FMA 包含两次浮点运算,所以对外宣传按 2 倍来算。如果你的代码里乘法和加法没有成功融合,实际吞吐会直接掉一半。

怎么确认融合发生了?把编译后的 SASS 拿出来看就行。编译时加-keep选项,或者用nvdisasm打开生成的 cubin 文件,搜FFMA字样。如果看到大量FFMA而不是单独的FMUL和FADD,说明编译器干得不错。我早期一版 kernel 性能一直上不去,查了半天发现某些分支路径里编译器没有做 FMA 融合,改写代码结构之后立刻好了。

2.2 float4 的账本:指令少四倍,带宽不白费

做 CUDA 优化的人一提到向量化,第一个想到的就是float4。这玩意本质是 CUDA 内置的一种向量类型,内含四个 float,内存中连续排列,要求 16 字节对齐。

为什么它能提速?要点有三个。

第一是减少指令数量。同一个线程读 4 个 float,用标量方式要发 4 条 LDG 指令,用float4只要 1 条 LDG.128。指令发射数量少了,计算流水线就有更多slot留给 FFMA。

第二是充分利用缓存行和访存带宽。GPU 的全局内存访问以 128 字节为一个事务,一个 warp 有 32 个线程,如果每个线程按 float4 读取,正好是 32 × 16 = 512 字节,也就是 4 个整齐的内存事务;如果每个线程只读一个 float,则需要 32 个线程共同拼凑才能填满 cache line。虽然连续标量访问在现代 GPU 上也能合并,但 float4 让合并变得更整齐、更容易命中问题的本质。

第三是流水线压力显著降低。在矩阵乘法的内层循环里,如果每个线程每次只取一个 float,那么为了凑够足够的计算量,内存流水线会被频繁塞满请求。换成 float4 之后,同样数量的数据只需要四分之一的内存请求次数,memory pipeline 的排队延迟随之下降。

我习惯把 float4 想象成快递:一个地址同时到了四个包裹,你非得分四次下楼去拿,这不是折腾自己吗?一次拿上来,后面处理就有余力了。

2.3 对齐是个隐形门槛,不注意直接崩

用float4最烦的一点是它要求 16 字节对齐。malloc和cudaMalloc分配的内存默认对齐到 256 字节,一般没问题。但当你用指针偏移去访问一个二维数组的某一行时,就需要小心了。

比如你有一个float* A,行数是 K,行主序布局下访问A[row * K + col]。如果你想把这个地址当成float4*用,就必须保证(row * K + col) * 4字节能被 16 整除,也就是row * K + col是 4 的倍数。如果 K 不是 4 的倍数,那恭喜你,某些行会天然错位。

解决方式一般是两种:一是在分配缓冲区时按 K 向上取整到 4 的倍数(也就是 padding),二是干脆把矩阵拷贝到一个对齐的新缓冲区里。千万别硬着头皮 reinterpret_cast,未对齐的 float4 访问在 CUDA 里会直接导致非法地址错误,而且这种错误还不容易定位。我后来在代码里养成了一个习惯,凡是涉及 float4 的地方,先写一个断言:

assert(reinterpret_cast<uintptr_t>(ptr) % 16 == 0);

别小看这一行断言,它帮我抓到了无数潜在的边界 bug。

3. 从 800 GFLOPS 到 10000 GFLOPS 的一条可复现路线

3.1 基线版本:列主序 B 矩阵让访存模式崩了

我的出发点是一个“最朴素”的 kernel:每个线程算一个输出元素,三重循环直接怼,没有任何分块。由于我的数据接口里 B 矩阵是列主序存放的,也就是B[k + col * K]这种索引方式,问题一下子就来了。

在同一 warp 里,线程的col是连续的,但访问 B 的时候地址步长是K × 4字节。这意味着 32 个线程访问的地址彼此相距巨大,根本无法合并。GPU 的全局内存系统为每个线程单独发起了内存事务,一个 warp 的数据请求被拆成了一堆碎片,带宽利用率惨不忍睹。

这版代码的性能惨到什么程度呢?只有大约 800 GFLOPS。虽然这个数字对很多 CPU 来说已经是天文数字,但是对 GPU 来说是在暴殄天物。它浪费的不是算力,而是内存系统把所有时间都花在搬运没用的数据上了。

这版失败的 kernel 给了我很深一个教训:性能优化永远先看访存模式,再看计算模式。一个错误的数据布局能让千万亿次算力原地趴窝。

3.2 第一步优化:先改布局,别急着上高深技巧

面对这个惨淡的基线,我做的第一件事不是引入共享内存,也不是上 float4,而是先把 B 矩阵的存放方式改成行主序。

改写之后,最内层循环里线程访问 B 的方式变成了B[kk * N + col],同一 warp 的 32 个线程地址连续,全局内存访问立刻重新合并。就这一下,性能从 800 GFLOPS 飙到了约 2600 GFLOPS,三倍多的收益,代码逻辑几乎没变。

这一步的意义在于把问题分成两层:数据布局问题和计算组织问题。数据布局不合理的时候,后面所有优化都建立在沙子上。就像做饭之前先确认食材是新鲜的,否则刀工再好也白搭。

顺带说一句,如果你没法改调用方传进来的 B 矩阵布局,还有一个常见做法是在 kernel 外面用一个cudaMemcpy2D把列主序矩阵转成行主序,或者专门写一个转置 kernel。转置本身也有开销,但一次转置能换来几十倍的计算加速,这笔账永远划算。

3.3 第二步优化:共享内存 Tiling + float4 向量化

数据布局理顺以后,下一步要解决的是“数据复用”的问题。朴素 kernel 里每个线程直接访问全局内存,它算完一个输出就不再复用任何数据。但实际上矩阵乘法的数据重用机会极大:A 矩阵的一行可以被多个线程复用,B 矩阵的一列也是。

所以我引入了共享内存分块,把 A 和 B 按小块先搬到共享内存里,再由所有线程从共享内存读取计算。共享内存在 GPU 片内,延迟比全局内存低一个数量级,而且带宽高得多。所谓 tiling,就是每次只处理一个 64×64 的输入块,把它切成小片搬进共享内存,算完再换下一片。

同时我把搬数据的过程向量化了,一次float4拷贝四个 float,减少了加载指令数量。大致的 kernel 骨架如下:

#define BLOCK_M 64 #define BLOCK_N 64 #define BLOCK_K 16 __global__ void sgemm_tile_v2(const float* A, const float* B, float* C, int M, int N, int K, float alpha, float beta) { __shared__ float As[BLOCK_M][BLOCK_K]; __shared__ float Bs[BLOCK_K][BLOCK_N]; int tx = threadIdx.x; // 0..63 int ty = threadIdx.y; // 0..63 int rowBase = blockIdx.y * BLOCK_M; int colBase = blockIdx.x * BLOCK_N; float acc = 0.0f; for (int k0 = 0; k0 < K; k0 += BLOCK_K) { __syncthreads(); // 向量化加载 A 和 B 的子块到共享内存 float4 aVec = *reinterpret_cast<const float4*>(&A[rowBase + ty][k0 + tx * 4]); float4 bVec = *reinterpret_cast<const float4*>(&B[k0 + ty * 4][colBase + tx]); // 这里需要保证 K 和索引按 4 对齐,示意代码省略边界处理 // ... __syncthreads(); #pragma unroll for (int kk = 0; kk < BLOCK_K; ++kk) { acc += As[ty][kk] * Bs[kk][tx]; } } // 写回 C,省略 beta 处理 C[(rowBase + ty) * N + (colBase + tx)] = acc; }

这个版本跑出来约 5300 GFLOPS。注意我全程用__syncthreads()控制共享内存读写顺序,这个 barrier 放错位置会导致数据竞争,计算结果直接花掉。我在调试这个版本的时候就吃过亏——一个__syncthreads()漏放,结果矩阵右上角总有随机值,排查了整整一下午。

这一步的收益来自两方面:一是全局内存访问次数大幅减少,因为数据从全局内存搬进共享内存之后可以被 64×64 个线程反复复用;二是 float4 让搬数据本身更高效了。

3.4 第三步优化:寄存器微块、循环展开与 FMA 指令密度

到了 5300 GFLOPS,计算瓶颈开始转移。用nsight compute一查,issue active已经比较高了,说明指令发射速率开始饱和。这时候再用“一个线程算一个元素”的结构就不行了——每个 FFMA 都对应着两次共享内存读取,共享内存 load 指令占据了太多发射槽位。

解法是寄存器微块。具体来说,让每个线程算一个 4×4 或者 8×8 的微块,也就是把多个输出元素放在寄存器里。这样每次从共享内存取一个 A 元素和一个 B 元素,可以同时更新多个累加器,大幅提高数据复用率。

比如一个线程算 4×4 的微块,它从共享内存读一个 A 向量和 B 向量到寄存器,就能更新 16 个累加器。共享内存的 load 指令数量降到原来的八分之一,FFMA 的密度也上来了。

float acc[4][4] = {}; float aReg[4], bReg[4]; #pragma unroll for (int kk = 0; kk < BLOCK_K; ++kk) { // 从共享内存按 float4 加载 *reinterpret_cast<float4*>(aReg) = *reinterpret_cast<const float4*>(&As[ty][kk * 4]); *reinterpret_cast<float4*>(bReg) = *reinterpret_cast<const float4*>(&Bs[kk * 4][tx]); #pragma unroll for (int i = 0; i < 4; ++i) { #pragma unroll for (int j = 0; j < 4; ++j) { acc[i][j] += aReg[i] * bReg[j]; } } }

寄存器微块还有一个额外好处:多个独立累加器增加了指令级并行度。FFMA 虽然吞吐高,但它有大约四五个周期的延迟。如果只有一个累加器,每步计算都得等上一步结果回来;有 16 个独立累加器的话,流水线可以同时塞满多个 FFMA,延迟被彻底隐藏。

这版性能达到了约 8800 GFLOPS。此时计算和访存的平衡已经很接近上限了,剩下的空间更像是在和流水线细节搏斗。

3.5 第四步优化:双缓冲与异步拷贝,跨过万级门槛

最后这 1800 GFLOPS 的提升是最费劲的。瓶颈从“计算密度”变成了“等待共享内存填充”。虽然 tiling 已经把全局内存流量压低了很多,但是每次计算新的 K 分块时,线程还是得停下来等__syncthreads()之后共享内存的数据就绪。

解决办法是双缓冲。核心思想是:在计算当前分块的同时,让下一个分块的数据提前从全局内存搬到共享内存。这样计算和访存重叠,把隐藏在同步点后的空闲时间全利用起来。

现代 CUDA 架构提供了cp.async指令,也就是异步拷贝。在 Ampere 及之后架构上,我直接用__pipeline_memcpy_async把全局内存的数据异步搬到共享内存,手动控制__pipeline_commit和__pipeline_wait_prior。

__shared__ float As[2][BLOCK_M][BLOCK_K]; __shared__ float Bs[2][BLOCK_K][BLOCK_N]; // 预取第0块 __pipeline_memcpy_async(&As[0][ty][tx * 4], &A[...], 16); __pipeline_commit(); for (int k0 = 0; k0 < K; k0 += BLOCK_K) { int cur = (k0 / BLOCK_K) & 1; int next = cur ^ 1; __pipeline_wait_prior(0); __syncthreads(); // 预取下一块 if (k0 + BLOCK_K < K) { __pipeline_memcpy_async(&As[next][ty][tx * 4], &A[...], 16); __pipeline_commit(); } // 使用当前块计算 ... __syncthreads(); }

这一步跑完,成绩实打实地站上了 10600 GFLOPS。相比最初的 800 GFLOPS,整体提升了十三倍。而且快跑进 4096³ 单次 13 毫秒以内了,这个规模下已经是相当可观的手工优化水准。

3.6 关键版本的性能与瓶颈变化记录

这是整个调优过程中的各阶段数据,每版 kernel 都做了正确性验证,误差控制均在 1e-4 以内。

版本关键优化点GFLOPS相对基线提升主要瓶颈
v0 基线无,列主序 B 直接硬算8001x非合并访存,带宽浪费
v1 布局调整B 转行主序,消除非合并访存26003.25x数据复用不足,计算密度低
v2 共享内存分块Tiling + float4 向量化搬数据53006.6x共享内存 load 指令过多
v3 寄存器微块每线程 4×4 累加器,FMA 指令密度提升880011x同步等待,访存未与计算重叠
v4 双缓冲异步拷贝cp.async 预取下一分块1060013.3x逼近纯手工 SIMT 优化天花板

NSight Compute 的数据也能看出瓶颈的变迁:v0 是stall_long_scoreboard(等待全局内存回包)占绝对主导;v2 变成了stall_short_scoreboard(等待共享内存);v3 是stall_barrier(同步等待);到 v4,各种 stall 占比变得非常均匀,说明已经很难再从简单手段里榨出性能了。

4. CUDA 环境与向量化实践的高频坑

4.1 换 CUDA 版本前先想清楚三件事

做 CUDA 优化绕不开环境配置。我见过太多人一上来就装最新版 CUDA,结果项目编不过、运行报错,最后发现是版本匹配问题。

决定能不能用某个 CUDA Toolkit 版本,核心有三件事:显卡的计算能力(Compute Capability)、驱动支持的 CUDA 版本、Toolkit 自身的版本号。

先说驱动。nvidia-smi右上角会显示一行CUDA Version: 13.0之类的信息,这代表当前驱动能支持的最高 CUDA 运行时版本。驱动是向后兼容的,只要你的 Toolkit 版本不高于这个数字,通常就能跑。反过来,驱动太老、Toolkit 太新,运行时就会报找不到 CUDA 入口的错误。

再说计算能力。每个 NVIDIA GPU 都有一个算力代号,比如 sm_86(3060Ti/3080)、sm_89(4060/4070)、sm_120(5070)。编译 kernel 时-arch=sm_89里的数字必须能被 Toolkit 识别。如果你的 GPU 是 5070(sm_120),而 CUDA Toolkit 还是 12.0,那么即使驱动支持,编译也会报Unsupported gpu architecture 'compute_120'。这种情况不是你的代码有问题,是 Toolkit 太老,必须升级到能识别 sm_120 的版本。

最后是 Toolkit 到底装哪个版本。我的建议是:别追新,先看你的深度学习框架、第三方库的要求。PyTorch 官方轮子自带 CUDA runtime,对系统级的 Toolkit 其实依赖不强;但如果你要自己编译 cuDNN 相关代码,就需要配套版本。

4.2 nvcc -V 和 nvidia-smi 显示不一致不是 Bug

这是个极其常见的迷惑点。nvcc -V显示的是你安装的 CUDA Toolkit 版本,比如 11.8;而nvidia-smi右上角显示的是驱动支持的最高 CUDA 版本,比如 13.4。两个数字不一样,太正常了。

如果你没在nvcc在 PATH 里配置好,更会出现这种情况:编译时用的是系统自带的旧版 Toolkit,运行时驱动却能支持新版本。我见过有同学因此把驱动卸了重装,折腾了半天结果根本不是问题。确认编译环境时用which nvcc和nvcc -V,确认运行环境时用nvidia-smi,两条链别混在一起。

4.3 几个高频报错的排查思路

我把平时遇到的典型报错整理成了一个速查表,每一项都是自己或者身边的人踩过的坑:

报错特征常见原因处理办法
nvcc fatal: Unsupported gpu architecture 'compute_120'Toolkit 版本过老,不认识新算力升级到支持对应 sm_xx 的 Toolkit 版本
/usr/local/cuda/include/cuda_runtime.h: No such fileCUDA_HOME没设置,或头文件路径不对检查CUDA_HOME是否指向实际安装目录,比如/usr/local/cuda-12.4
libcudart.so.12: cannot open shared object file运行时找不到 CUDA 动态库把/usr/local/cuda/lib64加进LD_LIBRARY_PATH,并跑一次ldconfig
PyTorch 启动时UserWarning: GPU is not compatible with the PyTorch versionPyTorch 自带的 CUDA runtime 不支持新显卡算力升级 PyTorch 到更新的 CUDA 编译轮子,而不是去重装系统 CUDA
CUDA_MALLOC一直失败或者 allocation 为 0显存被占满,或环境变量CUDA_CACHE_DISABLE被设置查nvidia-smi显存占用,清掉僵尸进程;确认缓存目录可写
编译gpu_burn这类项目时,报错指向/usr/local/cuda/include/cuda找不到编译脚本里的 CUDA 路径和你实际安装路径不一致设置CUDA_HOME或修改 Makefile 里的 CUDA 路径变量
OpenCV 装了opencv-python但cv2.cuda模块不存在官方预编译 wheel 默认不启用 CUDA 支持需要自己从源码构建带 CUDA 的 OpenCV,或使用第三方预编译包并核对版本搭配

另外提醒一个反直觉的坑:很多人以为装了 CUDA Toolkit 就等于装了显卡驱动,其实 Toolkit 只是编译器加运行时库,不包含驱动。驱动必须单独从 NVIDIA 官网装,而且先装驱动再装 Toolkit 是更稳妥的顺序。

5. 冲过万级 GFLOPS 之后的几点心得

如果你也想冲 10000 GFLOPS 这条线,我给的最实际建议是:别一上来就抄 CUTLASS 或者 cuBLAS 的源码。那些库为了适配所有架构做了大量模板抽象,直接读会一头雾水。先用手写 kernel 把一个简单的 SGEMM 优化链路完整走一遍,再回头看那些工业级代码,你会发现处处都能对上号。

我在这轮优化里最大的体会是:性能瓶颈会随着优化不断移动。最初是访存模式,然后是数据复用,接着是指令密度,最后是流水线等待。每个阶段都要用 profiler 重新确认瓶颈在哪,不能靠猜。NSight Compute 里的 stall 原因分布是非常可靠的诊断指标,比肉眼猜性能强太多了。

还有一个小技巧:在做性能对比时,每个优化只改一个变量。我 v1 到 v2 的版本只加了共享内存分块,v2 到 v3 只加了寄存器微块,v3 到 v4 只引入了异步拷贝。如果一次性把所有优化都堆上去,出了问题你根本不知道是哪一步导致的。控制变量法在进行性能优化时同样成立。

这次跑过 10600 GFLOPS 之后,我又顺手试了一下把分块尺寸从 64×64 调成 128×128,结果不升反降,因为共享内存占用太大导致 occupancy 掉得厉害。想冲更高性能,下一步还得从 Tensor Core 和更精细的 warp 调度入手。但那是另一条路的故事了。

RELATED READING

延伸阅读

更多一线实战笔记与深度复盘,助您持续精进