GEMM优化实战：从矩阵乘法到GPU算力利用率提升

矩阵乘法GEMMGeneral Matrix Multiply这个算子估计不少做深度学习的人天天见它但很少真正把它当回事。我一开始也是这么想的直到有一次在某个跨平台系统上跑一个图像检测模型发现前向推理时间全部耗在了一个看似平平无奇的矩阵乘上调了一整天模型结构都没用。从那时候开始我就决定专门深耕一下GEMM优化这个项目就是我注册的DeepGEMM实验项目目标很单纯把矩阵乘法在GPU上的计算效率压榨到接近硬件的理论峰值。这个项目适合三类人看一是做推理引擎或者算子库的人二是想搞懂CUDA和kernel优化的后端工程师三是那些好奇为什么同一张显卡跑同一个模型别人的实现就是比你快三倍的调优玩家。我会把自己的实测过程、踩坑记录和排错思路全部展开讲不吹不黑能复现的尽量给足数据。1. 我为什么愿意花时间去做一个GEMM项目1.1 GEMM在深度学习里的真实地位很多人觉得卷积就是卷积全连接就是全连接注意力机制就是注意力机制但真把计算图摊开看它们底层干的基本都是同一件事大量的矩阵乘加运算。卷积的im2col变换之后是矩阵乘全连接本身就是矩阵乘Transformer里那套QKV投影和attention score计算本质上也是一堆GEMM。所以GEMM优化的价值特别直白你只要把GEMM的算力效率提上去整个模型的前向、后向、推理、训练全都能受益。它不是某一个点上的小优化而是托底的玩意儿。我之前在某公司做性能优化的时候测过一个语言模型GEMM类算子差不多占了总耗时的60%~70%。你要是能让这部分的FLOPS利用率从45%涨到70%整体推理提速的幅度肉眼可见比你去微调几层网络结构要管用得多。DeepGEMM的定位就是把这件事做透。它不追求什么花哨的模型想法就老老实实地把高维矩阵乘的调度、分块、访存、向量化这些底层功夫做到位然后把这个过程沉淀成一套可以复用的实现思路。1.2 DeepGEMM想要解决的核心矛盾直接说痛点GEMM的数学定义很简单C A * B C顶多再补个偏置什么的但真正在硬件上跑起来你会发现绝大多数时间根本不是在算而是在等数据。我举个实际点的例子。假设你要算一个 4096 x 4096 乘以 4096 x 4096 的矩阵乘原始的数据量差不多是 4096 * 4096 * 4 字节单看 A 和 B 就有 128MB 左右。GPU 的共享内存通常只有几十到一两百KB寄存器总量也就几百KB。这意味着你不可能把整个矩阵都塞进片上必须设计一套策略一块一块地搬、一块一块地算。这个过程一旦配合不好计算单元就在空转等你从显存把数据拉回来。所以DeepGEMM要解决的核心矛盾简单说就是怎么在有限的高速缓存空间里最大化数据的复用同时保证计算流水线始终是满的。这里面牵扯到分块尺寸的选择、访存顺序的设计、指令发射的编排还有特别关键的在什么样的情况下用什么样的硬件指令。这些后面我都会展开聊这里先把思路立起来。2. 动手之前先摸清GEMM的瓶颈在哪2.1 数学很简单难的是数据搬运我面试过不少人聊到GEMM优化第一反应都是这不就是三层for循环么然后就开始背cublas里面用了什么什么黑科技。但真让他说说瓶颈在哪很多人说不清楚。其实GEMM优化从头到尾都围绕一个核心指标计算访存比Arithmetic Intensity。来算一笔账。假设我们把一次乘加操作记作 2 FLOPS一次乘法和一次加法那么计算单位时间能发出去多少FLOPs由硬件主频决定但数据搬运的吞吐则由内存带宽决定。如果你在一个循环里不断从显存取两个数、做一个乘加、再存一个数哪怕理论算力再高你的实际吞吐也被带宽卡死了。这种情况叫memory-bound解决办法就一句话提高数据的复用率。怎么提呢就是分块。在矩阵乘里A矩阵的一个元素会被用来和B矩阵一整行的元素相乘B矩阵的一个元素也会被A矩阵一整列的元素相乘。你可别小看这个复用要是能把A的一个块留在寄存器里反复用这个块每次被加载后能参与的计算次数就是块尺寸的量级。这样就能把计算访存比从1次访问对应2次计算拉到1次访问对应几百次计算整体性能才会有质的飞跃。DeepGEMM的整套设计本质上是围绕如何把计算访存比拉高这一个目标服务的。你可以把共享内存和寄存器想象成一个小型的、速度极快的手工工台GEMM优化就是在做一道流水线管理题原料矩阵分块什么时候进料、什么时候加工、什么时候把成品部分和送回仓库每一步都要卡到点子上。2.2 什么样的算子才算快很多人喜欢拿“跑到多少TFLOPs”当噱头但单纯说绝对算力其实意义不大因为显卡的型号不同、精度不同、输入尺寸不同天花板就不一样。做GEMM优化一定要看一个更客观的指标利用率也就是实际FLOPS除以硬件理论峰值FLOPS。举个例子。某块显卡的理论FP32算力是20TFLOPs你实测算一个 2048 x 2048 的矩阵乘跑出来的实际算力是11TFLOPs那利用率就是55%。之前cublas这种商业库在新卡上优化到位的GEMMFP32的利用率大致能到80%~90%的区间。你要是能把自己的kernel优化到70%以上就已经算是相当不错的水平了。但有一点得提醒你利用率和矩阵规模有强相关。矩阵越大、分块的边际开销占比越小利用率越容易做高反过来几百乘几百的小矩阵往往瓶颈在启动开销和调度开销上利用率怎么做都上不去。所以DeepGEMM在做性能对比的时候我会固定几组典型尺寸大尺寸看访存调度能力小尺寸看启动和并行粒度设计不然只看一个大矩阵的测试结果很容易自我感觉良好。2.3 搞清楚硬件的脾气写kernel之前必须把手上的硬件规格翻出来看一遍。DeepGEMM的早期版本我主要跑在一块消费级GPU上它有几点特征很影响设计决策显存带宽决定了从显存往片上搬数据的极限速度。共享内存大小限制了单次分块能占用的片上缓冲上限。寄存器文件总量决定了每个线程能同时持有的累加器数量。线程束warp大小一般是32SIMD宽度为32意味着一个指令同时处理32个线程的数据。二级缓存行为哪些数据能留在L2里复用也会影响小矩阵的性能。我一开始犯过个错误直接把某个很强悍的kernel实现从高端显卡搬到普通显卡上结果性能掉得离谱。原因就是高端卡共享内存更大、带宽更宽原来需要四批次搬运的数据在新卡上要拆成八批次来搬调度开销翻倍。所以硬件规格不是看一眼就完事的后面每个关键参数的选择几乎都要回到这张表上来校验。3. 核心优化手段逐一拆解3.1 分块把大矩阵切成可以“一口吃下”的尺寸GEMM优化的第一步是把大矩阵拆成小矩阵块让一个block甚至一个线程负责一块子矩阵的计算。这一步叫Tiling我把它理解为给搬运工分活。你的缓存和寄存器就那么大把活分小了每个人都能在自己的一亩三分地里高效周转。分块尺寸怎么选我直接给一组我用下来比较稳的起步参数以FP32计算为例block内用 16x16 的微块micro-tile每个线程计算一个 4x4 的小结果块。这样算下来一个block有256个线程每个线程要持有16个寄存器做累加器加上加载数据用的临时寄存器整体寄存器压力还算可控。很多人一开始会犯一个毛病分块越大越好觉得这样数据复用率高。但实际跑起来你会发现分块太大会挤压线程并行度block数量变少调度器可能喂不饱所有计算单元。分块太小呢每个线程算的东西太少访存又压不下来。这是个天平问题必须用实测来选。这里放一个我摸索出来的经验值表格参数是在我先前的实验环境上调优出来的具体数字可以当作探索的起点不要当成定论矩阵规模微块大小线程块大小实测利用率FP32512 x 5128x8128线程约61%1024 x 102416x16256线程约72%2048 x 204816x16256线程约78%4096 x 409632x32512线程约82%注意这里不是越大越好。512规模的矩阵你硬上32x32微块反而因为块数太少、尾部效应太重利用率会往下掉。选分块参数说白了就是按你的实际业务尺寸来量体裁衣。3.2 双缓冲让搬运和计算重叠起来分块解决了数据能不能放下的问题但还有一个同样致命的问题搬运是耗时的。假设你一个分块的计算时间是10微秒从显存把下一块数据搬进共享内存要5微秒如果这两件事是串行的你的计算单元就有三分之一的时间在发呆。解决办法叫双缓冲Double Buffering也叫软件流水线。通俗解释一下你去餐厅吃饭一个服务员只顾着收桌收完才让下一个客人进来双缓冲就是安排了两个服务员一个在给当前客人点菜另一个已经在收拾隔壁桌了。对应到kernel里就是让计算单元处理当前这块数据的同时数据加载单元提前把下一块数据搬到另一个缓冲里。等当前计算结束无缝切换过去。实现上我会在共享内存里开两块缓冲区交替使用。代码草图是这样的逻辑// 伪代码展示双缓冲思路 __shared__ float bufA[2][BLOCK_SIZE]; __shared__ float bufB[2][BLOCK_SIZE]; // 预加载第一块 load(A, bufA[0]); load(B, bufB[0]); for (int k 0; k K_tiles; k) { int cur k % 2; int next (k 1) % 2; if (k 1 K_tiles) { // 提前预取下一块 load(A, bufA[next]); load(B, bufB[next]); } // 计算当前块 compute(bufA[cur], bufB[cur]); }这套逻辑看着不难但真正写的时候有几个坑预取指令不能阻塞当前计算必须走异步拷贝或者用独立的加载方式缓冲区切换时还要加同步防止早来的线程把下个块数据覆盖掉。我第一版双缓冲没加好同步结果同样的输入有时候算对有时候算错排查了一天才定位到是race condition。3.3 指令选择向量化与矩阵单元的取舍说完了调度层面的策略再往下一层就是指令级的选择。这是DeepGEMM这类项目最容易拉开差距的地方。GPU上做GEMM指令选择大致有这么几档普通乘加指令FFMA/FMA每个线程一个周期做一个乘加对应的是标量或简单向量计算。向量指令比如一次处理4个数、8个数的SIMD乘加能有效摊薄指令发射和读取的开销。矩阵指令专门为矩阵乘设计的硬件单元比如某些型号上的张量核心Tensor Core能一次完成一个小矩阵的乘加吞吐量比普通CUDA核心高一个数量级。以FP16为例我在支持矩阵指令的GPU上用16x16x16的低精度矩阵指令去算比手动展开FFMA大概能快2倍以上。但问题在于矩阵指令对数据布局有严格限制比如要求数据按特定内存布局对齐否则自动填充都没用。你用矩阵指令之前可能需要把数据先做一次内存转置这个转置开销甚至会吃掉指令带来的优势。我的取舍原则是大矩阵、批量大、精度要求又不苛刻的时候优先上矩阵指令中小矩阵、输入形状不规则的时候反而是精心编排的FFMA向量化版本更稳因为启动开销和布局转换开销更可控。3.4 融合策略减少中间结果落地的开销GEMM在真实模型里很少是单独跑的前面通常跟一个量化或变换后面跟一个偏置与激活函数。如果你一个算子一个算子地跑每一个都要把中间结果写回显存下一个算子再把它读回来一来一回就是两次显存访问。在memory-bound的算子链里这是白白扔掉的性能。DeepGEMM里我把偏置加法和激活函数直接并到GEMM的输出阶段去做。也就是说GEMM计算完一小块结果还在寄存器里热乎着的时候顺手把偏置加上、激活函数算掉然后直接把最终结果写回显存。这样中间结果压根不会经过显存。注意这种做法有一定危险性融合逻辑会吃掉更多的寄存器和指令周期如果激活函数特别复杂反而可能拖慢GEMM主循环。所以我的经验是只融合偏置加法、ReLU、GELU这类便宜的操作那种特别重的自定义算子不要硬融。实测中融合ReLU带来的收益在5%~15%之间具体看算子链的访存压力有多大。4. 实测数据与性能对比4.1 测试环境和基准配置没有实测数据的GEMM优化文章就是耍流氓。DeepGEMM的测试环境没有用特别顶级的卡就用一台普通的单卡机器测的这样跑出来的数据对大多数开发者更有参考意义。软件栈方面用的是CUDA 12.x编译选项里开了-O3和--use_fast_math对精度影响见后面说明所有计时都通过CUDA事件精确测核函数耗时排除CPU侧启动开销。基准我选了三个对照组第一组不用任何优化的朴素三重循环实现纯看最差是什么水平。第二组只做分块、不做双缓冲和指令优化的中等优化版。第三组DeepGEMM完整优化版包括双缓冲、向量化、输出融合。顺便说一句如果拿cublas出来比我们这些手写kernel基本是打不过的但它10年前可能并没有那么强这句话还得看具体版本与卡型号。所以手写优化kernel的定位不是取代官方库而是搞清楚性能天花板在哪里在特殊需求场景里做到可控和可定制。4.2 不同规模下的表现下面是DeepGEMM在 2048 x 2048 的FP32矩阵乘上的实测数据利用率已经按该卡的FP32理论峰值折算实现版本耗时实际FLOPS利用率朴素三重循环约42ms0.41 TFLOPS约2%中等优化分块约8.1ms2.12 TFLOPS约10%DeepGEMM完整版约1.05ms16.4 TFLOPS约78%从2%到78%差不多39倍的差距每一步优化都是成倍成倍地抠出来的。不过这里要提醒一点别盯着绝对数字看换一张卡、换一个精度这些数字全部要重测。利用率这个相对指标才是真正衡量kernel质量的标尺。再测了一组 512 x 512 的小矩阵完整版利用率就只剩55%左右了。原因前面说过小矩阵的启动开销占比大、并行块数少、分块尾部浪费明显。这部分优化的方向就和FP32大矩阵不一样得往减少block初始化开销和合理占用GPU多核方向去调。4.3 性能剖析时间究竟花在哪光看整体耗时是不够的还得知道时间都去哪儿了。DeepGEMM里我用性能分析工具抓过几次每次抓完都有新发现。举一个具体例子早期版本里我的分块和双缓冲都上了但实测L1缓存命中率只有67%意味着大量该在片上复用的数据还是跑去了L2甚至显存。进一步分析发现问题出在数据加载的访问模式上——相邻线程访问A矩阵的内存地址跨度太大导致缓存行没有被完全利用。改掉地址映射方式之后命中率拉到89%左右整体性能涨了11%。还有一次抓出来的是指令发射瓶颈编译器把一些小循环展开了但展开之后寄存器溢出严重部分数据被临时写到local memory里去了。这个问题特别隐蔽因为编译器不会报错但性能就是上不去。解决方式是手动减少循环展开的层级并把累加器数量控制到和寄存器文件匹配的水平。所以性能剖析在GEMM优化里的地位跟debug在软件开发里的地位差不多千万别跳过去。5. 常见问题与排查实录5.1 性能不升反降先查访存模式有朋友照着一个优化教程改了kernel结果比原来的版本还慢。我自己也经历过后来排查的结论是他用的分块尺寸不适合自己的显卡共享内存大小导致一个block的数据要分两批搬运每一批之间还有一次同步等待。如果你也遇到优化后性能不升反降优先查三件事第一共享内存有没有超量看看有没有触发隐式的local memory或者bank conflict第二线程束的访问模式是否连续有没有在共享内存里产生大量的bank冲突第三双缓冲到底有没有真正和计算重叠起来用性能分析工具看看流水线图如果加载和计算还是串行的肯定是同步或者异步搬移用错了。5.2 结果不对数值精度问题GEMM的精度问题在大矩阵里非常经典累加顺序不一样结果就不一样。FP32矩阵乘累加4096次浮点误差会随着次数累积你要是用--use_fast_math开了快速数学库FP32可能会被编译器偷偷改成近似计算误差更大。DeepGEMM的解决办法是把累加器提升到FP32精度如果原始数据是FP16输入部分累加也可以考虑用FP32的累加器这样正确率基本能对齐pytorch的默认输出。另外有一种情况是共享内存里出现未初始化的尾数。分块尺寸不整除矩阵边界的时候越界的部分必须用0填充否则那些垃圾值参与计算结果就会突然变NaN。这个bug的表现是小矩阵正常大矩阵偶发出错排查起来头大。处理边界的时候一定要写清填充逻辑并且建议在kernel入口做一次断言检查。5.3 输入规模特殊导致的分支问题实际业务里矩阵不总是规规整整的幂次方尺寸。比如序列长度是128或者300特征维度是768等等。这种不整除的尺寸会让分块尾部出现大量闲置线程性能损耗严重。DeepGEMM处理这种情况会用边界感知的分块策略——正常块走快速路径尾部块走慢速但正确的路径用一个if (tile_in_range)判断分叉代价是多一点分支指令但能保证整体性能不崩。踩过一次坑早期图省事把所有block都按最大边界处理显存分配超大小矩阵跑起来也浪费。后来改成动态计算实际需要的block数量启动开销直接少了30%左右。6. 做GEMM优化的一点个人体会讲了这么多方法和数据最后说点掏心窝的话。GEMM优化这个方向最磨人的地方不是它有多难而是它特别琐碎。每一个优化手段单独看都好理解但堆在一起的时候互相之间的依赖和冲突就会冒出来。你调好了访存发现寄存器又溢出了你压了寄存器发现指令流水又空了。它不是一个线性叠加的过程而是一个需要反复权衡、反复实测的循环。我实际做下来最大的感受是性能优化的每一步都建立在测量之上。任何理论分析、任何纸面上的复杂度推演最后都必须落到性能分析工具和精确计时上。DeepGEMM的成果不是我一次性设计出来的而是跑了不知道多少轮profiling、改了多少个版本才磨出来的。另外有个建议如果你是第一次接触GEMM优化别一上来就上矩阵指令、双缓冲这些重武器。先写一版朴素实现然后只加一个优化手段比如分块看看性能提升了多少再加一个比如双缓冲再测一次。保持每一步都是单一变量你才能真正积累出直觉知道哪个手段在什么场景下值多少收益。我刚开始就是贪快一口气把所有优化全写进去结果出了问题根本不知道是哪一步引入的。最后再分享一个小技巧优化GEMM的时候尽量拿一个和业务真正匹配的矩阵尺寸做基准别只用标准的2048乘2048。很多项目实际跑的是768、1024、1280这种半规整尺寸在这些尺寸上手工调整分块参数带来的收益往往比在中规中矩的大矩阵上跑出的漂亮数字更有实际价值。DeepGEMM后续的迭代方向我也在往这种非规整尺寸针对性调优上走。