从GEMM到DeepGEMM：CPU向量化与GPU矩阵指令级优化实践

一聊到底层性能优化，很多人第一个想到的就是GEMM。原因很简单：卷积、全连接、注意力机制，拆到最底层全是矩阵乘法；矩阵乘法的快慢，直接决定一个模型在真实场景里的延迟和吞吐。最近我把一个叫DeepGEMM的算子库从CPU向量化一路做到GPU专用矩阵指令，跑了几组对比数据，文章想把这些优化思路、踩过的坑、不同硬件上的调参经验完整整理出来。如果你在做算子优化、模型部署，或者对底层计算感兴趣，这篇内容应该能帮你少走不少弯路。

DeepGEMM并不复杂，它的目标只有一个：在尽量通用的条件下，把不同形状的矩阵乘法压到硬件计算峰值附近。对，就是“峰值”，不是“比naive快多少”。大多数人在学习阶段写的GEMM，其实都是三层for循环——那也是矩阵乘法的标准解法，但它在现代硬件上的效率通常只有百分之二三十。真正能吃的峰值要动的东西很多：数据在内存里怎么排、分块多大、用哪条指令、怎么流水线、怎么处理精度和边界。这篇文章我会按照实际开发顺序来写，从问题本质到分块策略，再到指令级优化、实测数据，最后是踩坑记录。

1. GEMM为什么是深度学习的基石，DeepGEMM到底在优化什么

1.1 从卷积、注意力到底层的同一个运算

先说清楚GEMM的定义：给定矩阵A（M行K列）和矩阵B（K行N列），计算C = A × B，其中C是M行N列。单看这个式子非常简单，但它覆盖的场景非常多。

卷积怎么变GEMM？最传统的方式是im2col。比如一个3×3卷积，输入是N×C×H×W，把每个窗口内的元素拉成一列，得到一个较大的中间矩阵，然后和卷积权重矩阵做矩阵乘法。这种方式会引入额外内存，但底层运算形式就是GEMM。注意力机制更直接：Q和K相乘得到attention score，score和V相乘得到输出，这两步都是GEMM。全连接层就更不用说了，forward阶段就是一次大矩阵乘法。

所以整个深度学习世界里，GEMM性能的每一点提升，都会传导到几乎所有模型上。这也是为什么各大厂商都要把矩阵乘法算子反复打磨。DeepGEMM做的不是发明新的数学算法，而是把处理器厂商已经提供的计算能力，通过合理的软件调度真正用起来。

以经典场景为例，计算量是2MNK FLOPs。当M=4096、N=4096、K=4096时，不考虑任何优化，总计算量约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，再在高速片上空间里反复计算，让公共数据被所有线程反复读取。

一个常见的分块设计是让BM=BN=128，BK=8或16。A的块大小是128×16，B是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进来。如果BM=BN=128，一共要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下的数据。深度学习里的矩阵形状变化很大：训练场景常见大矩阵，推理场景经常是M=1的小batch；有些模型K维度特别长，有些N维度特别小。DeepGEMM的验证方式是把不同shape分成几类，分别调参。

我的基准组设置是这样的：naive三层循环作为基线；使用分块和向量化但没有矩阵指令的版本作为中间对照；官方提供的BLAS库作为成熟参考；DeepGEMM是最后的目标实现。测试统一使用BF16输入、FP32累加，每个shape重复多次取中位数。数据吞吐用TFLOPS衡量，同时算一个相对于硬件峰值的利用率。

表格里这些结果是在我手头某个数据中心级GPU和某消费级GPU上跑出来的，数据有一定代表性，但不代表任何官方性能数字，差异在具体硬件上会有变化。

4.2 几个典型shape下的表现

我选了三个有代表性的场景：大正方形矩阵(M=4096,K=4096,N=4096)、Transformer推理常见的小batch(M=1,K=1024,N=1024)、以及长K场景(M=1024,K=16384,N=1024)。

在大正方形场景，DeepGEMM比官方BLAS略快，主要原因是分块尺寸和流水线配置针对这个shape做了扫描；在M=1的推理场景，由于访存占比高，官方BLAS往往没有对单batch做特别优化，DeepGEMM用了一套专门的小矩阵内核，把K维度充分利用起来，所以提速超过20%。在长K场景，两者差距不大，因为瓶颈已经变成了带宽和归约开销。

不同shape下，利用率差异非常明显。大正方形场景利用率超过80%是很常见的，M=1场景可能只有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场景里，累加项太多，误差依然会积累。实测中，K=16384时，某些矩阵的FP32累加结果和FP64参考值之间的相对误差能达到1e-3量级，这在某些数值敏感场景下不可接受。

解决办法有三个层次。第一层是保证累加寄存器是FP32，不要让中间结果存回BF16；第二层是split-K，把K维度切分成多段，每段独立算partial sum，最后做一次归约，这样每段的累加项变少，误差下降；第三层是对于一些条件特别差的情况，用某种补偿求和，比如Kahan求和，但代价是计算量增加，通常只在误差确实超标时启用。

DeepGEMM默认是前两层：FP32累加加split-K。因为输入数据大部分来自训练好的模型权重，归一化之后数值范围比较集中，第一层和第二层已经能把误差控制在1e-4以内。设置里的误差检查开关会在运行时检测输出范围，如果发现异常，再提醒换用高精度路径。

5.2 非整数倍分块的边界处理

真实矩阵的尺寸很少正好是128的整数倍。M=257、K=768、N=1000这类shape非常常见。边界处理如果做不好，轻则读越界，重则计算结果错误。

主流做法是padding：把矩阵在内存中补到分块的整数倍，多出来的区域填0。这样计算内核完全不用关心边界，走一条整齐的路径。缺点是会浪费一些计算和内存带宽，通常不超过一个块的规模。

还有一种是predication：在每个线程加载数据时，用条件判断判断坐标是否越界，越界则用0填充。这种方式不浪费额外内存，但判断逻辑会插入到每个内层循环中，影响流水线效率。我在实测中发现，如果边界非常不规则（比如M=1025，差一行），predication方案和padding方案差距不大；但如果矩阵是M=1000这种和128差很多的形状，padding会浪费约25%的计算量，这时候predication更合适。

更优的做法是同时提供两种路径：大部分区域用整齐内核，尾巴区域用predicated内核处理。DeepGEMM最终选择了混合方案——如果矩阵尺寸与分块对齐度超过90%，就用padding；否则用到边界的predication。这样可以兼顾代码复杂度和性能。

5.3 编译器“聪明”过头：自动向量化与fast-math的坑

编译器自动优化有时候是好事，有时候是灾难。我在CPU阶段就遇到过一件事：打开编译器的高优化等级后，编译器把我写好的循环自动变换了，生成的代码性能反而下降。原因是我原本的累加器分配方式已经经过调优，编译器擅自改动了调度顺序，打破了我刻意做的指令依赖分离。

这类问题可以使用编译器提供的优化禁用指令，或者把关键循环写成内联汇编/intrinsic，告诉编译器“这里不要动”。还有一种情况更隐蔽：开启fast-math后，编译器会把一些近似数学函数替换成精度更低的快速版本，甚至把不安全的浮点操作合并，比如把ab + ac改写成a*(b+c)。大多数场景下没问题，但如果你的数值敏感性测试在其中一两个用例上突然超标，先检查fast-math是不是罪魁祸首。

GPU端也有类似问题：编译器会自动调整shared memory的访问顺序，有时反而引发bank conflict。我最后用内联PTX的方式固定了数据搬运的关键指令，性能稳定下来。建议在做这类底层优化时，对热循环保持对最终生成汇编的掌控力，不要完全交给编译器。

5.4 实测中的几个细节经验

再写几个零散但很实际的点。

第一个是数据对齐。A、B的起始地址和每行步长如果不是按对齐边界设置，SIMD和矩阵指令的效率会明显下降。我在代码里统一用对齐分配，并且在打包时把每行步长padding到对齐倍数。这个细节对性能的影响大致有5%到10%。

第二个是预热。GPU内核第一次启动会有编译和初始化开销，测性能时如果直接把第一次运行计入结果，数字非常难看。实际测试要跑若干次预热后再计时，取中位数。

第三个是attention场景的扩展。矩阵乘法算完之后通常还要做softmax，如果分别写两个内核，中间结果要写回全局内存再读出来，M=1时尤其浪费。DeepGEMM的后续方向是把GEMM和后续elementwise操作融合，让softmax直接消费寄存器里的结果。这个优化能让M=1场景再提升约15%。

最后还有一个朴素但重要的原则：不要盲信任何一张benchmark表。GEMM性能强烈依赖具体硬件、矩阵shape、精度组合和内存状态。我写的这些数字只代表手头设备上的情况，换一块卡可能排名就会变。最好的方法是在你的目标硬件上把候选实现跑一遍，用数据做选择。这也是算子优化这份工作最有意思的部分——任何经验都必须经过实机验证。