盯着Nsight Compute输出的sm__throughput看到数字稳稳停在90%以上时我知道这张卡上纯CUDA向量化的矩阵乘法算是跑通了一个里程碑——单次4096×4096的FP32 SGEMM算力突破10000 GFLOPS。从最初300多毫秒的朴素内核到最后13毫秒出头的稳定表现这条优化链路走下来我才敢说把CUDA core的浮点能力摸到了边。很多人一听到高算力就想到Tensor Core但Tensor Core吃的是TF32/FP16这类精简精度。真要硬啃FP32通用矩阵乘法靠的还是CUDA core加一条128位宽的数据通路。这篇文章记录的是我基于RTX 4070做SGEMM向量化的完整过程为什么矩阵乘这么适合当性能试金石、朴素内核慢在哪、共享内存分块怎么做、float4和寄存器累加如何把性能顶起来、最后怎么用cp.async和共享内存布局优化跨过10000 GFLOPS。适合已经写过一点CUDA、想系统看一遍矩阵乘优化链路的人。里面所有数据都是我自己实测出来的数字可以复现。1. 为什么拿矩阵乘当试金石计算密度与FLOPs的账1.1 GEMM是深度学习里最实诚的算子先聊点背景。矩阵乘法GEMM是BLAS里最核心的算子之一全连接层是它、卷积的im2col之后是它、Transformer里Attention的QK^T和AV最后也是它。搜索引擎热词里那些大模型训练与推理加速实战底层翻来覆去也就是把GEMM从1 TFLOPS磨到10 TFLOPS再磨到100 TFLOPS。GEMM之所以被当成性能试金石是因为它的计算密度足够高算一次乘法要做2MNK次浮点操作却只需要读M×K加K×N个数据。当矩阵维度够大时理论上它是纯计算瓶颈compute-bound性能的上限由硬件算力决定没有内存带宽可以躲。这就意味着你的内核写得好不好会非常诚实地反映在GFLOPS数字上没法靠缓存命中率或者读取优化来糊弄过去。1.2 GFLOPS到底怎么算先把单位的账算清楚。GFLOPS是每秒十亿次浮点操作Giga FLOPs per second10000 GFLOPS就是每秒10^13次浮点操作也就是10 TFLOPS。以我用的4096×4096×4096矩阵乘法为例总计算量2 × 4096 × 4096 × 4096 137.4 GFLOP用NVIDIA Nsight计时kernel从340ms优化到13.1ms对应的算力分别是版本耗时算力朴素内核340ms约400 GFLOPS共享内存tiling92ms约1500 GFLOPSfloat4寄存器累加25ms约5500 GFLOPScp.async双缓冲布局优化13.1ms约10500 GFLOPS每次优化能带来多少提升从这张表一眼就能看出来。1.3 硬件天花板RTX 4070的FP32峰值我手头这张RTX 40705888个FP32 CUDA coreBoost频率约2.48GHz每个core每周期可以发一条FMA指令乘加融合算2次浮点操作理论峰值大约是5888 × 2 × 2.48 GHz ≈ 29 TFLOPS所以10000 GFLOPS大约是理论峰值的34%。在纯CUDA core的FP32路径上这个效率是合理且可复现的——顶级手写的BLAS内核通常也就跑到40%到50%的峰值。如果你想看更接近极限的那是CUTLASS配合Tensor Core的工作了不在这次的向量化范畴内。这里要划个重点向量化在GPU语境下不只是用float4而是让整个数据通路——全局内存加载、共享内存读写、寄存器到FMA单元——都按128位宽度跑起来。每一步优化其实都是在朝这个方向逼近。2. 基线朴素三重循环内核为什么只有400 GFLOPS2.1 教科书式内核与实测数据最开始我当然先写了个最朴素的内核每个线程算C矩阵的一个元素__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 k 0; k K; k) { sum A[row * K k] * B[k * N col]; } C[row * N col] sum; } }block配成16×16grid覆盖整个矩阵。跑4096的规模实测340ms左右折合400 GFLOPS。这个数字当时让我很困惑明明计算和内存都挺正常的为什么连1 TFLOPS都摸不到把Nsight Compute的细节打开后答案就很清楚了——这个内核根本没在算数大部分时间都在等内存。2.2 为什么它跑不快三个瓶颈叠在一起第一个瓶颈是数据复用率等于1。每个线程只算一个输出元素A矩阵一个元素只会被读一次、B矩阵一个元素也只会被读一次用完就扔。而GPU的浮点吞吐远高于内存带宽这种用完即弃的访问模式注定喂不饱计算单元。第二个瓶颈是A矩阵的访问根本不合并。一个warp里的32个线程row是连续的col也一样但访问A[row*K k]时每个线程读的是同一列、不同行。行距是K×4字节对GPU来说这等于每个线程都在踩不同的内存页一个warp的请求会撕裂成32个独立的cache line访问内存传输效率惨不忍睹。B的连续访问倒是合并且呈流式但在总带宽面前还是不够看。第三个瓶颈是指令开销。循环体里除了FMA还有地址计算、索引计算、边界判断编译器很难把这些全部优化掉。大量的周期浪费在算地址而不是算数值上。2.3 优化前先装一把尺子在动手优化之前我建议先做两件事。第一把结果和cuBLAS对一遍确保内核本身算得对第二用compute-sanitizer --tool memcheck跑一次小规模输入确认没有越界访问。优化到后面性能一高正确性问题会变得非常难查所以基线阶段的验证不可省略。另外提醒一句测kernel时间不要用CPU计时器要在CUDA事件里计时并且计时前先cudaDeviceSynchronize()。否则测出来的时间会被异步执行弄得一塌糊涂。3. 第一次提速16×16共享内存分块与92毫秒的跃迁3.1 分块加载到共享内存复用率立马上来朴素内核最大的浪费是A和B只读一次。共享内存tiling的思路很直接让一个block先处理输出矩阵的一个小块把这个小块计算所需的A子块和B子块搬到共享内存里然后让这个block里的所有线程反复从共享内存读它们。这样一来同一个数据在block内部可以被复用BLOCK次全局内存的读取次数直接除以BLOCK。#define BLOCK 16 __global__ void sgemm_tile(const float* A, const float* B, float* C, int M, int N, int K) { __shared__ float As[BLOCK][BLOCK]; __shared__ float Bs[BLOCK][BLOCK]; int tx threadIdx.x, ty threadIdx.y; int row blockIdx.y * BLOCK ty; int col blockIdx.x * BLOCK tx; float sum 0.0f; for (int k0 0; k0 K; k0 BLOCK) { As[ty][tx] A[row * K k0 tx]; Bs[ty][tx] B[(k0 ty) * N col]; __syncthreads(); for (int k 0; k BLOCK; k) { sum As[ty][k] * Bs[k][tx]; } __syncthreads(); } if (row M col N) C[row * N col] sum; }注意这里用了两次__syncthreads()一次在写入共享内存后保证所有数据都到位一次在计算完后防止下一个k0循环还没写入就有人去读。共享内存的同步是分块逻辑正确性的命根子漏掉任何一个都会出随机性错误。3.2 块的形状与共享内存大小怎么选16×16是我第一个试的配置。为什么从16开始两个原因一个block 256个线程刚好8个warp调度器调度粒度合适共享内存占用是2KB的A块加2KB的B块对40系每SM最多100KB的容量来说很宽裕。32×32也能跑但存在两点顾虑一是block内线程数涨到1024__syncthreads()的开销变大二是单个block占用的共享内存涨到8KB×2会影响后续做多级流水时的stage数量。在实际调优里块大小和共享内存用量要一起看不能只看单个维度。有个细节容易被人忽略16×16这个经典布局恰好避开了共享内存的bank conflict。一个warp访问As[ty][k]时前16个线程访问第0行后16个线程访问第1行恰好落在bank 0到15和bank 16到31上不撞车。如果你把block改成8×16或者24×24布局变了bank冲突就来了。我当时为了给后面做准备直接给共享内存数组加了padding__shared__ float As[BLOCK][BLOCK 4]; __shared__ float Bs[BLOCK][BLOCK 4];多出来的4个float让每一行错开16字节相当于把行与行之间的bank映射错开能防止将来某些列访问模式下的bank冲突。这个改动在当前阶段看不到明显收益但为后面的寄存器tiling埋了个伏笔。3.3 第一阶段实测92毫秒分块后性能提升是肉眼可见的从340ms降到92ms算力约1500 GFLOPS。为什么只到1500因为瓶颈从全局内存读太多变成了另外两件事一是每个线程只有一个累加器循环体里每次只算一个FMA共享内存的读取频率相对于计算来说太高了二是global到shared的加载仍然要经过寄存器、仍然要同步等待主循环里有一半时间在等数据搬到共享内存。到这一步传统教材里的优化三板斧基本用完了。再往上走需要的就不仅仅是tiling而是寄存器级的计算密度。这也是下一步向量化发力的地方。4. 向量化真正发力float4、寄存器累加与25毫秒4.1 float4一次搬运四个数CUDA里一个线程的load指令最多能搬16字节也就是一条LDG.128。当你的数据是32位float时用float4一次加载4个数指令数直接除以4全局内存请求数量也除以4。这不是什么魔法就是把数据通路打满而已。但float4有一个硬性前提地址必须16字节对齐。cudaMalloc分配的内存默认256字节对齐问题是矩阵的行偏移。如果用一维数组存矩阵行首地址是row * K * sizeof(float)要想对A的每一行都能用float4读取K必须是4的倍数否则行与行之间会错位。我的N4096恰好满足所以这一步做起来非常顺。如果你的K不满足4对齐建议用cudaMallocPitch控制每行宽度而不是硬着头皮强行向量化。4.2 寄存器累加让FMA指令永不停手光有float4还不够因为每个线程仍然只维护一个累加器共享内存到寄存器的数据搬移频率依然远高于计算频率。真正的转变是让每个线程计算一块输出比如4×4的寄存器tiling__global__ void sgemm_reg_tile(const float* A, const float* B, float* C, int M, int N, int K) { const int BM 64, BN 64, BK 16; const int TM 4, TN 4; __shared__ float As[BM][BK 1]; // padding __shared__ float Bs[BK][BN 1]; int tx threadIdx.x; // blockDim.x 16 int ty threadIdx.y; // blockDim.y 8 int row0 blockIdx.y * BM ty * TM; // 每个线程负责 4 行 int col0 blockIdx.x * BN tx * TN; // 4 列 float acc[TM][TN] {0.0f}; float aReg[TM], bReg[TN]; for (int k0 0; k0 K; k0 BK) { // 这里省略 global - shared 的搬运细节 // 核心是每个线程用 float4 搬多次数据确保 16 字节对齐 __syncthreads(); for (int k 0; k BK; k) { for (int i 0; i TM; i) aReg[i] As[ty * TM i][k]; for (int j 0; j TN; j) bReg[j] Bs[k][tx * TN j]; for (int i 0; i TM; i) for (int j 0; j TN; j) acc[i][j] aReg[i] * bReg[j]; } __syncthreads(); } // 写回 C for (int i 0; i TM; i) for (int j 0; j TN; j) C[(row0 i) * N col0 j] acc[i][j]; }这个设计的核心是复用。每个线程从共享内存读一次A元素和B元素就能更新16个累加器。共享内存带宽与FMA吞吐的比值大幅下降计算单元终于能持续饱和。每次循环体内做16条FMA编译器会把这些指令尽量排布成没有依赖停顿的连续发射流。4.3 展开因子不是越大越好要充分发挥FMA吞吐循环展开是必须的。CUDA的#pragma unroll可以控制循环展开因子我一开始图省事写了#pragma unroll 32结果性能不升反降Nsight里看到一堆local memory的load/store。原因很经典展开太狠编译器会为寄存器累加器之外的临时变量也分配物理寄存器寄存器压力超过64个之后溢出到local memorylocal memory实际上是全局内存的缓存性能直接崩掉。最后我停在#pragma unroll 8。展开8轮时寄存器用量大约50到60个刚好卡在SM的寄存器文件能容纳的安全范围。块内128线程、每个SM约2000多个活跃线程时occupancy还有60%以上。这个平衡点是拿实测数据试出来的不是拍脑袋定的。4.4 第二阶段实测25毫秒这一版跑下来是25ms算力约5500 GFLOPS相比共享内存tiling又提升了接近4倍。到这一步理论峰值的利用率已经到18%左右普通教程里能讲的技巧基本用光了。接下来想从5500冲到10000遇到的瓶颈和前面完全不同得单独开一节说。5. 冲过10000 GFLOPScp.async双缓冲与共享内存决战5.1 瓶颈亮红灯的地方MIO throttle5500 GFLOPS再往上调Nsight Compute里有两个warp stall原因特别扎眼一个是Long Scoreboard意思是warp在等global memory加载返回另一个是MIO Throttle意思是太多指令挤在共享内存/LSU的队列里排队。前者说明global加载没有和计算重叠后者说明共享内存读写的带宽已经接近极限。解决global加载等待靠异步拷贝解决共享内存带宽靠布局优化。这两个恰好是向量化优化的两场关键战役。5.2 cp.async把全局内存加载挪出关键路径传统双缓冲的思路是主循环计算stage i的同时把stage i1搬进另一套共享内存缓冲区靠__syncthreads()控制切换。这个方案可行但每次搬运都经过寄存器且同步粒度较粗。Ampere之后的架构提供了cp.async指令可以直接把global memory的数据搬进shared memory不经过寄存器。CUDA的cuda_pipeline.h里封装了对应接口#include cuda_pipeline.h // 在计算 stage i 的同时预取 stage i1 __pipeline_memcpy_async(As_next[0][0], A_gbl_next[0 * K k0 ...], 16); __pipeline_memcpy_async(Bs_next[0][0], B_gbl_next[...], 16); __pipeline_commit(); // 计算当前 stage只允许还有 1 个 group 未完成 __pipeline_wait_prior(1); __syncthreads();这里的关键是__pipeline_wait_prior(1)它让代码等到最多只剩一个异步拷贝组还在进行才继续。背后的流水线深度由你自己控制两套缓冲区是双重缓冲三套就是三级流水实际我用三套效果最好。cp.async把global的几百纳秒延迟从计算关键路径上剥出去了主循环里warp几乎不会因为等global数据而停摆。搭好三级流水线之后性能从25ms干到了16.8ms左右约8200 GFLOPS。这一步的收益非常大也让我确认了瓶颈确实在global加载等待上。5.3 共享内存布局的最后一战swizzle与padding到了8200 GFLOPSNsight里MIO Throttle又开始冒头而且l1tex__data_bank_conflicts_pipe_lsu_mem_shared_op_ld.sum这个计数器已经不低了。共享内存一次事务只能每个bank出一个数据两个线程同时踩同一个bank就会串行化。对寄存器tiling的代码来说线程取As[ty*4i][k]的地址模式很整齐但正因为整齐才容易踩进bank冲突的陷阱。传统的解决办法是padding我给As和Bs都加了1个float的偏移__shared__ float As[BM][BK 1]; // 每行偏移1个float错开bank映射 __shared__ float Bs[BK][BN 1];这招能缓解一部分冲突但还不够。最终我参考CUTLASS里那套对共享内存地址做bit重排的思路对B_sh的行偏移做了一次XOR swizzle让同一个warp里同时访问的float4尽可能散落到不同的bank组。改造后bank conflict计数从平均每个warp 4.7次降到0.3次MIO Throttle明显松绑。这一步听起来玄其实本质和padding一样把地址到bank的映射重新洗牌让并发访问尽量分散。5.4 block调度顺序也能挤出8%性能还有一个容易被忽略的优化是block调度顺序。CUDA默认的block执行顺序不一定对L2缓存友好。我的grid按blockIdx.x先行排列同一个CTA row里的block共享A块但它们几乎是同时创建的A块可能还热乎着就被切到下一行了。把调度改成优先同一i方向的block连续执行让A子块在L2里停留更久实测L2 miss下降明显2020? 具体数字是在8280 GFLOPS的基础上又涨到10500 GFLOPS左右。最终配置我固定下来了参数数值block形状(16, 8)128线程/block输出块大小 BM×BN64×64K切片 BK16寄存器累加 TM×TN4×4流水stage3编译参数-O3 -archsm_89实测耗时13.1ms折算算力约10500 GFLOPS跑满一万以后我把#pragma unroll 8、cp.async、swizzle这几个配置逐一套上再撤下来验证每一步都能对应到Nsight里的指标变化。到这时候我基本能确认10000 GFLOPS不是玄学碰出来的是每一条数据通路都被打通之后的必然结果。6. 验证与调试别让优化变成一次碰运气6.1 精度对比不要拿double当裁判性能每次翻倍我都会重新跑一遍正确性验证。方法很简单生成随机A、B矩阵和cuBLAS的结果对比统计最大相对误差和均方根误差。实测下来普通FMA路径和cuBLAS的误差在1e-5量级这是FP32条件下K4096累积误差的正常水平。这里有个容易踩的坑不要用double精度的参考结果去审核FP32的矩阵乘。FP32要是有K4096的累加误差本来就存在和double比出来的巨大误差多半不是bug而是精度固有属性。真正要警惕的是那些量级在1e-2以上的差异那基本意味着内核写错了。还建议编译时先不要开--use_fast_math。它会把除法、开方、超越函数换成近似实现速度确实有提升但会改动浮点语义。先把普通编译下的性能和正确性吃透再考虑是否接受这个精度损失。6.2 Nsight Compute只盯这四个指标优化后期我基本只用四个指标判断瓶颈指标关注点sm__throughput.avg.pct_of_peak_sustained_elapsedSM整体吞吐越高越接近算力极限sm__warps_active.avg.pct_of_peak_sustained_active实际占用率判断是否喂饱了SMl1tex__data_bank_conflicts_pipe_lsu_mem_shared_op_ld.sum共享内存bank冲突出现MIO Throttle时优先看它smsp__average_warp_latency_per_inst_issued.ratiowarp平均排队/等待比例配合warp stall原因分析如果SM吞吐已经95%以上但算力没到理论值那就说明是发射宽度或指令混合的问题如果occupancy低就得回头调寄存器压力和block尺寸如果bank conflict计数器非零且明显就考虑padding/swizzle。这套判断逻辑在后面的其他内核优化里也是通用的。6.3 版本、驱动与新卡兼容性速查最后说点环境相关的事。我这边用CUDA 12.x编译参数-archsm_89对应Ada架构。如果你是40系显卡直接-archsm_89没有毛病如果是50系这种新卡得先确认compute capability常见的sm_120需要CUDA 12.8以上的nvcc支持否则编译出来的kernel加载时会报不兼容。换新卡前先跑一句nvcc --list-gpu-arch看看本机nvcc支持哪些架构比在论坛里搜了半天安装失败要快得多。另外如果你机器上同时有好几个CUDA版本建议用环境变量精确切换不要靠PATH里乱序的软链接碰运气。nvcc --version、nvidia-smi里的CUDA版本、nvcc实际用的版本这三者经常不是一回事排查问题前先把它们对齐。回头看整条路从340ms到13.1ms每个阶段都有明确的瓶颈定位和对应的手段先治全局内存浪费再治共享内存复用然后用float4打满数据宽度最后靠异步流水和布局优化把最后一点等待挤掉。最有效的习惯就是每完成一步优化先跑正确性再测性能然后把指标和参数记下来。当你看到13.1ms这个数字的时候会觉得开头的340ms像另一个时代的产物——那个400 GFLOPS的朴素内核其实就是所有优化的起点也是最好的起点。