资讯动态

手写CUDA SGEMM从200到2600 GFLOPS:Warp Tiling与共享内存优化实录

发布时间:2026/10/26 13:05:53 来源:尧图企业网站定制
我最初以为“矩阵乘法跑赢 cuBLAS”是那种只有专家才能碰的题目直到自己花了两周时间把一块消费级显卡上的 SGEMM 从 200 GFLOPS 一路优化到 cuBLAS 的 83%我才意识到这件事离普通 CUDA 开发者并没有那么远。这个优化过程的核心不是玄学而是 Warp Tiling、共享内存布局和异步拷贝这三件事。这篇文章就是把我的完整实录整理出来从为什么矩阵乘法值得手写到每一版 kernel 的改法再到最后怎么用 profiler 找瓶颈全部说清楚。适合正在学 CUDA、准备高性能计算方向或者面试被问到“你怎么优化 GEMM”的朋友。1. 怎么想到要和 cuBLAS 比1.1 矩阵乘法是显卡的试金石矩阵乘法也就是 C A × B是所有基础线性代数里最有代表性的一道题。它的计算量一眼就能算出来如果 A 是 M×KB 是 K×N那么 C 是 M×N总计算量是 2·M·N·K 次浮点运算。以 4096×4096×4096 为例一次全尺寸 SGEMM 就是 2×4096³ 137.4 GFLOP 的运算量。这个数字很大但 GPU 本身擅长做并行乘加所以它成了检验一块显卡真实算力的标准项目。很多面试官喜欢问矩阵乘法并不是因为它简单而是因为它能同时考察线程编排、存储层次、访存模式、指令级并行和性能分析。你把你自己的 kernel 和官方库放在同一块 GPU 上跑差距就是你和 GPU 优化高手之间的距离。我第一版 naive kernel 大概只有 0.2 TFLOPS那时候和 cuBLAS 的差距是十倍还多没有任何讨论的空间。所以我把矩阵乘法当作“硬碰硬”的试金石不用张量核心不用 FP16就用最普通的 FP32 乘加指令看纯手写代码到底能贴近官方库到什么程度。这个目标听起来简单但做起来要拆的东西非常多。1.2 cuBLAS 到底赢在哪一开始我也天真地以为 cuBLAS 只是“官方库所以优化得多”。真正看下去才发现它赢在几个层面。第一是分块策略极其丰富同一个尺寸它会测试多种 block tile、多种 BK、多种调度方式甚至对不同的 K 段采用不同的 kernel。第二是共享内存和寄存器都用得非常狠通常每个线程负责一个 8×8 甚至更大的输出子矩阵而不是每个线程只算一个输出点。第三是它有非常好的软件流水线能通过预取隐藏访存延迟让访存单元和计算单元尽量同时工作。如果不看源码你会觉得这些都很神秘。但它们本质上就是我们手写 kernel 时做的那些事只不过官方库把这些事情做到了极致并且用了大量 profiling 数据去调参。拿 FP32 SGEMM 来说官方库在没有张量核心参与的路径上也能把一块高带宽显卡的浮点利用率推到很高的水平。所以“到达 cuBLAS 的 83%”实际上意味着我已经把手工 GEMM 做得像模像样了剩下的差距更多来自参数枚举、汇编级微调这类体力活。2. 设计一个 Warp Tiling 的 SGEMM2.1 从线程、Warp 到 Block 的三层映射矩阵乘法并行化的第一步是把输出矩阵 C 切成很多小块让不同的线程块去算不同的小块。我选的 block tile 是 128×64也就是一个线程块负责 C 上的 128 行、64 列范围。整个线程块一共有 128 个线程正好 4 个 warp。网格划分就是 (N/64, M/128)如果有 4096×4096 的输出矩阵就会生成 32×64 2048 个 block每个 block 独立跑。关键的一步是“Warp Tiling”不要把线程块里的任务直接切成 128 个独立线程去算而是先用 4 个 warp 去切分行方向每个 warp 负责 32×64 的输出区域。再往下一层每个 warp 里的 32 个 lane 又分成 4 组每组 8 个 lane每个 lane 负责一个 8×8 的小块。这样一层层套下来最终每个线程算 64 个输出值。这样做的好处非常明显。每个线程的 64 个结果都在寄存器里复用同一个 A 值要和 8 个不同的 B 值做 8 次乘加同一个 B 值也要和 8 个不同的 A 值做 8 次乘加。和“每个线程只算一个输出”的朴素做法相比数据从寄存器读一遍就被复用了很多次共享内存和全局内存的访存次数都大幅下降。GPU 最怕的不是计算而是数据搬来搬去Warp Tiling 本质上就是把“搬数据”这件事尽量挡在寄存器外面。2.2 共享内存里的轨道布局有了线程映射还要解决数据供应的问题。全局内存虽然大但延迟很高所以矩阵 A 和 B 的块要先进共享内存。我选的内层分块大小是 BK16也就是循环里每次取 A 的 128×16 一小块取 B 的 16×64 一小块。共享内存里安排两块缓冲区__shared__ float As[2][BK][TILE_M 1]; __shared__ float Bs[2][BK][TILE_N 1];这里 TILE_M128TILE_N64加 1 是为了避免 bank conflict。共享存储器一共有 32 个 bank每隔 4 字节换一个 bank。如果行没有任何 padding那同一行相邻元素刚好落在相邻 bank 上看起来没问题但当你按行方向取不同数据时可能多个线程访问同一个 bank造成串行化。加一个 float 的 padding 后行的步长从 128 变成 129129 对 32 取余是 1这样连续行的起始 bank 位置就不一样了冲突概率会低很多。共享内存大小也顺便算一下As 是 16×129 个 floatBs 是 16×65 个 float一个 stage 加起来是 3104 个 float约 12 KB。双缓冲就是约 24 KB完全能塞进常见的 48 KB 静态共享内存范围。这个尺寸设计不是随便拍的它让一个 SM 上还能放好几个 block保证足够的并行度去掩盖延迟。2.3 双缓冲与异步拷贝矩阵乘法循环的长度是 K/BK 次每次都要把新的 A、B 小块从全局内存搬进共享内存。最笨的办法是“先搬完再算再搬下一块”但这会让计算单元在等数据时发呆。我的做法是双缓冲加异步拷贝算第 i 块的同时用 cp.async 把第 i1 块从全局内存预取到另一块共享内存中。cp.async 是 Ampere 之后架构提供的异步拷贝指令它和普通 global load 最大的区别是数据从全局内存到共享内存的过程不需要经过寄存器也不会把当前线程的计算流卡住。你可以把它想象成在厨房做菜时水龙头已经在另一个水池里放水而不是你站在水池边等水满了才开始切菜。这里有个容易搞错的点双缓冲不只是把数组加一个[2]就完事了。同步逻辑必须处理好。我用的是两阶段 pipeline每个迭代先把下一 stage 的拷贝发出去再等待当前 stage 的数据就绪然后用__syncthreads()保证整个 block 里的线程都看到这一 stage。如果同步写错轻则计算结果错重则死锁。后面实操部分我会把伪代码写出来。3. 实操落地内核代码与关键参数选择3.1 第一版先保证正确我没有一上来就写花里胡哨的版本而是先写了一个最简单的共享内存版本每个线程只算 C 的一个点。线程块仍然是 128×64每个线程加载 A 的一个值和 B 的一个值乘加后写回 C。它的优点是正确性极好验证一旦跑出来的结果和 cuBLAS 一致就说明分块逻辑、索引换算、循环边界都没问题。这个版本性能当然很惨。因为每个输出点只做一次乘加却要从共享内存读两次数据访存开销远大于计算。我实测同一块卡上大约只能跑到 0.8 TFLOPS 左右naive 版本就更不用说了。不过这个版本是后面所有优化的“正确性底座”后面每次改动都会和它做同一份随机矩阵的 diff。很多人会跳过这一步直接上 register tiling结果出了 bug 也不好定位。我的建议是第一版哪怕再慢也要写完整因为矩阵乘法调索引本身就很烦如果一开始就混合了 tile 映射、双缓冲、cp.async任何一个小错都会让你排查到崩溃。3.2 第二版把一个 Warp 当成计算单元第二版开始引入真正的 Warp Tiling。我给每个线程分配 8×8 的输出矩阵用 8 个寄存器存 A 的 8 个值8 个寄存器存 B 的 8 个值64 个寄存器存 C 的输出。这 64 个寄存器是性能的关键因为它们在最内层循环里反复被读和累加。内层循环的代码逻辑大概是for (int k 0; k BK; k) { float a[TM]; // 从 As 读 8 个 A 值 float b[TN]; // 从 Bs 读 8 个 B 值 for (int i 0; i TM; i) for (int j 0; j TN; j) c[i][j] a[i] * b[j]; }每读一组 a 和 b就做 8×864 次乘加。在这 64 次乘加中a 和 b 都在寄存器里不再产生新的访存。这就是“寄存器级数据复用”的威力。第二版只是把朴素版本改成这个结构性能就立刻翻了好几倍实测能到 1.6 TFLOPS 左右。但这时候还没有加双缓冲每次迭代都要等共享内存加载完才能继续算卡在访存延迟上的时间很多。3.3 第三版铺平内存访问与控制 Bank Conflict第三版要做的事情很具体让共享内存的访问尽量不冲突同时用向量化 load 把带宽打满。先说 bank conflict。共享内存一次周期只能服务若干个 bank如果同一 warp 里的 32 个线程恰好访问了同一个 bank 上的不同地址硬件就会把它们拆成多次访问性能直接按倍数下降。我在 2.2 里说的 padding 就是为这个服务的。具体到我的布局里As 的行数是 128加一个 padding 后变成 129B 同理。这样在访问连续行时不会所有线程都撞到同一组 bank 上。还有一点容易被忽略同一 warp 里如果有好几个线程在读同一个共享内存地址硬件是可以做广播的不产生冲突。所以在写索引的时候我会刻意让某些线程组访问同一个地址减少真正被访问的“地址数”。另外我还在加载阶段用 float4 一次性读 4 个 float。对齐好的 float4 可以把 load 指令数降到原来的四分之一指令发射压力小很多。代价是 padding 和对齐要小心如果行宽不是 4 的倍数float4 就会错位。实际我选择让每一行都按 4 的倍数起步宁可浪费一点点空间也要保住向量化。3.4 第四版预取数据把等待藏起来第四版加入异步拷贝和双缓冲。这一步让性能第一次有了“跑起来很爽”的感觉。原先每个 K 循环迭代都是“等数据 → 计算 → 等数据”现在变成“一边算第 k 块一边加载第 k1 块”。计算单元大多数时候都在做乘加而不是空等。代码层级的伪代码如下// 预取第一块 cp_async(As[0], A offset0, ...); cp_async(Bs[0], B offset0, ...); cp_async_commit(); for (int k0 0; k0 K; k0 BK) { int cur (k0 / BK) 1; int next cur ^ 1; if (k0 BK K) { cp_async(As[next], A offset_next, ...); cp_async(Bs[next], B offset_next, ...); } cp_async_commit(); cp_async_wait1(); __syncthreads(); // 用 As[cur], Bs[cur] 计算这一个 128×64 小块 compute_warp_tile(As[cur], Bs[cur], ...); __syncthreads(); }这里cp_async_wait1表示允许还在路上的最新一组异步拷贝未完成但更早的那组必须完成。这正好保证当前要用的 cur 块已经就绪。第二次__syncthreads()是防止某个 warp 提前进入下一轮把别人还在用的共享内存覆盖掉。加了这版之后我在 4096³ 的 FP32 用例上跑到了大约 2.6 TFLOPS离这块卡上 cuBLAS 的 3.15 TFLOPS 已经不远了。83% 的成绩基本就是这个版本测出来的。4. 性能实测83% 是怎么来的4.1 跑分环境与口径先说跑分环境。我是在 WSL2 里配的 CUDA 环境编译用 nvcc运行时用同一个 CUDA 工具包所以没有出现 Windows 驱动和 Linux 工具链打架的问题。看性能之前建议先确认两件事nvidia-smi里显示的驱动版本以及nvcc --version或者cudart的版本这两个版本不必完全一致但你的代码编译和运行要使用同一套 runtime。矩阵尺寸选的是常规的 MNK4096数据类型为 FP32。性能计算方式是TFLOPS 2 * M * N * K / 时间(秒) / 1e12我用同一尺寸跑 100 次取中位数避免频率波动带来的干扰。对正确性我还会和 cuBLAS 的结果做相对误差比较绝对误差控制在 1e-3 以内不然优化半天跑出错误结果就白忙了。4.2 几个版本的数据对比我把自己手上几个版本的实测数据整理了一下一张表看下来就清楚每一步优化到底值多少钱版本关键做法性能(GFLOPS)相对 cuBLASnaive直接三重循环每线程 1 个输出1805.7%共享内存版分块 每线程 1 个输出82026%Warp Tiling每线程 8×8 输出寄存器复用165052%Warp Tiling padding消除主要 bank conflict205065%加法缓冲与 cp.async异步预取隐藏访存延迟260082.5%cuBLAS官方库全状态调优3150100%注意这些数字都是同一个运行环境下测的不同显卡、不同尺寸、不同架构结果会不一样。比如矩阵尺寸换成 1024³cuBLAS 的优势依然存在但比例可能有变化换成非对齐的 1000 之类尺寸官方库还会表现出更大的调参优势。所以别把 83% 当成一个固定答案它更多是告诉你可以做到“很接近”。4.3 用分析器找瓶颈跑完分之后我并没有因为“看起来还行”就停手而是用分析器把每个版本的瓶颈找出来。这里最常用的工具是 ncu。我会重点看几个指标共享内存 bank conflict 数量、寄存器溢出情况、SM 上 warp 的平均占用率以及访存等待占比。第一版 naive 的问题很容易看很多周期都花在 memory 等待上SM 里的计算单元利用率很低。Warp Tiling 版本之后计算占比明显上升。但我也遇到过很诡异的现象代码看起来应该很快实际帧数却突然掉一半最后用 profiler 一查发现寄存器溢出大量局部变量被压到本地内存里性能直接崩塌。另一个值得关注的指标是stall类型。如果大量 warp 都是long_scoreboard这种等待全局访存返回的停顿说明预取还不够如果barrier停顿占比高说明__syncthreads()太频繁或者调度粒度不对。这些指标读起来稍微有点门槛但绝对是手写 kernel 的“仪表盘”比瞎猜靠谱得多。5. 常见问题与排查心得5.1 共享内存 Bank Conflict 的“隐形降速”共享内存 bank conflict 是最容易隐蔽的问题。它的可怕之处在于程序不会报错结果也正确就是性能提不上去。我第一次从 1650 涨到 2050 的跨越就是靠修 bank conflict 拿到的。怎么排查如果你不想装 profiler可以先手工推索引假设 warp 里有 32 个线程同时访问共享内存地址addr(tid)看addr % 32的分布重复越多冲突越重。最典型的错误是行宽恰好是 32 的整数倍导致同一列上的线程全落在同一个 bank。padding 到“和 32 互质”是成本最低的修法。不过也要注意加了 padding 之后共享内存容量会稍微变大。如果 BK 很大、双缓冲又开着可能超过静态共享内存限制。这时可以用动态共享内存或者在 launch 时设置更大共享内存配置。我一般先看 profiler 的l1tex__data_bank_conflicts_pipe_lsu_mem_shared_op_ld.sum那个值如果成百上千肯定要先修这里。5.2 寄存器溢出性能直接从山峰落回盆地寄存器溢出是我这次踩过最大的坑。每个线程 8×8 的输出意味着至少 64 个寄存器要用来保存 C 的值再加上 A、B、地址、循环变量很容易逼近 128 甚至 255 个寄存器上限。寄存器是有限的如果被压到本地内存所有看似漂亮的优化都会瞬间作废。我当时的症状是加上 cp.async 之后性能不升反降。排查发现为了写双缓冲逻辑我引入了太多临时变量编译器又使用了比较激进的展开策略把寄存器用爆了。解决方法有三个方向一是减少线程块的 tile让每线程输出从 8×8 降到 4×8二是把某些地址计算放到循环外复用三是用__launch_bounds__(128, 4)这类约束限制编译器不要放飞自我必要时-maxrregcount也要会用。如果你想看寄存器到底用了几路编译时加-Xptxas -v它会直接打印每个 kernel 的寄存器数、shared memory 数和 spill 情况。我看到spill stores和spill loads不为 0 的时候就会回去调整。5.3 CUDA 多版本、WSL2 环境带来的坑我在 WSL2 里做测试时确实也遇到过 CUDA 版本不匹配的问题。最常见的是机器上装了多个 CUDA 工具包软链指到一个版本但某些第三方库或者 Python 包是拿另一个版本编译的跑起来就报版本不一致错误。遇到这种问题不要急着重装驱动先确认你的PATH和LD_LIBRARY_PATH指向哪个版本再用nvcc --version和实际运行的cudart对一下。另外很多人装 CUDA 是为了编译 OpenCV 或者跑深度学习框架。这些库对 CUDA 版本和配套库版本有严格绑定比如某个框架版本可能要求特定的 CUDA 和配套数学库版本。我的经验是能不动环境就不动环境项目尽量用容器或者独立工具链隔离版本。WSL2 的好处是可以用 Windows 侧驱动Linux 侧只装 CUDA 工具包省去了非常多驱动兼容问题但坏处是如果你从网上复制安装命令很容易装出一堆半吊子环境。5.4 别把线程数和性能划等号最后一类“问题”其实来自思维惯性。很多新手觉得GPU 优化就是堆线程数线程越多越快。我在实测里验证过当共享内存和寄存器资源被占满后再多开线程反而会因为 block 调度不过来、资源排队严重而变慢。正确的思路是“让每个线程做更多有用功”而不是“让更多线程各做一点点”。把 C 矩阵拆成 128×64 的块每线程 8×8 输出已经让总共 4096 个 block 在 2048 个线程的 GPU 上有了足够的调度余地。真正决定上限的是每个 SM 上的寄存器、共享内存和调度带宽。我最终停留在这个配置不是因为不能把 tile 调大而是再调大就撞到寄存器墙性价比不高了。6. 关于 Warp Tiling 的延续方向如果你也想把这套思路继续往前推有几个方向是可以再走的。一是把每线程 tile 从 8×8 改成 16×8 甚至 16×16进一步加大寄存器复用但同时对共享内存 load 的向量化要求会更高。二是引入更复杂的 swizzle 模式在共享内存索引上做 XOR 操作让 float4 向量化 load 时也能规避 bank conflict这比 padding 更省空间。三是切分 K 维度把一个大矩阵按 K 方向拆成几段交给不同线程块并行计算后再做累加这在 K 特别大时很有用也是 cuBLAS 的常用策略之一。还有一个方向是尝试使用张量核心路径比如 TF32 或者 FP16 累加。不过那就是另一个世界了指令、布局和精度策略都不一样。这篇实录里讲的 83%是在纯 FP32 乘加指令下得到的如果你把张量核心打开手写 kernel 和官方库的差距会重新拉开优化目标也会变成如何把数据喂得更快更满。我现在写矩阵乘法已经不会因为它只到 cuBLAS 的 83% 而觉得失败。真正值钱的是你清楚自己的内核慢在哪、快在哪也知道下一个百分点要去哪个方向找。如果非要说还有一条最值得分享的经验那就是先把 kernel 跑对再开 profiler最后才去抠指令级优化。顺序反过来你会在第一天就怀疑人生。

读完文章,也想定制专属网站?

尧图设计师 24 小时内与您沟通定制方案

免费获取报价 →
↑