DeepGEMM 高性能矩阵乘法算子库设计与优化实战

1. 从 DeepGEMM 这个名字说起它到底在解决什么问题第一次看到 DeepGEMM 这个项目名很多人会下意识把它归到“深度学习框架”那一类里去觉得又是一个大而全的训练推理平台。实际上完全不是。DeepGEMM 是一个专注到极致的高性能矩阵乘法GEMM算子库核心目标只有一个把深度学习里最耗时的那类矩阵运算在特定硬件上压榨出接近理论峰值的性能。它不负责模型定义不负责数据加载也不负责分布式调度它只干一件事——让矩阵乘得快、稳、可复用。矩阵乘法在深度学习里的地位用一句话概括就是“算力的主要消耗者”。一个 Transformer 结构里注意力机制里的 QK^T、softmax 之后的加权求和、前馈网络里的两层全连接全都是矩阵乘法。卷积网络里卷积在 im2col 之后同样可以转化成矩阵乘法。所以一个 GEMM 库的性能直接决定了整个模型训练和推理的吞吐上限。DeepGEMM 就是冲着这个瓶颈去的。它适合谁来参考三类人最值得花时间研究。第一类是算子开发工程师需要写 CUDA 或类似底层代码想学习怎么把分块、流水线、寄存器复用这些技巧组合到极致。第二类是推理框架开发者需要在部署阶段替换掉默认的 GEMM 实现换取更低的延迟。第三类是性能调优爱好者想理解一个 GEMM 内核从朴素实现到接近硬件峰值中间到底要跨过哪些坑。哪怕你平时只写 Python读懂 DeepGEMM 的设计思路也能帮你在选型时判断“这个库到底值不值得接进来”。需要提前说明的是DeepGEMM 这类项目通常高度依赖具体硬件架构不同代际的 GPU 在 SM 数量、共享内存大小、Tensor Core 指令集上差异很大。所以本文讨论的是通用设计方法论和实操思路具体参数需要你根据自己的硬件实测调整。下面我会从整体设计、核心细节、实操流程、问题排查四个层面把这类 GEMM 库的完整面貌拆开讲。2. 整体设计与思路拆解为什么 GEMM 要这样写2.1 朴素实现为什么慢先看清敌人要理解 DeepGEMM 的设计得先知道一个“天真”的矩阵乘法有多慢。假设计算 C A × BA 是 M×KB 是 K×NC 是 M×N。最朴素的写法是三重循环对 C 的每个元素遍历 K 维度做乘加。这个写法的问题在于每次取 A 和 B 的元素都要访问全局内存而全局内存的延迟是几百个时钟周期算力全被内存等待吃掉了。一个合格的 GEMM 实现核心思想是提高计算访存比也就是让每次从内存搬进来的数据被尽可能多次地复用。这就引出了分块Tiling策略把大矩阵切成小块块内数据先搬到片上高速存储共享内存、寄存器然后在片上反复计算。片上存储的带宽比全局内存高一到两个数量级复用次数上去了性能自然就上来了。DeepGEMM 这类库的设计本质上就是在“怎么切块、怎么搬数据、怎么排流水线”这三个问题上做极致优化。切块大小决定了复用率搬数据的方式决定了延迟隐藏效果流水线排布决定了计算单元会不会空转。2.2 分块策略的取舍大块不一定好分块大小不是越大越好这里有个很实际的权衡。块越大数据复用次数越多但需要的共享内存和寄存器也越多。共享内存是有限的一块 GPU 上每个 SM 通常只有几十到一百多 KB。如果块大到共享内存放不下就得降级到全局内存性能直接崩掉。寄存器更紧张每个线程能用的寄存器数量有硬上限超了就会溢出到本地内存那比全局内存还慢。所以实际选块大小的时候通常遵循一个经验法则让共享内存占用控制在可用容量的 60% 到 70%留出余量给流水线的双缓冲甚至三缓冲。比如某代 GPU 每个 SM 有 100KB 共享内存那双缓冲场景下每个缓冲最多用 30KB 左右块大小就要据此反推。这个计算过程我在后面实操部分会给出具体例子。另一个容易被忽略的点是边界处理。矩阵维度往往不是块大小的整数倍最后一块会有“残缺”。处理边界有两种常见做法一种是在内核里加判断越界就不算另一种是提前把矩阵 padding 到整数倍。前者省内存但有分支开销后者无分支但多占内存和带宽。DeepGEMM 这类追求极致性能的库通常会根据维度是否对齐走不同的内核路径对齐时走无分支的快路径不对齐时走带判断的慢路径。2.3 流水线设计让计算和搬运重叠GEMM 内核里有两类操作搬数据从全局内存到共享内存从共享内存到寄存器和算数据乘加。如果这两类操作串行执行那计算单元在搬数据的时候就是闲置的利用率上不去。解决办法是软件流水线把循环体拆成多个阶段让第 i 次迭代的计算和第 i1 次的搬运同时进行。实现流水线的关键工具是异步拷贝指令。较新的 GPU 架构提供了从全局内存直接异步拷贝到共享内存的指令拷贝发起后 CPU 侧准确说是线程侧不用等待可以继续往下走等真正要用数据时再同步。这样搬运的延迟就被计算掩盖了。DeepGEMM 这类库会把主循环展开成“预取下一块 计算当前块”的结构配合异步拷贝把内存延迟藏得严严实实。流水线的深度也有讲究。双缓冲两个共享内存缓冲交替用是最基本的能隐藏一部分延迟。三缓冲甚至更深能隐藏更长的延迟但占用更多共享内存。实际选几级缓冲要看内存延迟和单块计算时间的比值。如果单块计算时间远大于内存延迟双缓冲就够了如果计算很快、内存很慢就得上更深的流水线。2.4 为什么不用现成的库自研 GEMM 的价值有人会问cuBLAS、cuDNN 这些现成库已经很强了为什么还要自己写 GEMM原因有几个。第一特定形状的优化空间。通用库要照顾所有形状而你的模型可能只用几种固定形状针对这几种形状做特化性能能再上一个台阶。第二融合机会。GEMM 后面经常跟着 bias 加法、激活函数通用库只能分开调用中间结果要写回内存再读出来自研内核可以把这些融合进去省掉一次往返。第三可控性。遇到性能问题自研代码能直接定位到具体哪一行通用库只能当黑盒调参。DeepGEMM 的定位就是给需要这类特化和融合能力的场景提供一个高质量起点。它不是要取代通用库而是在通用库覆盖不到或者覆盖不好的地方补位。3. 核心细节解析与实操要点把性能榨出来的关键3.1 共享内存布局避免 bank conflict共享内存虽然快但有个坑叫bank conflict。共享内存被划分成若干个 bank同一时刻如果多个线程访问同一个 bank 的不同地址就会串行化性能打折。GEMM 里从共享内存读 A 和 B 的时候如果布局没设计好很容易踩这个坑。解决办法通常是padding。比如一个块是 32×32 的 float 数组按行存每行 32 个元素。如果直接存第 0 列和第 32 列会落在同一个 bank 上假设 32 个 bank访问时冲突。如果在每行末尾多留一个元素变成 32×33那第 0 列和第 32 列就错开了冲突消失。这个 padding 只多占一点点共享内存但能显著提升读取效率。DeepGEMM 这类库在共享内存布局上会做精细设计有时候还会用swizzle一种地址重排来进一步打散访问模式。swizzle 的原理是把地址的某些位做异或让原本连续的访问变得分散从而避开 bank 冲突。这个技巧在 Tensor Core 场景下尤其重要因为 Tensor Core 对数据布局有特定要求布局不对根本喂不进去。3.2 寄存器分块每个线程算多少共享内存解决的是“块级复用”寄存器解决的是“线程级复用”。一个线程从共享内存读数据到寄存器后应该尽可能多地参与计算而不是读一次算一次。这就引出寄存器分块每个线程负责计算 C 的一个小片比如 8×8那它需要从共享内存读 8 个 A 元素和 8 个 B 元素然后做 64 次乘加。读 16 个、算 64 次复用比是 4比读一次算一次强多了。寄存器分块的大小要平衡。分块越大复用比越高但寄存器压力也越大。一个 8×8 的分块需要 64 个累加器寄存器加上 A、B 的临时寄存器很容易超过 255 个的上限。实际调的时候通常从 4×4 或 8×8 起步看寄存器占用和性能曲线找到拐点。这里有个实操心得用编译器报告看寄存器占用。编译 CUDA 代码时加上--ptxas-options-v编译器会打印每个内核用了多少寄存器、多少共享内存、有没有溢出。如果看到寄存器数接近上限或者有 spill溢出就说明分块太大了得往下调。这个报告是调优的第一手依据比盲目试快得多。3.3 Tensor Core 的接入数据类型和指令选择现代 GPU 的算力大头在 Tensor Core 上普通 CUDA Core 的浮点算力只是零头。所以一个高性能 GEMM 库必须用上 Tensor Core。Tensor Core 的接入涉及几个关键决策用什么精度、用什么指令、数据怎么排布。精度方面Tensor Core 支持 FP16、BF16、TF32、INT8 等多种输入类型累加通常是 FP32。选哪种取决于你的模型需求。推理场景常用 FP16 或 INT8训练场景常用 BF16 或 TF32。不同精度的吞吐差异很大INT8 通常是 FP16 的两倍左右TF32 是 FP16 的一半左右。这个换算关系在选型时很关键。指令方面不同代际的 GPU 有不同的 Tensor Core 指令比如 mma 指令有 m16n8k16、m16n8k8 等多种形状。选指令要看你的分块形状能不能对上。如果分块是 16×8 的那 m16n8 系列就合适如果分块是 16×16可能要用两条指令拼。DeepGEMM 这类库会针对不同架构准备多套内核运行时根据硬件选择。数据排布是 Tensor Core 最容易踩坑的地方。Tensor Core 要求操作数按特定格式放在寄存器里比如 A 操作数要求每个线程持有特定的几个元素。如果从共享内存读出来的数据排布不对就得先做一次 shuffle 或者转置这个开销不小。所以高性能库通常会在写共享内存的时候就按 Tensor Core 要求的格式写读出来直接能用。3.4 边界与对齐那些不起眼但致命的细节前面提过边界处理这里展开说几个具体坑。第一个是维度不对齐导致的性能悬崖。假设你的块大小是 128矩阵维度是 127那最后一块只有 1 列有效其余全是浪费。更糟的是如果内核没处理好这 1 列有效数据可能触发大量分支判断性能比对齐时差好几倍。实测中见过维度从 128 变成 129性能直接掉 40% 的情况。第二个坑是指针对齐。很多向量化加载指令要求地址按 16 字节或 32 字节对齐。如果矩阵的起始地址没对齐或者行 stride 不是对齐的向量化加载就用不了只能退化成标量加载带宽利用率大打折扣。解决办法是在分配内存时就保证对齐比如用cudaMalloc分配时它本身会返回 256 字节对齐的地址但如果你自己做切片就要小心保持对齐。第三个坑是K 维度不是指令 K 的整数倍。Tensor Core 指令的 K 通常是 8 或 16如果矩阵的 K 是 12那最后 4 个元素没法用一条完整指令处理得单独补。这个补的过程如果写得不好会引入额外开销。常见做法是把 K padding 到整数倍多出来的部分填零这样内核逻辑统一代价是多算一点点。4. 实操过程与核心环节实现从零搭一个 GEMM 内核4.1 环境准备与基线测量动手之前先把环境和基线搭好。你需要一块支持 Tensor Core 的 GPU、对应版本的 CUDA 工具链、以及一个能编译和运行 CUDA 的工程。工程结构建议简单点一个.cu文件放内核一个主程序负责分配内存、初始化数据、调用内核、校验结果、计时。基线测量很重要别跳过。先跑一个朴素的三重循环版本记录它的耗时。再跑一个只用共享内存、没用 Tensor Core 的版本记录耗时。最后跑一个用 Tensor Core 的版本。这三个数据能帮你判断每一步优化带来了多少收益也能在后续调优时作为参照。我见过有人一上来就写 Tensor Core 版本结果性能还不如朴素版就是因为基础的分块和流水线没做对。计时的时候注意用CUDA event而不是 CPU 计时。CPU 计时会把内核启动开销算进去对于短内核误差很大。CUDA event 能精确测量 GPU 侧的执行时间。另外要跑多次取平均第一次运行往往有冷启动开销不能算数。4.2 分块参数的计算过程假设你的 GPU 每个 SM 有 100KB 共享内存你要用双缓冲那每个缓冲最多 50KB。再假设你用 FP16 数据每个元素 2 字节那每个缓冲能放 25600 个元素。如果 A 块是 BM×BKB 块是 BK×BN两个块加起来不能超过 25600。取 BMBN128那 128×BK BK×128 256BK ≤ 25600BK ≤ 100。取 BK64则 A 块 128×648192 个元素B 块 64×1288192 个合计 16384 个元素占 32KB双缓冲 64KB留了 36KB 余量安全。寄存器方面假设每个线程算 8×8 的输出需要 64 个累加器。加上 A、B 操作数的临时寄存器大概 100 个出头。如果编译器报告超过 128 个就要考虑把分块降到 8×4 或 4×8。这个计算不是一次到位的要反复试用编译器报告和实测性能一起判断。线程块大小也要算。BM×BN128×128 的输出如果每个线程算 8×8那需要 (128/8)×(128/8)256 个线程。256 是 2 的幂正好是常见的线程块大小。如果算出来不是 2 的幂或者超过 1024就要调整分块。4.3 内核主循环的写法主循环的结构大致是这样外层循环遍历 K 维度每次处理一个 BK 大小的块。循环体里先把 A 和 B 的当前块从全局内存搬到共享内存用异步拷贝然后从共享内存读到寄存器做 Tensor Core 矩阵乘累加到寄存器里的 C 分块。循环结束后把 C 分块写回全局内存。关键点在于异步拷贝和计算的交错。正确的顺序是在循环开始前先发起第一块的拷贝循环里先等待上一块拷贝完成发起下一块拷贝然后计算当前块。这样拷贝和计算就重叠了。如果顺序写反了变成“等待拷贝、计算、再发起拷贝”那就没有重叠性能差一截。写这个循环的时候同步原语要用对。异步拷贝的完成要用cp.async.wait_group之类的指令等待不能用普通的__syncthreads()后者会等待所有线程粒度太粗。另外要注意共享内存的读写之间要有同步避免一个线程还在读、另一个线程已经写了。这个同步通常用__syncthreads()但位置要放对放多了会拖慢放少了会出错。4.4 结果写回与 epilogue 融合主循环算完C 分块还在寄存器里要写回全局内存。这一步看似简单其实也有优化空间。如果 C 的写回是分散的、不对齐的带宽利用率会很低。常见做法是让每个线程负责写回连续的一段或者用向量化的写指令一次写 4 个或 8 个元素。更重要的是epilogue 融合。实际模型里GEMM 后面往往跟着 bias 加法、激活函数ReLU、GELU 等、甚至 dropout。如果分开做C 要先写回全局内存再读出来做 bias 和激活再写回去两次往返。融合的做法是在寄存器里就把 bias 加上、激活算完然后一次性写回。这样省掉一次全局内存往返对于大矩阵能省不少时间。融合的难点在于 bias 和激活的访问模式。bias 通常是按 N 维度广播的每个输出列对应一个 bias 值。在寄存器分块里每个线程负责的列是固定的所以 bias 可以提前加载到寄存器循环里直接用。激活函数如果是逐元素的直接在寄存器上算就行没有额外开销。这些细节决定了融合能不能真正带来收益。5. 常见问题与排查技巧实录踩过的坑都在这5.1 性能不达预期从哪查起性能不达预期是最常见的问题排查要有顺序。第一步确认瓶颈在哪。用性能分析工具看内核的占用率、内存吞吐、计算吞吐。如果计算吞吐很低但内存吞吐很高说明是内存瓶颈要优化数据复用如果两者都低说明是延迟瓶颈要加深流水线如果计算吞吐接近峰值但整体还是慢可能是启动开销或尾部效应。第二步检查分块参数。用编译器报告看寄存器占用和共享内存占用。如果寄存器溢出降分块如果共享内存没用满可以尝试加大分块。这一步往往能解决大部分问题。第三步检查边界处理。把矩阵维度设成块大小的整数倍看性能有没有提升。如果有明显提升说明边界处理有优化空间。可以针对对齐情况写一个专门的内核不对齐时再走通用内核。第四步检查同步和流水线。用分析工具看内核里有没有长时间的 stallstall 的原因是什么。如果是等待内存说明流水线不够深如果是等待同步说明同步放多了。5.2 结果不对数值精度和索引错误结果不对分两类一类是数值精度问题一类是索引错误。精度问题通常出现在用低精度计算时比如 FP16 累加导致溢出或精度损失。解决办法是用 FP32 累加或者用 Kahan 求和之类的补偿算法。如果模型对精度敏感就别用 INT8老老实实上 FP16 或 BF16。索引错误更隐蔽。常见的是共享内存的索引算错导致读到了别的块的数据。排查方法是把矩阵设小比如 4×4然后手动算一遍期望结果和内核输出对比。如果对不上就在内核里加打印看每个线程读到了什么、算出了什么。这个过程很枯燥但最有效。还有一个坑是线程块和矩阵维度的映射搞反。比如把 M 和 N 弄混了结果矩阵转置了。这种错误在小矩阵上不容易发现因为转置后形状可能还对得上。排查方法是故意用非方阵比如 M8、N16如果输出形状不对就是映射错了。5.3 常见问题速查表现象可能原因排查方法解决方向性能只有峰值 10%没用 Tensor Core 或分块太小看计算吞吐和指令类型接入 Tensor Core加大分块性能随维度波动大边界处理低效对比对齐和非对齐维度写专门的对齐内核寄存器溢出分块太大看编译器报告降分块或减少临时变量结果偶尔错同步缺失或竞态加同步后看是否消失补__syncthreads()精度不达标低精度累加对比 FP32 结果改 FP32 累加或补偿启动开销大内核太小或启动频繁看内核执行时间和总时间比合并内核或增大批量5.4 几个独家避坑心得第一个心得先用小矩阵验证正确性再用大矩阵测性能。小矩阵容易手算验证大矩阵才能看出性能差异。很多人一上来就用大矩阵结果错了都不知道错在哪。第二个心得保留一个参考实现。写一个朴素但绝对正确的版本每次改完优化版本都拿参考实现对比结果。这样能第一时间发现回归。第三个心得性能调优要单变量。一次只改一个参数测完记录再改下一个。同时改多个参数出了问题不知道是哪个引起的。第四个心得注意 GPU 的温度和频率。长时间跑性能测试GPU 会降频后面的数据不可比。测试之间留冷却时间或者用固定频率模式。第五个心得别迷信理论峰值。理论峰值是理想条件下的上限实际能到 70% 到 80% 就已经很不错了。追求 100% 是浪费时间不如把精力放在实际场景的端到端优化上。6. 工具选型与扩展思路GEMM 之外还能做什么6.1 什么时候该自研什么时候该用现成的自研 GEMM 有成本不是所有场景都值得。判断标准很简单如果你的模型只用几种固定形状且这些形状在通用库里性能不理想那就自研。如果形状多变或者通用库已经能满足需求那就别重复造轮子。另一个判断维度是融合需求。如果你的 GEMM 后面跟着复杂的 epilogue通用库不支持融合那自研的收益就很大。反之如果就是纯 GEMM通用库往往已经调得很好了。还有一个现实因素团队能力。自研 GEMM 需要懂 CUDA、懂硬件架构、懂性能分析的人。如果团队里没有这样的人强行自研可能还不如用现成的。这时候可以考虑基于开源库做二次开发而不是从零写。6.2 从 GEMM 扩展到其他算子GEMM 做透了很多其他算子也能触类旁通。比如卷积在 im2col 之后就是 GEMM可以直接复用 GEMM 内核。注意力机制里的 QK^T 和加权求和也是 GEMM只是中间夹了个 softmax。归一化层虽然本身不是 GEMM但里面的矩阵运算可以借鉴 GEMM 的分块和流水线思路。扩展的时候要注意不同算子的访存模式不一样。GEMM 是规则访问卷积在 im2col 之后也是规则的但注意力里的 softmax 涉及跨行归约访存模式更复杂。直接套 GEMM 内核可能不是最优需要针对性地调整。6.3 后续可以深入的方向如果 GEMM 已经调得差不多了可以往几个方向深入。一是稀疏 GEMM利用矩阵的稀疏性跳过零元素能大幅提升有效算力。二是混合精度不同层用不同精度在精度和速度之间找平衡。三是自动调优写一个搜索框架自动遍历分块参数、流水线深度等找到最优配置。四是跨平台适配把内核移植到不同厂商的硬件上这需要抽象出硬件无关的接口层。这些方向每一个都够写一篇长文这里只是点个方向。核心还是那句话先把一个 GEMM 做透理解了分块、流水线、寄存器复用这些底层逻辑再看其他算子就是同一套方法论的迁移。我个人在实际操作中的体会是GEMM 优化最难的从来不是写代码而是建立性能直觉。你得能看着一组参数大概判断出性能会在什么量级瓶颈可能在哪。这个直觉只能靠大量实测积累没有捷径。每次调优都记录参数和结果时间长了看到新形状就能快速给出合理的初始配置这才是真正的经验。