手写GPU GEMM内核：从朴素实现到接近硬件峰值

“GEMM”这个缩写，凡是碰过深度学习框架源码或者写过推理引擎的人都不会陌生：General Matrix Multiplication，通用矩阵乘法。我在开发模拟项目“DeepGEMM”时，想做的其实是一件很朴素的事：不依赖任何第三方库，手写一个能在主流GPU上把矩阵乘法性能压榨到接近硬件理论峰值的内核，同时把整个优化思路、踩坑过程记录下来。这篇文章就是基于那次实践整理出来的完整复盘，适合正在学习CUDA内核优化、或者工作中需要手工调GEMM性能的开发者阅读。我会从“为什么GEMM值得折腾”讲起，一路拆到分块、双缓冲、张量指令这些底层细节，最后给出可直接参考的代码骨架和排障清单。

1. 为什么非要折腾GEMM

不少朋友看到“手写GEMM”第一反应是：调库不就行了吗？现实中确实如此，大多数场景下直接调用现成的高性能数学库就够了。但一旦你开始做推理引擎定制、模型算子融合，或者要支持某些库不友好的形状，事情就变了。

1.1 GEMM在深度学习里的实际面孔

先想一个问题：一个Transformer模型从输入到输出，计算量到底花在哪？我做过一次粗糙的统计，拿一个中等规模的GPT风格模型为例，前向推理里超过90%的乘加运算都落在各种矩阵乘法上。注意力里的QKV投影是GEMM，注意力输出投影是GEMM，FFN的两个全连接层还是GEMM。就连看起来毫无关系的卷积操作，在绝大多数深度学习框架里也会被转化成GEMM来执行，也就是所谓“隐式GEMM”的做法。

换句话说，GEMM就是深度学习推理引擎的地基。地基没打好，上层优化做得再漂亮，最终也跑不快。那些优秀的推理引擎之所以能在各种硬件上保持很高利用率，靠的往往就是一套精心优化的GEMM内核，再加上合适的算子调度策略。这也是我决定做DeepGEMM这个项目的根本原因：把这块地基踩透，很多上层问题会豁然开朗。

1.2 GEMM的数学本质和三个关键特征

GEMM的定义非常简单：给定矩阵A，维度为M×K，矩阵B，维度为K×N，计算C = A × B + C，最终C的维度为M×N。这里的“+C”叫累加（accumulate），是很多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; } }

代码看起来短，但里面藏着三个对性能至关重要的特征。第一，循环体的核心操作是一次浮点乘加（FMA），一共执行2×M×N×K次浮点运算。第二，乘法操作本身虽然简单，但数据访问在存储层面很复杂：A的访问是行连续，B的访问是列连续。换句话说，内层循环里B的每次访问都需要跨一整行，这会直接打爆缓存。第三，整个计算过程完全可并行：C矩阵的每个元素之间互不依赖，天然适合大规模并行计算。这三个特征，决定了GEMM优化的主战场不是“如何算”，而是“如何让数据以最快速度送到计算单元里”。

1.3 为什么现成的库函数不够用

我经常遇到有人说“直接用cuBLAS不就行了吗，性能还那么好”。这话在通用场景下没错，但放到实际工程里，有几个现实问题。第一，不是所有形状都能吃到最优内核。很多库对某种形状（比如超大矩阵）优化得非常好，但对细长矩阵、宽短矩阵、或者batch size很小的情况，性能会明显打折。第二，当你做算子融合时，可能需要把GEMM的中间结果留在寄存器或者共享内存里直接做下一个操作，这时候库函数提供的黑盒接口就帮不上忙了。第三，有些场景需要支持自定义的数据布局、mask、或者稀疏特性，通用库很难都覆盖。更重要的是，手写一个GEMM是学习GPU计算优化的最佳入口：它够简单、够核心、优化手段丰富，又没有复杂的算法分支。做完一轮GEMM优化，你对访存、并行、指令、流水线这些概念的理解，会有一个质的飞跃。

2. 为什么朴素三重循环跑不出性能

我先在GPU上把上面那段伪代码直接翻译成内核跑了一遍，结果非常“感人”：性能只有理论峰值的几个百分点。这里面的差距，就是GEMM优化要解决的核心矛盾：存储墙。

2.1 计算强度：一次乘加要搬多少个字节

我们把每个浮点乘加拆开看。A矩阵的一个元素和B矩阵的一个元素相乘，加到累加器上。按照FP32精度来算，每个元素是4字节。也就是说，在内层循环里做一次乘加，至少需要从内存里读8个字节（A元素4字节加B元素4字节）。如果数据没有被缓存命中，这些字节还得从全局内存里搬。而现代GPU一个计算单元（SM）一个周期可以执行很多次浮点运算，假设某个SM每个周期能吞吐128次FP32乘加，而它从全局内存读取的带宽摊下来每个周期可能只有几个字节：

• 计算需求：128 FLOP/周期
• 内存供给：假设一个周期内带宽只能提供约4字节数据，配合FP32乘加的计算强度需求8字节/FLOP，完全喂不饱。
• 结论：绝大部分计算单元都在等待数据，利用率自然惨不忍睹。

用Roofline模型的话来说，这个朴素循环的计算强度（arithmetic intensity）太低了。对FP32来说，每做一次乘加要取两个数，理论计算强度大约只有0.25 FLOP/Byte，连硬件平衡点（往往是几百甚至上千FLOP/Byte）的零头都不到。想要提高性能，最好的办法就是提高数据复用：让一个从内存载入的数据被多次使用。

2.2 三种存储层级与访问代价

GPU里数据的存储层级大概是这个关系：寄存器（最快）→ 共享内存（片上，快）→ L2缓存（片外统一缓存）→ 全局内存（最慢，直接面向显存）。速度差异有多大呢？可以这么理解：如果寄存器访问是1纳秒，共享内存大概是它的几倍延迟，全局内存的延迟可能是它的几十上百倍。所以GEMM优化的本质，归结为一句话：尽量让数据停留在最靠近计算单元的地方。

这个逻辑在日常生活中也很好理解。你不希望做饭时每炒一个菜都去一趟超市买菜，而是希望提前把菜全部搬到厨房案板上，甚至切好配好，做饭时才高效。GEMM里的“厨房案板”就是共享内存和寄存器，“去超市买菜”就是从显存读数据。好内核和差内核的差距，往往不是算得快不快，而是“跑了几趟超市”。

2.3 命中缓存的交叉访问模式

朴素内层循环里，A是按行连续访问的，这对缓存比较友好；但B是按列访问的，每个内层循环都会跳一行，缓存命中率极低。更麻烦的是，如果多个线程并行计算C的不同行，它们访问B的模式还会造成严重的Bank Conflict（共享内存冲突）或者缓存争抢。说白了，就是多个线程同时想访问同一批存储单元的不同地址，硬件被迫把一个请求拆成多次，白白浪费了大量周期。

这些东西都不深奥，但它们每一个都在告诉你同一个道理：直接硬算是不行的，必须设计新的数据流。

3. 核心优化手段：从分块到张量指令

针对上一节的问题，业界经过多年积累形成了一套组合拳。DeepGEMM 的内核实现基本就是把这套组合拳按顺序打了一遍。

3.1 分块：把大矩阵切进共享内存

第一板斧是分块（tiling）。我们不一次性把整个大矩阵塞进共享内存（也塞不下），而是把A和B分别切成小块，让一个线程块负责计算C中一个更大的块。例如一个线程块计算C的一个128×128的块，那么A矩阵只需要提供128×K的条带，B矩阵只需要提供K×128的条带。K那一维再切成很多小块，每次取一小段做累加。整个过程就是：

for (int k0 = 0; k0 < K; k0 += BK) { // 协作把 A 的 [M块:128][k0:k0+BK] 和 B 的 [k0:k0+BK][N块:128] 加载到共享内存 __syncthreads(); // 线程块内的子线程各自计算 C 的某一部分 }

好处一目了然：A和B的数据被读入共享内存之后，可以被整个线程块反复使用。如果我们把BK定成16或者32，那么每个共享内存元素可以被用来做多次乘加，数据复用率直接从1提升到数百。这一步是性能从“可怜”到“平庸”的关键。

3.2 两层tiling：共享内存tile之后还要寄存器tile

分块到共享内存还不够。共享内存本身也有带宽上限，而且线程要执行计算，还是得把数据从共享内存读进寄存器。所以第二个层次的tiling是寄存器tiling：每个线程只负责计算C块中的一小块，比如8×8，而不是一次算一个元素。每次从共享内存里取出A的一个小片和B的一个小片，直接更新寄存器里的累加器。累加器始终住在寄存器里，多个k迭代累积下来，一次数据读取可以发生多次乘加运算。这就是GEMM里常说的“outer product”思路：一个线程算一个8×8的小块，需要A的8个元素和B的8个元素，两两相乘得到64个结果。

我实际测试过，寄存器tile是决定内核性能上限非常关键的一环。如果没有这层tiling，每个线程只算一个元素，那么每做一次乘加都要从共享内存取两个数，共享内存带宽会成为新的瓶颈。有了8×8甚至16×8的寄存器tile，共享内存带宽压力就能大幅下降。

3.3 向量化与Bank Conflict

往共享内存里搬运数据时，也不能一条一条地搬。现代的GPU都支持向量化访存，比如一次加载4个float，也就是128 bit。DeepGEMM里，所有的全局内存读取我都尽量用float4类型来操作，这样一条加载指令就能搬16字节，带宽利用率和指令效率都会更高。

Bank Conflict则是共享内存特有的坑。共享内存被分成了很多个bank，如果多个线程同时访问同一个bank的不同地址，就会发生冲突，硬件会把访问串行化。比如我们希望线程t访问共享内存中地址为t的位置，如果每个地址恰好落在不同的bank里，那就完美；但如果它们的地址差是某种会映射到同一bank的值，性能就会骤降。常见的解决办法是填充（padding）：在共享内存数组的行长度上加一个或几个float的偏移，打乱列地址和bank的对应关系。这个细节特别细小，但实测对性能影响非常显著，我后面在排障章节还会专门展开。

3.4 用张量指令代替普通乘加

到了这一步，性能已经比朴素版本好很多了，但离硬件的理论峰值还有很大距离。主要原因是：普通的FP32乘加指令，每个线程每个周期能执行的运算数有限。现代旗舰级GPU都配备了专门的矩阵运算硬件单元，一条指令就能完成一个小的矩阵乘加，比如一次计算一个16×16×16的矩阵乘加。API层面对应的概念就是wmma或者mma这样的指令接口：

// 伪代码：硬件矩阵乘加指令 // 其中 a 是 16x16 切块，b 是 16x16 切块，c 是 16x16 累加器 mma_sync(c, a, b, c);

DeepGEMM 里，我主要针对半精度（FP16/BF16）输入做了矩阵指令适配。使用该指令后，性能会有一个质的飞跃：同样一个周期内完成的FLOPs可以翻好几倍。不过随之而来的问题是数据布局：这类指令对A、B的布局有严格要求，很多情况下需要提前把矩阵转换成特定的分块布局，而不是简单的行主序。这个“布局转换”是一个经常被忽略但极其影响性能的工程细节。

3.5 数据布局：什么时候做T变换

我用一个简单的例子说明布局对矩阵指令的重要性。假设硬件矩阵指令要求A的tile按[M,K]顺序排布、B按[K,N]排布且K维连续，而你的数据实际是行主序存储，K维在内存里恰好是连续的，那么恭喜你，布局可以直接喂给指令。但如果B是按列主序存的，或者你为了共享内存读取方便切成了按block存储的格式，那就要做一次转置或者重排。两次选择我都试过：一是在全局内存里提前重排，也就是常说的“布局预转换”，适合推理场景里权重固定不变的情况；二是在内核加载进共享内存的过程中同步重排。前者性能更好但需要额外的存储操作，后者灵活但会增加共享内存阶段的开销。DeepGEMM 走的是“加载时重排”方案，主要为了让代码支持更多输入格式。

4. 动手实现一个DeepGEMM内核

理论说了这么多，最终还是要落到代码上。这一节我给出一个简化但完整可用的内核骨架，设计思路对标主流高性能GEMM的实现路径。代码会省略一些工程细节，但主循环和数据流是完整真实可运行的。

4.1 工程结构与启动参数

我按如下结构组织项目：

• gemm.h：对外接口，模板化支持FP16、BF16、FP32。
• gemm_kernel.cu：CUDA内核实现。
• bench.cu：基准测试，对比朴素实现和一个参考实现。
• test.cu：正确性校验，随机生成矩阵与朴素版本对比。

内核启动参数是关键。我针对“M=8192，N=8192，K=8192”的典型大矩阵，线程块尺寸设为128×128，每个线程计算8×8的寄存器tile，K方向每次切16。这样算下来每个线程块包含256个线程：

• 线程块网格：grid.x = N / 128，grid.y = M / 128。
• 线程块内线程数：256。
• 每个线程的累加器：8行×8列 = 64个float。
• 共享内存中的A切片：128×16个半精度元素，B切片：16×128个半精度元素。

4.2 主循环与双缓冲

为了提高数据加载和计算的并行度，我采用双缓冲（double buffering）方案：共享内存里准备两份A和B的缓冲区，一份用于当前k块的计算，另一份用于后台预取下一块。这样计算单元不会因为等数据到来而闲下来。

下面是主循环的伪代码：

for (int k = 0; k < K; k += BK) { // 阶段1：让线程协作把A、B的当前切片加载到当前缓冲区 load_tile_to_smem(A, B, k, buffer[0]); __syncthreads(); // 阶段2：在计算当前缓冲区的同时，预取下一块到buffer[1] if (k + BK < K) { load_tile_to_smem(A, B, k + BK, buffer[1]); } // 阶段3：计算，从buffer[0]读取数据，更新累加器 compute_tile_from_smem(buffer[0]); __syncthreads(); }

真实内核里，预取往往靠异步拷贝指令（类似cp.async）把数据从全局内存直达共享内存，不占用寄存器，也不阻塞计算。实现上更复杂一些，但思路就是“让数据和计算像两条平行的流水线一样同时运转”。我在代码里用了基础的cp.async指令，发现这个改动带来的收益非常明显，尤其在K很大的情况下，预取能把内存延迟完全掩盖掉。

4.3 计算阶段的核心代码

计算阶段的大致逻辑是：每个线程根据自己在线程块中的位置，从共享内存的A、B切片中选取所需的8个和8个元素，然后循环K方向累积。下面是一个高度简化的示意：

// 简化后的计算核心：线程 tid 负责计算一个 8x8 的结果块 float acc[8][8] = {0.0f}; for (int kBlock = 0; kBlock < BK; ++kBlock) { float aReg[8], bReg[8]; // 从共享内存读取A切片中本线程需要的8个元素 // 从共享内存读取B切片中本线程需要的8个元素 for (int i = 0; i < 8; ++i) { for (int j = 0; j < 8; ++j) { acc[i][j] += aReg[i] * bReg[j]; } } } // 最后把acc写回全局内存C

当然这个示意版本没有包含地址映射。地址映射的推演是GEMM内核里比较容易把人绕晕的部分，也是很多材料略过不讲的点：线程tid如何映射到C块中的具体行和列。我给你的建议是画一张图：把线程块内的256个线程编号，把它们想成16行×16列的二维网格，再把每个线程负责的8×8累加器铺开，凑成整个128×128的C块。把这张对应关系先画出来，再去写代码，会顺手很多。

4.4 从固定形状扩展到通用形状

固定形状的内核能跑得很快，但真实场景里M、N、K都是任意的。DeepGEMM采用了两种策略：第一，把不满足128对齐的边角区域当成“单列/单行”处理，使用简单版本的内核或者标量循环；第二，对于中间大块，用高性能分块内核，边角用fallback内核补齐。这样组合起来，虽然边角性能略差，但总体不会有太大退化。很多成熟库也是这么干的：大规模主流形状走全优化内核，非主流形状走通用路径。

我做了一个测试：形状为8192×8192×8192时，主内核能发挥接近硬件峰值；而把形状改成8192×1×8192后，要么退化成矩阵向量乘积，要么继续用大内核导致大量线程空转，性能掉得很厉害。这说明GEMM的优化极度依赖形状，库函数内部其实维护着一堆针对不同形状的“内核选择表”，不是一把梭。

5. 实测与基准

写内核和验证性能是不可分割的。这里给出我实测过程中的一些数据和解读方式，方便你复现后自己做对比。

5.1 测试配置

我在某个主流GPU设备上做了测试。为避嫌不写具体型号，直接给相对数据比例。测试矩阵统一使用FP16输入、FP32累加，规模5120×5120×5120。对比对象包括：朴素三重循环内核、优化分块FP16内核、加入张量指令后的DeepGEMM内核。

5.2 对比结果

从这个表能明显看到，每个优化阶段消除的瓶颈都不同。朴素版本输在内存访问上，分块版本赢在数据复用，而最后的张量指令把计算吞吐拉满。如果你看到自己的内核性能百分比长期停在个位数，大概率是某一层瓶颈没解决，而不只是浮点运算不够快。

5.3 不同形状的敏感度

我又测试了一组形状差异很大的矩阵，包括M=1的GEMV形状、M=64的小批形状、M=4096的超大形状。结论比较典型：

• 超大矩阵：主内核性能稳定，利用率高。
• M=64、N=8192、K=8192：虽然也能用主内核，但边角浪费严重，性能约下降25%。
• M=1、N=8192、K=8192：主内核完全不合适，应该走专门优化过的矩阵向量乘实现。即便换路径，性能也比大矩阵GEMM低很多。

这个敏感度不是bug，它来自并行度不足。GEMM要快，需要足够多的线程块填满GPU的所有计算单元。M太小意味着线程块数量太少，大量硬件资源闲置。群里经常有人问“为什么我的GEMM在某形状下那么慢”，九成是这个问题。通用库通过维护多个内核来缓解，但总有一些形状怎么都难做好。

6. 常见问题与排查清单

好，代码也写了，性能也测了，接下来是实践中最容易被卡住的部分。我把过程中遇到的高频问题整理成一份速查表，方便你直接对照排查。

6.1 性能上不去的头号嫌疑：bank conflict

我在一次中期测试里发现内核性能比预期低了大概40%，怎么优化都上不去。后来用性能分析器逐个检查，发现是共享内存访问阶段出现了大量bank conflict。原因是我把A切片的行大小设置成与线程束访问模式恰好冲突的值，导致同一束内不同线程大部分时间都在访问同一个bank。解决办法十分直白：在共享内存声明里，把行宽从float tileA[128][16]改成float tileA[128][17]，多出来的一个列作为padding。就这么一行改动，性能接近翻倍。你可以用同样的方式快速验证自己是否踩坑。

6.2 精度问题的排查思路

GEMM里精度误差绝大多数来自累加顺序和中间精度。DeepGEMM在FP16输入时，强制使用FP32累加器，这样能保证结果和朴素FP32版本相差很小。如果你发现误差超出预期，建议先做一个“单块单线程”的对照实验：用完全相同的数据，先跑朴素三层循环，再跑优化内核，把两者输出差打印出来。如果差距集中在某些特定位置，优先怀疑布局转换或同步问题，而非浮点精度本身。

6.3 工具与指标

强烈建议装上性能分析工具，比如CUDA自带的Nsight Compute（ncu），它给出的指标对定位瓶颈非常关键。我日常主要看几个指标：

• 计算单元利用率（SM busy）：反映计算资源是否吃满。
• 内存吞吐（Memory Throughput）：反映显存带宽是否成为瓶颈。
• 共享内存bank冲突数（Shared Memory Bank Conflicts）：帮助定位共享内存访问踩坑。
• warp占用率（Occupancy）：判断线程块是否足够多，能否掩盖延迟。

把这几个指标组合起来，基本能定位90%的GEMM性能问题。

6.4 终极调试技巧

如果实在找不到性能瓶颈，还有一个笨办法：不断地“删功能”。先把双缓冲去掉，看性能变化；再把张量指令降级为普通乘加，看性能变化；再把分块改小一倍。通过这种二分法对比，你能快速定位到底是哪一层优化没生效，哪一层反而拖了后腿。这套方法我一直用到现在，比肉眼猜代码靠谱得多。

最后补一句

这次写DeepGEMM最大的收获，是我真正理解了为什么那些高性能库能跑得那么快：它们本质上不是在算矩阵，而是在管理数据流动。一开始我满脑子都是“怎么让乘加更快”，后来才意识到，真正的跷跷板在“怎么让数据在正确的时间、出现在正确的位置”。踩过几次坑之后，我对GPU计算的整体认知有了很大提升，也更愿意去读那些底层资料了。如果你也想试试手写GEMM，不用贪多，先在固定形状下跑通一遍分块+共享内存+寄存器tile，再逐步加张量指令和双缓冲。等你自己把第一个接近硬件峰值的GEMM内核调出来时，那种满足感绝对值得。