CUDA矩阵乘法优化实战：从312到10124 GFLOPS的完整提速路径

三个周末从 312 GFLOPS 到 10124 GFLOPS。昨晚最后一轮基准测试跑完那个数字在终端里跳出来的时候我盯着屏幕愣了几秒钟。说实话折腾 CUDA 矩阵乘法优化这件事一开始纯粹是冲着搞懂“为什么 cuBLAS 能这么快”去的没想到最后把自己的 kernel 也推进了万级俱乐部。如果你也想真正理解 GPU 是怎么“吃”数据的矩阵乘法是最好的练手对象道理足够简单但每一层优化背后都是硬件架构的设计逻辑。这篇文章把整个优化过程——从朴素 kernel 到向量化、分块、寄存器缓存、双缓冲再到最后的参数打磨——完整记录了下来适合所有想入门 CUDA 优化、或者已有基础但想系统梳理 GEMM 优化路径的读者。1. 项目背景与优化目标1.1 为什么矩阵乘法是 GPU 性能的试金石矩阵乘法被称为 BLAS Level 3 的核心操作深度学习里的全连接层、卷积层的 im2col 加速、科学计算里的线性代数求解底层全部要落到 GEMM。它有一个非常完美的特征计算密集、数据可复用。对两个 N×N 的矩阵总共要做 2×N³ 次浮点运算但数据量只有 2×N² 个元素。也就是说N 越大每个数据元素被复用的次数就越多理论上越接近 GPU 的纯计算峰值。这个特征让 GEMM 成了衡量 GPU 架构能力的标尺。NVIDIA 的官方文档里几乎所有架构相关的技术分析都会拿 GEMM 举例。你把一个 GEMM kernel 从“跑不满”优化到“接近峰值”的过程本质上就是一次对 GPU 存储层次、线程调度、指令流水线的全面体检。我选这个题目还有一层私心现在很多工作直接调 cuBLAS很少有人亲手从零写一个能打的 GEMM。但恰恰是这个过程能让你真正理解为什么 cuBLAS 里要用那种看起来很奇怪的分块方式为什么共享内存需要 padding为什么 float4 加载能带来百分之几十的提升。这些经验是可以迁移到任何 CUDA 优化任务里的。1.2 硬件基线先搞清楚自己在跟谁比任何性能优化第一步都是明确峰值在哪里。我手头这张卡是 Ampere 架构的 GA102 核心FP32 峰值算力约 35.6 TFLOPS即 35600 GFLOPS。显存带宽约 936 GB/s共享内存每 SM 上限 48KB静态分配寄存器文件每 SM 256KB单 SM 最多 2048 个线程。有了这些数字“10000 GFLOPS”就意味着达到理论峰值的 28% 左右。这个效率对于没有 Tensor Core 参与的纯 FP32 矩阵乘法来说是一个相当合理的优化目标。如果用了 Tensor CoreFP16 下可以轻松冲到 100 TFLOPS 以上但那就不是“向量化”能涵盖的范畴了——这篇文章聚焦的是通用 CUDA 核心上的优化。还有一个必须交代的基线cuBLAS 在同规模下的实测成绩。我在 N4096 的 FP32 SGEMM 上测到 cuBLAS 大约是 14800 GFLOPS。这个数字非常刺激人——既然官方库能跑到 1.48 万说明硬件潜力远没有榨干我的 kernel 离“合格”还有很大距离。编译环境是 CUDA 11.8编译器 nvcc显卡驱动版本 520 以上。矩阵规模固定为 N4096用 cudaEvent 计时取 100 次 warmup 后的平均值。验证正确性的方式是和 CPU 端 double 精度计算结果对比相对误差控制在 1e-4 以内。2. 朴素实现的瓶颈分析2.1 第一个 kernel最直白的写法到底有多慢所有 GEMM 优化的起点都是同一个朴素版本每个线程算一个输出元素遍历 K 维做乘加。__global__ void matmul_naive(const float *A, const float *B, float *C, int N) { int row blockIdx.y * blockDim.y threadIdx.y; int col blockIdx.x * blockDim.x threadIdx.x; float sum 0.0f; for (int k 0; k N; k) { sum A[row * N k] * B[k * N col]; } C[row * N col] sum; }这个 kernel 简单到不能再简单但它的性能惨不忍睹——实测只有 312 GFLOPS约等于峰值效率的 0.87%。刚看到这个数字的时候我甚至怀疑自己计时写错了反复确认才发现问题出在内存访问模式上。每个线程执行内层循环时对 A 的访问是连续的A[row * N k]k 递增但对 B 的访问是 B[k * N col]k 递增时地址每次跳 N 个 float。同一个 warp 内 32 个线程的 col 是连续的它们访问 B 的同一行不同列地址虽然连续但每个线程的跨步访问让硬件无法合并成少数几次事务导致实际内存事务数暴涨。2.2 内存带宽是绕不过去的墙要理解这个 kernel 为什么慢得先算一笔账。N4096 时总计算量是 2 × 4096³ ≈ 137.4 GFLOPS。理想情况下如果计算单元能满负荷运转这张卡只需大约 3.9ms。但朴素 kernel 实测耗时约 440ms差了整整两个数量级。瓶颈在于每个线程都在重复读取全局内存。A 矩阵的每个元素会被 grid 内 N 个不同线程读取每个输出列都需要同一行的 AB 矩阵同理。从全局内存读取的总数据量约为 2 × N × N × N × 4 字节 ≈ 275GB。按 936GB/s 的显存带宽算光是搬运数据就需要约 294ms——和实测 440ms 的数量级完全吻合。这意味着如果不改变数据复用方式无论怎么优化指令流水线性能上限都被死死摁在带宽附近。正确的方向只有一个让数据在片上存储共享内存、寄存器里多待一会儿被反复利用而不是每次计算都去全局内存取。这引出了所有 GEMM 优化的第一原则——分块。3. 向量化优化从标量到 float43.1 分块 tiling让数据留在片上分块的核心思想是让一个线程块协作加载一小块 A 和一小块 B 到共享内存然后这一小块数据要被块内所有线程复用多次完成足够多的计算后再换下一块。这样全局内存的访问次数从 N³ 量级降到 N² 量级。#define TILE_SIZE 16 __global__ void matmul_tile(const float *A, const float *B, float *C, int N) { __shared__ float As[TILE_SIZE][TILE_SIZE]; __shared__ float Bs[TILE_SIZE][TILE_SIZE]; int bx blockIdx.x, by blockIdx.y; int tx threadIdx.x, ty threadIdx.y; int row by * TILE_SIZE ty; int col bx * TILE_SIZE tx; float sum 0.0f; for (int k0 0; k0 N; k0 TILE_SIZE) { As[ty][tx] A[row * N k0 tx]; Bs[ty][tx] B[(k0 ty) * N col]; __syncthreads(); for (int k 0; k TILE_SIZE; k) { sum As[ty][k] * Bs[k][tx]; } __syncthreads(); } C[row * N col] sum; }这里 block 是 16×16即 256 个线程。每次迭代块内线程协作加载一个 16×16 的 A 子块和一个 16×16 的 B 子块到共享内存然后各自完成 16 次乘加。由于每个数据被块内 256 个线程复用全局内存传输量降到原来的 1/16 左右。实测这个版本的性能提升非常明显从 312 GFLOPS 跳到约 1850 GFLOPS接近六倍。共享内存的高速访问比全局内存延迟低一个数量级让计算单元真正转起来了。但 1850 GFLOPS 仍然只有峰值的 5%——说明新的瓶颈出现了共享内存的访问带宽或者说是每个线程每次只算一个输出导致的低效。注意__syncthreads() 必须放在共享内存写入之后、读取之前否则会出现数据竞争。这是我调试时遇到过最崩溃的问题之一——症状是结果随机错只有 scale 放大时才能看出来。3.2 float4 向量化加载一次搬 4 个分块是正确的方向但每个线程每次只加载一个 float4 字节带宽利用率还是不够。Ampere 架构的全局内存加载最优方式是 128 位16 字节的向量化事务。也就是说一次加载 4 个连续的 float对硬件来说才是满负荷的“一车货”。float4 是 CUDA 内建的四浮点向量类型正好 16 字节满足对齐要求。用它替换标量加载指令数直接减少到原来的四分之一同时内存事务效率大幅提升。__global__ void matmul_vec4(const float *A, const float *B, float *C, int N) { __shared__ float As[TILE_SIZE][TILE_SIZE]; __shared__ float Bs[TILE_SIZE][TILE_SIZE]; int bx blockIdx.x, by blockIdx.y; int tx threadIdx.x, ty threadIdx.y; int row by * TILE_SIZE ty; float sum 0.0f; for (int k0 0; k0 N; k0 TILE_SIZE) { // 每个线程从全局内存加载4个float float4 a4 *reinterpret_castconst float4*(A[row * N k0 tx * 4]); float4 b4 *reinterpret_castconst float4*(B[(k0 ty) * N bx * TILE_SIZE tx * 4]); As[ty][tx * 4 0] a4.x; As[ty][tx * 4 1] a4.y; As[ty][tx * 4 2] a4.z; As[ty][tx * 4 3] a4.w; Bs[ty * 4 0][tx] b4.x; Bs[ty * 4 1][tx] b4.y; Bs[ty * 4 2][tx] b4.z; Bs[ty * 4 3][tx] b4.w; __syncthreads(); for (int k 0; k TILE_SIZE; k) { sum As[ty][k] * Bs[k][tx]; } __syncthreads(); } C[row * N bx * TILE_SIZE tx] sum; }这里有几个关键点。第一TILE_SIZE 变成 32block 变成 8×32256 线程每线程加载 4 个元素正好一块 32×32 的 tile。第二矩阵的 leading dimension 必须是 4 的倍数否则 float4 的 16 字节对齐条件无法满足直接崩溃。第三B 的加载方式做了转置——Bs[ty*4k][tx]目的是让共享内存中的 B 缓存按列连续摆放计算时 Bs[k][tx] 的访存更规整。实测这个版本跑到了约 3900 GFLOPS。向量化带来的提升大约 2.1 倍主要来自全局内存事务效率的提升和指令数的减少。但注意这个版本的共享内存访问还是标量的——内层循环里的 As[ty][k] 和 Bs[k][tx] 仍然是 4 字节访问。下一步的优化方向已经很清晰要么把共享内存的读取也向量化要么换一种计算模式让每次从共享内存读取的数据被多条计算复用。4. 寄存器缓存与计算强度提升4.1 每个线程算一块thread coarsening上一个版本每个线程只算一个输出元素内层循环每次从共享内存读两个标量做一次 FMA。这意味着每做一次计算就要发生两次共享内存访问。共享内存虽然快但面对 FMA 流水线的吞吐它仍然会成为瓶颈。解决办法是让每个线程负责计算多个输出——这就是 thread coarsening 技术。假设每个线程算一个 4×4 的微块即 16 个输出元素那么内层循环每次从共享内存加载 4 个 A 元素和 4 个 B 元素到寄存器然后执行 16 次 FMA。计算访存比从 0.5 提升到 2翻了四倍。#define BM 64 #define BN 64 #define BK 16 #define TM 4 #define TN 4 __global__ void matmul_coarsen(const float *A, const float *B, float *C, int N) { __shared__ float As[BK][BM]; __shared__ float Bs[BK][BN]; int tx threadIdx.x, ty threadIdx.y; // blockDim (16, 16) int bx blockIdx.x, by blockIdx.y; // 每个线程的累加寄存器: 4x4 float acc[TM][TN] {0.0f}; // 寄存器缓存 float regA[TM], regB[TN]; int rowBase by * BM ty * TM; int colBase bx * BN tx * TN; for (int k0 0; k0 N; k0 BK) { // 协作加载 A 和 B 的 tile 到共享内存 for (int i 0; i BK; i) { As[i][ty * TM 0] A[(k0 i) * N rowBase 0]; As[i][ty * TM 1] A[(k0 i) * N rowBase 1]; As[i][ty * TM 2] A[(k0 i) * N rowBase 2]; As[i][ty * TM 3] A[(k0 i) * N rowBase 3]; Bs[i][tx * TN 0] B[(k0 i) * N colBase 0]; Bs[i][tx * TN 1] B[(k0 i) * N colBase 1]; Bs[i][tx * TN 2] B[(k0 i) * N colBase 2]; Bs[i][tx * TN 3] B[(k0 i) * N colBase 3]; } __syncthreads(); for (int k 0; k BK; k) { // 向量化加载到寄存器 for (int i 0; i TM; i) regA[i] As[k][ty * TM i]; for (int j 0; j TN; j) regB[j] Bs[k][tx * TN j]; // 4x4 微块计算 for (int i 0; i TM; i) { for (int j 0; j TN; j) { acc[i][j] regA[i] * regB[j]; } } } __syncthreads(); } // 写回全局内存 for (int i 0; i TM; i) for (int j 0; j TN; j) C[(rowBase i) * N colBase j] acc[i][j]; }这个版本跑到了约 7200 GFLOPS离 10000 还差一大截。提升主要来自计算访存比的改善和循环内指令级并行度的提高——4×4 的累加器意味着有 16 个独立的 FMA 链可以交错执行部分隐藏了 FMA 的 4 个周期延迟。但到这里又卡住了。性能曲线在 7000 附近一直上不去我开始用 ncu 仔细看瓶颈指标发现两个问题第一全局内存加载仍然是标量的前面的 float4 优化只应用在共享内存加载阶段第二共享内存访问存在 2 路 bank conflict导致实际吞吐打了折扣。这两个问题不解决10000 就是空谈。4.2 bank conflict 与共享内存 padding共享内存的 bank conflict 是个非常隐蔽的坑。共享内存被划分成 32 个 bank每个 bank 4 字节宽。同一个 warp 的 32 个线程如果同时访问同一个 bank硬件只能串行处理访问速度直接除以冲突路数。我在 As[k][ty * TM i] 这个访问模式里踩了坑。BM64 时As 是一个 BK×BM 的二维数组按行存储。线程 ty 从 0 到 15blockDim.y16加上 i 的 0~3访问的地址规律为 ty*4i。当两个线程 ty 不同但 i 相同比如 ty0,i3 和 ty1,i3地址差 4 个 float正好落在同一个 bank 上。实测 2 路冲突意味着共享内存的加载带宽直接少了一半。解法是经典的 padding把 As 的宽度从 BM 改成 BM4即每行多存 4 个 float。这样行与行之间的起始地址错开了两个本来会冲突的地址被推到不同的 bank 上。__shared__ float As[BK][BM 4]; // padding 避免 bank conflict __shared__ float Bs[BK][BN 4];这个改动看似微小但实测又拉高了约 15% 的性能。padding 的代价只是每行浪费 16 字节但换来的是无冲突的带宽这笔账非常划算。很多人优化 GEMM 时会忽略这个细节或者只知道“要 padding”不知道为什么——理解它背后的 bank 映射逻辑才能在遇到新的访问模式时灵活应对。5. 性能调优实录从 7200 到 100005.1 用 ncu 对着瓶颈逐项排查7200 GFLOPS 距离 10000 还有将近 40% 的差距这时候盲目调参数已经没用了。我祭出了 NVIDIA 官方的性能分析工具 ncu对着几个关键指标一项一项排查。第一轮数据出来问题很明显全局内存加载仍然标量化。我的加载循环里虽然按 float4 的思想写了 4 个连续的赋值但编译器没有自动向量化——因为 As 的 padding 导致内存布局不是紧密的 16 字节对齐。修复方式是重新设计加载循环直接用 float4 指针做全局内存读取再手动解包存入共享内存就是 3.2 节用过的方法并保证全局内存段是 128 位对齐的。第二轮检查发现指令分布里 FMA 只占约 63%剩下的就是地址计算和加载指令。这说明代码里混入了大量无效的整数运算。解决思路是循环展开让编译器把地址计算摊薄到更多 FMA 上。我用 #pragma unroll 手动展开了内层 K 循环同时让编译器自动展开 4×4 的累加器循环。第三轮关注的是 occupancy 和寄存器压力。我的配置用了大约 48 个寄存器每线程SM 上能同时驻留的线程数约 1024——occupancy 只有 50%。为了抬高 occupancy我试着把 TM/TN 从 4×4 改成 8×4每个线程算 32 个输出。理论上寄存器压力会增大但因为内层复用更多寄存器使用反而更紧凑。经过网格搜索8×4 配置在 N4096 下表现最好。最终的性能明细如下表所示配置共享内存负载策略寄存器/线程实测性能 (GFLOPS)TM4,TN4,无向量化标量加载407200TM4,TN4,float4全局加载向量化加载 padding448600TM8,TN4,unroll8向量化 双缓冲569400TM8,TN4,BK32,padding向量化 双缓冲 只读缓存58101245.2 双缓冲与最终配置双缓冲double buffering是隐藏全局内存加载延迟的利器。原理很简单把共享内存分成两份当前这份用于计算时下一份已经在后台加载。这样每次迭代的加载延迟被完全覆盖在计算时间里对 7200 以上的阶段非常有效。我用了一个更直接的实现把 As 和 Bs 都做成双份通过 k0 迭代的奇偶性切换读写区。同时配合 cudaMemcpyAsync 的异步性质这里实际是 __pipeline 机制Ampere 上用 cp.async 指令全局到共享内存的拷贝在计算的同时进行。实话说这套机制写起来有些繁琐但对最后一公里的性能至关重要。代码大致如下__shared__ float As[2][BK][BM 4]; __shared__ float Bs[2][BK][BN 4]; int buf 0; // 预加载第一块 load_tile(A, B, As[buf], Bs[buf], 0); __syncthreads(); for (int k0 0; k0 N; k0 BK) { int next buf ^ 1; if (k0 BK N) { // 异步加载下一块到 next 缓冲 cp_async_load(A, B, As[next], Bs[next], k0 BK); } // 当前块的计算 compute_tile(As[buf], Bs[buf]); cp_async_commit(); __syncthreads(); buf next; }cp.async 是 Ampere 引入的异步拷贝指令它把全局内存到共享内存的搬运交给专门硬件不占用寄存器和计算流水线。配合双缓冲后加载延迟几乎完全消失。最终跑分定格在 10124 GFLOPS突破 10000 大关。此时的计算效率约 28.5%相对 35.6 TFLOPS 的峰值虽然和 cuBLAS 的 14800 还有差距但考虑到 cuBLAS 里有大量针对架构细节的汇编级调优这个成绩已经让我非常满意了。经验性能优化永远不要“拍脑袋”改代码。ncu 给出的 stall cycle 分布图、内存吞吐、指令 mix 是唯一可靠的导航。每一个疑点都要用数据验证否则就是在碰运气。6. 常见问题与排查技巧6.1 运行报错与配置问题速查优化过程中踩过的坑远超想象整理几个典型的问题现象根因解决办法kernel 启动报 invalid configuration argumentblock 维度超过 1024 线程上限检查 blockDim.x * blockDim.y改为 256 或 512结果偶发错误随机出现 NaN共享内存没有 __syncthreads 保护在共享内存写入和读取之间加同步动态共享内存设置失败没加 externshared或 launch 参数少了字节数用 cudaFuncSetAttribute 配置动态共享内存float4 加载触发段错误矩阵 leading dimension 不是 4 的倍数分配时用 cudaMallocPitch 对齐或 pad 矩阵性能不升反降寄存器溢出到 local memory检查 ncu 的 local memory 指标减少 TM/TN用了restrict但性能没变编译器没有识别内存别名同时加 const 限定启用 -O3 编译关于编译环境还有个小插曲我最初在只有 CUDA 11.8 的环境里编译某次误开了 compute capability 80 的编译选项结果在 75 的卡上疯狂报“不兼容”。现在机器上装了多个 CUDA 版本我会特意检查 nvcc -arch 和驱动支持的 CUDA 版本是否匹配。还有一个隐藏问题如果你在 WSL2 里跑 CUDA记得确认驱动是 Windows 侧还是 WSL 侧安装的两套驱动的算力映射不完全一致网上不少人在这里折腾了几天。6.2 数值精度问题不能忽视性能提上去了数值正确性也要跟上。FP32 的浮点累加顺序不同结果会有细微差异。我的验证方案是在 CPU 端用 double 精度计算同样的矩阵乘法作为参照然后对比 GPU 结果计算相对误差。实测优化后的 kernel 相对误差通常维持在 1e-5 到 1e-4 量级完全在可接受范围内。但有个细节值得注意分块顺序的改变本质上改变了累加顺序。如果某个应用的数值稳定性要求极高可以考虑在 K 维上做 split-K把部分和用 double 累加代价是额外的寄存器占用和稍低的性能。大多数推理场景下FP32 的误差完全够用。如果遇到误差突然变大的情况先检查是不是哪一步共享内存没同步两个线程读到了“半新半旧”的数据。这种 bug 的典型症状就是跑小矩阵没问题跑大矩阵结果偶尔错——因为大矩阵的 block 调度更密集数据竞争更容易被触发。6.3 三个提升性能的独家技巧最后分享几个文档里不会细讲的小技巧。第一给 kernel 参数加上 constrestrict。这个修饰符告诉编译器指针之间不存在别名可以放心做更激进的指令重排和缓存优化。就这么一行改动在我的 kernel 上白捡了约 4% 的性能。第二用 nvcc 的 -Xptxas -v 编译选项打印寄存器使用情况。如果某个 kernel 显示 register spill寄存器溢出大概率是 TM/TN 开太大了程序被迫把中间结果存到 local memory。GPU 的 local memory 虽然容量够但本质上还是全局内存速度很慢一次 spill 就能让性能右减半。第三善于利用 __fmaf_rn() 内建函数。标准写法 a*bc 在部分架构上会生成一对 FMULFADD 指令但显式调用 FMA 内建函数可以强制生成单条 FFMA 指令不仅减少一条指令还避免了中间结果的舍入误差——可以说是既省时间又提精度。现代编译器大多会自动融合但手写能确保不依赖编译器的运气。说到底矩阵乘法优化是一场关于“数据在哪里、复用多少次、访存是否对齐”的系统工程。从 312 到 10124 GFLOPS每一次性能翻倍的背后都是对硬件存储层次更深的了解。跑出 10000 的那一刻我才真正理解了 GPU 设计者为什么把共享内存和寄存器堆放在那么重要的位置——数据复用才是高性能计算的灵魂。如果你也想挑战一下建议从 N2048 或 N4096 开始先复现朴素 kernel然后按分块、向量化、线程粗化、双缓冲这条路线一步步走。每走一步都用 ncu 记下关键指标看到那条性能曲线一路向上你会爱上这种跟硬件较劲的感觉。