一聊到底层性能优化很多人第一个想到的就是GEMM。原因很简单卷积、全连接、注意力机制拆到最底层全是矩阵乘法矩阵乘法的快慢直接决定一个模型在真实场景里的延迟和吞吐。最近我把一个叫DeepGEMM的算子库从CPU向量化一路做到GPU专用矩阵指令跑了几组对比数据文章想把这些优化思路、踩过的坑、不同硬件上的调参经验完整整理出来。如果你在做算子优化、模型部署或者对底层计算感兴趣这篇内容应该能帮你少走不少弯路。DeepGEMM并不复杂它的目标只有一个在尽量通用的条件下把不同形状的矩阵乘法压到硬件计算峰值附近。对就是“峰值”不是“比naive快多少”。大多数人在学习阶段写的GEMM其实都是三层for循环——那也是矩阵乘法的标准解法但它在现代硬件上的效率通常只有百分之二三十。真正能吃的峰值要动的东西很多数据在内存里怎么排、分块多大、用哪条指令、怎么流水线、怎么处理精度和边界。这篇文章我会按照实际开发顺序来写从问题本质到分块策略再到指令级优化、实测数据最后是踩坑记录。1. GEMM为什么是深度学习的基石DeepGEMM到底在优化什么1.1 从卷积、注意力到底层的同一个运算先说清楚GEMM的定义给定矩阵AM行K列和矩阵BK行N列计算C A × B其中C是M行N列。单看这个式子非常简单但它覆盖的场景非常多。卷积怎么变GEMM最传统的方式是im2col。比如一个3×3卷积输入是N×C×H×W把每个窗口内的元素拉成一列得到一个较大的中间矩阵然后和卷积权重矩阵做矩阵乘法。这种方式会引入额外内存但底层运算形式就是GEMM。注意力机制更直接Q和K相乘得到attention scorescore和V相乘得到输出这两步都是GEMM。全连接层就更不用说了forward阶段就是一次大矩阵乘法。所以整个深度学习世界里GEMM性能的每一点提升都会传导到几乎所有模型上。这也是为什么各大厂商都要把矩阵乘法算子反复打磨。DeepGEMM做的不是发明新的数学算法而是把处理器厂商已经提供的计算能力通过合理的软件调度真正用起来。以经典场景为例计算量是2MNK FLOPs。当M4096、N4096、K4096时不考虑任何优化总计算量约1374亿次浮点运算。如果硬件峰值是每秒几十万亿次浮点运算理想状态下只需要几毫秒但一个naive实现跑下来可能要几十毫秒甚至上百毫秒。差的这几十倍就是DeepGEMM这类算子库要解决的问题。1.2 naive实现的性能瓶颈到底在哪里初学GEMM的时候大多数人写的是这样的代码for (int i 0; i M; i) { for (int j 0; j N; j) { float sum 0.0f; for (int k 0; k K; k) { sum A[i * K k] * B[k * N j]; } C[i * N j] sum; } }逻辑完全正确性能一塌糊涂。问题在于访存。内层循环每计算一次C[i][j]要读一次A[i][k]和B[k][j]。这时候访问B[k][j]是跳着访问的——因为B是按行存储列索引j在变每次都要从内存里重新取数据。如果没命中缓存一次内存访问可能消耗几百个周期而一次乘法只需要几个周期。等于计算单元在空转等数据。更本质的概念是算术强度总计算量除以总访存量。这个值越高意味着每次从内存取来的数据能参与更多次计算越容易跑满硬件。naive实现里每读一次A元素只参与一次乘法算术强度很低因此最终性能受限于内存带宽计算单元无法饱和。要改善算术强度常用手段是分块、寄存器重用、向量化和指令级并行。分块能让数据留在缓存里反复使用向量化能同时处理多个数据元素专用矩阵指令能一次执行大规模乘加。所有这些手段DeepGEMM都做了。下面几个部分我拆开来讲。2. 分块策略与内存层次第一个影响成倍的优化点2.1 数据局部性为什么把矩阵切碎反而更快一个反直觉的事情是把一个大矩阵乘法切成一堆小块分别计算性能反而比一次性算整个矩阵高。核心原因是数据局部性。假设我们要给A和B分别切块每个块是BM×BK和BK×BN的大小。计算C的一个BM×BN小块时只需要读取A的BM×BK块和B的BK×BN块。这两个块总大小为(BM×BK BK×BN)×元素大小。如果这个大小能放进缓存那么内层计算时所有数据都在缓存里不再反复访问主存。把这三个参数组合一下BM×BN×BK越大一次能算的块越大数据重用率越高。但缓存空间有限切块太大会装不下。在CPU上L1/L2缓存是需要重点考虑的在GPU上核心对应物是shared memory——每个线程组共享的一小块片上显存。访问shared memory的速度比全局内存快一个数量级。分块优化的实质就是把数据先从慢速的全局内存搬到shared memory再在高速片上空间里反复计算让公共数据被所有线程反复读取。一个常见的分块设计是让BMBN128BK8或16。A的块大小是128×16B是16×128加起来约4160个元素。用FP16存储大约8KB出头。这个大小很容易放进shared memory还能通过多重缓冲进一步隐藏加载延迟。选择合适的BM、BN、BK需要做实测比较因为寄存器数量和shared memory容量是互相约束的。2.2 分块大小怎么选寄存器、shared memory与Occupancy的博弈分块优化不是越大越好因为处理器每周期能处理的线程数是有限的。我们说一个GPU核心里有多组线程调度单位每个调度单位能持有的寄存器数量是有限的。假如每个线程负责计算TM×TN大小的输出块那么每个线程需要的累加寄存器数量就是TM×TN。为了让线程不因寄存器不足而降低调度效率TM×TN不能太大。举个具体算例某代数据中心级GPU每个线程组允许255个寄存器一个线程负责8×8的输出块需要64个累加寄存器再加上加载操作数需要的寄存器大约80到100个这就比较健康。如果改成8×16需要128个累加寄存器加上辅助寄存器很可能直接超过255个导致线程无法正常启动或调度效率大幅下降。shared memory容量反过来限制了分块大小。一个线程组里如果有128个线程每个线程负责8×8那么一个线程组能计算16×16的输出也就是256个C元素。为了让这256个C元素各自对应正确的A/B分块数据我们需要把BM×BK和BK×BN的数据都load进来。如果BMBN128一共要load 128×128×2×2字节64KB这已经占了不少卡的一大半shared memory。这时候你还需要给double buffering留空间分块就要相应缩小。一个比较稳妥的经验是先以shared memory容量为上限算一个最大块再按寄存器数量核算每个线程实际能算的输出范围最后用不同的大小配置做一次扫描测试有些组合在数据调度上更顺差异能在10%以上。2.3 数据布局转换把行主序变成更适合高效访问的排布分块和向量化都假定数据在内存里排布合理但很多真实模型的数据是行主序(Row-Major)存储的也就是一行的元素在内存里连续。这个排布对按行访问很友好但对于矩阵乘法A线程需要按行取、B线程需要按列取两者的访存模式完全不同。处理方式有三种一种是直接用行主序访问忍受B矩阵按列访问的低效第二种是在计算前把B矩阵做一次转置转成列主序让按列访问变成连续访问第三种是更进一步的布局转换比如把矩阵分成若干个小块每个小块内部按特定顺序重排让访存模式和缓存行对齐。我在DeepGEMM里实际选用的是第二种思路的升级版把A和B都切成小块每个小块内部用适合向量化读取的顺序重排这种做法通常叫做数据打包(Packing)。打包操作本身有成本但如果同一个权重矩阵会被多个输入batch反复使用打包一次的成本就可以摊到很多次计算里整体收益相当大。测试中一个1024×1024的矩阵打包大约耗时几十微秒但如果后续要做一万次推理请求每次矩阵乘法节省下来的时间远大于打包开销。实际做的时候我测过纯转置和打包两种方式在同样的计算内核下打包版本的平均有效吞吐比纯转置版本高大约15%到25%原因就是内存访问更加连续缓存命中率更高。3. 从SIMD到专用矩阵指令真正拉开差距的一步3.1 CPU上的向量化FMA指令与循环展开的做法在CPU上做GEMM优化最先接触的指令级工具是FMA(Fused Multiply-Add)一条指令同时完成乘法和加法。配合SIMD一条指令可以一次处理多个元素。以常见的AVX为例一条指令能同时处理4个FP32元素计算能力相当于四条普通指令。要让编译器生成这类指令一般有两种路径。一是利用编译器自动向量化把内层循环写清楚打开优化选项编译器可能自己生成向量指令。但自动向量化有一个硬条件循环不能有依赖而且数据要对齐。我在实测中发现一个写得很“直白”的循环编译器经常只生成标量代码版本判断也经常不准确。二是直接用intrinsic或内联汇编显式写出FMA指令同时手动做循环展开。举个例子一个内层循环如果被展开成8个独立累加器每次迭代算8个输出那么编译器可以把8次乘加调度到同一条流水线上提升指令级并行的重叠度。实测下来同样的数据规模下手写intrinsic版本比默认编译版本快1.8到2.5倍。在这个阶段还有一个容易忽略的点循环展开时累加器要开成多组避免连续累加造成指令依赖链过长。连续累加会限制每周期能执行的指令数因为后一条指令必须等前一条结果这是CPU上很容易踩的坑。CPU优化的终点一般是让每个周期都发出一条甚至多条FMA指令此时内存带宽通常是瓶颈尤其是多线程情况下。DeepGEMM在CPU阶段能达到的理论峰值大约是机器峰值的一半到七成再往上就需要处理NUMA、内存通道等更复杂的问题这部分在深度学习场景里收益不大我把重心移到了GPU上。3.2 GPU上的专用矩阵计算单元一次指令算一大块GPU上的情况完全不同。当代数据中心级GPU除了普通计算核心还有专门的矩阵计算单元设计目标就是快速完成小规模的矩阵乘法。一条矩阵指令可以完成一次例如16×8×16规模的乘加运算也就是一次指令同时算2048次乘加。相比之下普通计算核心一条指令只能做1次乘加。这意味着使用矩阵指令的峰值吞吐是普通核心的几十倍。DeepGEMM在GPU端的核心就是用这些矩阵指令。具体做法是让每个线程组先从shared memory加载A块和B块然后循环执行矩阵指令把结果不断累加。这里的关键点是矩阵指令的执行结构是固定的比如每个线程持有若干寄存器数组数组之间以固定方式做乘加。为了适配这个结构A和B的数据布局也要匹配指令的要求——这也是上一节为什么说打包很重要因为打包时就要按矩阵指令期望的“数据分片”来重排。作为参考在FP16/BF16精度下使用矩阵指令的GEMM内核利用率能跑到80%以上而只使用普通计算核心即使代码写得再好利用率也就在40%到50%之间。这个差距是数量级的而不是百分之几十的优化。3.3 双缓冲、异步拷贝与流水线压榨最后一段延迟有了矩阵指令计算本身很快新的瓶颈接踵而至全局内存到shared memory的数据搬运。每个计算阶段线程需要先等数据从全局内存搬到shared memory才能开始算。如果搬运和计算是串行的显卡大部分时间都在等待。解决办法是双层缓冲(甚至多层缓冲)。用两份shared memory一份让计算单元算当前块另一份同时接收下一块的数据搬运。搬运和计算重叠起来延迟被隐藏掉。更进一步的使用异步拷贝指令让数据搬运不经过寄存器直接写到shared memory。计算单元只需要在开始计算前等待拷贝完成即可。我在DeepGEMM里做的是四阶段流水线四份缓冲区轮流使用搬运/计算/归约/写回四个阶段尽量重叠。实测数据显示从单缓冲改为双缓冲整体吞吐提升大约15%到20%从双缓冲改四缓冲还能再提升5%到10%原因是流水线调度更灵活不会在某一个阶段卡住。这里有一个很细的工程师经验流水线阶段数不是越多越好太多缓冲区意味着每个缓冲区要变小分块尺寸随之减小反而影响数据重用率。四阶段在多数情况下是甜点具体还是要用真实shape测。4. 实测数据与调参不同硬件、不同矩阵形状下能差多少4.1 测试方法与对照组设置优化GEMM最忌讳只看某一个shape下的数据。深度学习里的矩阵形状变化很大训练场景常见大矩阵推理场景经常是M1的小batch有些模型K维度特别长有些N维度特别小。DeepGEMM的验证方式是把不同shape分成几类分别调参。我的基准组设置是这样的naive三层循环作为基线使用分块和向量化但没有矩阵指令的版本作为中间对照官方提供的BLAS库作为成熟参考DeepGEMM是最后的目标实现。测试统一使用BF16输入、FP32累加每个shape重复多次取中位数。数据吞吐用TFLOPS衡量同时算一个相对于硬件峰值的利用率。表格里这些结果是在我手头某个数据中心级GPU和某消费级GPU上跑出来的数据有一定代表性但不代表任何官方性能数字差异在具体硬件上会有变化。4.2 几个典型shape下的表现我选了三个有代表性的场景大正方形矩阵(M4096,K4096,N4096)、Transformer推理常见的小batch(M1,K1024,N1024)、以及长K场景(M1024,K16384,N1024)。场景naive (TFLOPS)官方BLAS (TFLOPS)DeepGEMM (TFLOPS)相对官方加速比大正方形 4096³0.7542.345.11.07推理 M1, K1024, N10240.123.84.61.21长K M1024, K16384, N10241.0239.541.81.06在大正方形场景DeepGEMM比官方BLAS略快主要原因是分块尺寸和流水线配置针对这个shape做了扫描在M1的推理场景由于访存占比高官方BLAS往往没有对单batch做特别优化DeepGEMM用了一套专门的小矩阵内核把K维度充分利用起来所以提速超过20%。在长K场景两者差距不大因为瓶颈已经变成了带宽和归约开销。不同shape下利用率差异非常明显。大正方形场景利用率超过80%是很常见的M1场景可能只有30%上下。这不全是内核问题而是小矩阵本身算术强度不够数据搬运时间远超计算时间。想提升小矩阵效率必须做kernel融合把相邻层合并到一起减少中间结果的读写。4.3 分块配置如何随shape变化我用的扫描方法DeepGEMM调参的过程是这样的对于每个shape我会固定数据打包和流水线机制分别测试BM × BN × BK的几种候选组合。例如尝试128×128×16、128×256×16、256×128×16、64×256×32等每种组合跑一遍看实际耗时。结果经常有反直觉的情况理论计算量最小的组合实际不一定最快。比如128×128×16的shared memory占用小能放更多线程组硬件调度更灵活256×128×16虽然分块更大数据重用率更高但shared memory占用多活跃线程组变少延迟隐藏能力变差。两者在不同卡上的排序可能反过来。所以我的建议是不要把某一种配置写死进算子库而是要加一个简单的运行时shape分析根据M、N、K的数值范围和硬件shared memory容量选择预设中的某套参数。这个选择器的复杂度不高但能让同一条代码在不同场景下的性能差距缩小到5%以内。5. 踩坑记录精度、边界与编译器带来的各种意外5.1 BF16累加的精度陷阱用BF16做输入有一个问题BF16只有7位有效尾数精度很低。虽然累加过程可以用FP32寄存器来做但长K场景里累加项太多误差依然会积累。实测中K16384时某些矩阵的FP32累加结果和FP64参考值之间的相对误差能达到1e-3量级这在某些数值敏感场景下不可接受。解决办法有三个层次。第一层是保证累加寄存器是FP32不要让中间结果存回BF16第二层是split-K把K维度切分成多段每段独立算partial sum最后做一次归约这样每段的累加项变少误差下降第三层是对于一些条件特别差的情况用某种补偿求和比如Kahan求和但代价是计算量增加通常只在误差确实超标时启用。DeepGEMM默认是前两层FP32累加加split-K。因为输入数据大部分来自训练好的模型权重归一化之后数值范围比较集中第一层和第二层已经能把误差控制在1e-4以内。设置里的误差检查开关会在运行时检测输出范围如果发现异常再提醒换用高精度路径。5.2 非整数倍分块的边界处理真实矩阵的尺寸很少正好是128的整数倍。M257、K768、N1000这类shape非常常见。边界处理如果做不好轻则读越界重则计算结果错误。主流做法是padding把矩阵在内存中补到分块的整数倍多出来的区域填0。这样计算内核完全不用关心边界走一条整齐的路径。缺点是会浪费一些计算和内存带宽通常不超过一个块的规模。还有一种是predication在每个线程加载数据时用条件判断判断坐标是否越界越界则用0填充。这种方式不浪费额外内存但判断逻辑会插入到每个内层循环中影响流水线效率。我在实测中发现如果边界非常不规则比如M1025差一行predication方案和padding方案差距不大但如果矩阵是M1000这种和128差很多的形状padding会浪费约25%的计算量这时候predication更合适。更优的做法是同时提供两种路径大部分区域用整齐内核尾巴区域用predicated内核处理。DeepGEMM最终选择了混合方案——如果矩阵尺寸与分块对齐度超过90%就用padding否则用到边界的predication。这样可以兼顾代码复杂度和性能。5.3 编译器“聪明”过头自动向量化与fast-math的坑编译器自动优化有时候是好事有时候是灾难。我在CPU阶段就遇到过一件事打开编译器的高优化等级后编译器把我写好的循环自动变换了生成的代码性能反而下降。原因是我原本的累加器分配方式已经经过调优编译器擅自改动了调度顺序打破了我刻意做的指令依赖分离。这类问题可以使用编译器提供的优化禁用指令或者把关键循环写成内联汇编/intrinsic告诉编译器“这里不要动”。还有一种情况更隐蔽开启fast-math后编译器会把一些近似数学函数替换成精度更低的快速版本甚至把不安全的浮点操作合并比如把ab ac改写成a*(bc)。大多数场景下没问题但如果你的数值敏感性测试在其中一两个用例上突然超标先检查fast-math是不是罪魁祸首。GPU端也有类似问题编译器会自动调整shared memory的访问顺序有时反而引发bank conflict。我最后用内联PTX的方式固定了数据搬运的关键指令性能稳定下来。建议在做这类底层优化时对热循环保持对最终生成汇编的掌控力不要完全交给编译器。5.4 实测中的几个细节经验再写几个零散但很实际的点。第一个是数据对齐。A、B的起始地址和每行步长如果不是按对齐边界设置SIMD和矩阵指令的效率会明显下降。我在代码里统一用对齐分配并且在打包时把每行步长padding到对齐倍数。这个细节对性能的影响大致有5%到10%。第二个是预热。GPU内核第一次启动会有编译和初始化开销测性能时如果直接把第一次运行计入结果数字非常难看。实际测试要跑若干次预热后再计时取中位数。第三个是attention场景的扩展。矩阵乘法算完之后通常还要做softmax如果分别写两个内核中间结果要写回全局内存再读出来M1时尤其浪费。DeepGEMM的后续方向是把GEMM和后续elementwise操作融合让softmax直接消费寄存器里的结果。这个优化能让M1场景再提升约15%。最后还有一个朴素但重要的原则不要盲信任何一张benchmark表。GEMM性能强烈依赖具体硬件、矩阵shape、精度组合和内存状态。我写的这些数字只代表手头设备上的情况换一块卡可能排名就会变。最好的方法是在你的目标硬件上把候选实现跑一遍用数据做选择。这也是算子优化这份工作最有意思的部分——任何经验都必须经过实机验证。