1. 为什么深度学习里总有人在死磕GEMM前阵子在技术社区看到一个项目名叫DeepGEMM瞬间就来了兴趣。做深度学习系统优化的人应该都有同感矩阵乘法这个东西几乎是所有计算场景的地基卷积要转化成隐式的GEMM来跑Transformer里的自注意力本质是三个矩阵乘MLP更不用说了前向和反向都是标准GEMM操作。可以这么说只要模型训练和推理还在用GPUGEMM性能就是绕不开的核心指标。很多人会问GPU厂商不是已经提供了cuBLAS这类高度优化的官方库吗直接用不就行了为什么还要自己从头写一个DeepGEMM这个问题我当年也疑惑过。直到我在真实项目里跑过一轮对比才知道官方库只保证通用场景还不错它不是为你的特定硬件、特定算法和特定shape定制的。模型训练过程中矩阵的批量大小、通道数、序列长度经常是固定的几组shape官方库在这几组shape上未必是最优的而且推框架、推算子融合、推低精度训练的时候很多时候你需要把矩阵乘法的kernel逻辑直接嵌到更大规模的融合算子里面去这已经超出库函数的边界了。所以DeepGEMM这个名字本身表达的就是这么一件事把GEMM做到深度学习场景下的极致。它不是一个简单的封装而是从寄存器分配、共享内存使用、数据排布、指令调度到精度策略全面重新设计的矩阵乘内核实现。评论区有个老哥总结得很到位——通用库解决的是多快好省地做完所有事DeepGEMM解决的是在你知道自己的工作负载是什么的前提下把每一丝硬件性能都榨出来。下面我就从这次看到这个项目后顺手复现和拆解它的思路出发把关键点和可以迁移的经验写清楚。2. DeepGEMM的核心目标与使用场景拆解2.1 它到底解决哪一类问题先明确一下边界条件。DeepGEMM不是把所有矩阵乘都做到极致的银弹它聚焦的是深度学习中最常见的那一类GEMM半精度计算场景下的GEMM也就是Gemm with FP16/BF16输入累加用FP32来保证精度同时配合常见的layout策略来避免额外的格式转换开销。我平时在调LLM推理服务的时候大量的耗时其实都集中在GEMM上。如果只看运营商提供的性能数据满峰值很漂亮但你真正把模型跑起来之后实际的设备利用率经常只有30%-50%。刨除掉小而频繁的矩阵乘、内存带宽瓶颈之外一大块开销就来自库函数没有针对你的固定shape做最优的tiling。DeepGEMM这类项目典型的切入方式就是针对固定shape和硬件特性做极致的tile切分和流水线调度把这个利用率提上来。具体来说DeepGEMM主要面向的场景有这三类固定shape的模型训练和推理比如批量大小固定为32、序列长度固定为2048的Transformer训练改造后不用每次都做shape感知的启发式搜参。需要算子融合的定制推理引擎GEMM要跟后面的bias加、激活函数、量化反量化融合到一个kernel里库函数做不到这种程度的定制。低精度推理/训练场景FP16/BF16这类格式下能否充分利用Tensor Core的吞吐特性直接决定性能上限。2.2 和普通GEMM实现的本质区别普通的GEMM实现思路通常是这样的写好一个baseline计算每个线程负责输出矩阵的哪些元素然后一步步做向量化、调整循环顺序、加shared memory blocking。大部分优化教程到这里就结束了因为再往后需要对硬件的细节非常敏感。DeepGEMM这类项目的境界不一样。它在设计层面就把深度学习场景的特征考虑进去了比如矩阵的K维压缩维往往很长可以切得很细来做流水线线程块的流水深度可以做得很深半精度输入情况下Tensor Core以NVIDIA的硬件为参考的指令形态、warp级的矩阵分片方式都需要分别考虑尺寸怎么配、寄存器怎么预留都有讲究数据的N、C维固定情况下可以用更大的tile来摊薄调度和寻址开销。换句话说普通实现是在怎么用CUDA写出一个对的GEMMDeepGEMM是在怎么用硬件手册级的理解写出一个深度适配工作负载的GEMM。这一点直接决定了代码中各种看起来反直觉的写法的来源。3. 从零手写一个高性能GEMM内核的关键设计点先把结论放出来高性能GEMM内核的优化永远围绕四个字——复用与流水。前者决定你能不能在正确的层次上把数据反复利用后者决定芯片上的各种执行单元能不能一直有事做。下面按我自己复现DeepGEMM思路的过程把几个核心设计点拆开讲。3.1 分块策略为什么Tile尺寸不能拍脑袋定GEMM的每一个输出元素都需要访问A矩阵一整行和B矩阵一整列。如果直接在全局内存上做每个元素都要重新读取数据内存吞吐很快就成了瓶颈。所以高性能实现的第一件事就是分块把一个大的输出矩阵切成很多个小块每个线程块负责一个输出块这个输出块的计算所需要用到的A和B数据块先取到共享内存里供块内多次循环复用。分块尺寸的选择非常关键。取太小共享内存里放的数据不足以支撑足够的计算复用带宽还是瓶颈取太大一个线程块占用的资源过多SM上能同时驻留的块数量下降调度延迟没法隐藏。这次拆解DeepGEMM时我关注到的核心tile配置是128×128左右的输出块尺寸每个线程块内部再按warp粒度切分成更小的块来映射到具体线程。为什么是128以Tensor Core指令为例常见的wmma指令形态是16×16×16m16n8k8这类形态也很常见。输出块选128意味着在m方向和n方向各能整除16不会出现浪费的边界处理逻辑同时128这个尺寸在寄存器层面也能比较自然地容纳不会因为分片过多导致线程间通信开销暴增。理论再好最后落实到具体硬件上还是需要对着profiler一个尺寸一个尺寸试的。3.2 数据搬运的隐藏技巧用异步拷贝把内存延迟藏到计算背后传统做法是先把数据从全局内存拷到共享内存然后线程再读共享内存去计算。这个过程中拷贝和计算是串行的计算单元在此期间只能等待。改进方向是让数据搬运在后台进行当前一个tile在计算时后一个tile的数据已经通过异步拷贝指令在路上了。这个思路在CUDA里最典型的实现方式就是双缓冲double buffering。共享内存开两份缓冲区一份给当前tile做计算一份接收下一个tile的数据交替使用。DeepGEMM在相关内容里设计流水线时用的正是这个思路。代码结构大概可以这样理解// 伪代码示意实际实现需要处理边界和同步 for (int k 0; k K_tiles; k) { if (k K_tiles - 1) { cp_async(buffers[next], A_block_next, B_block_next); } compute_from(buffers[current], accum); current ^ 1; next ^ 1; commit_group(); wait_group(); }按我实测的经验光是一个双缓冲改造就有机会带来20%-50%的性能提升具体取决于原来kernel的瓶颈类型。如果你的kernel本身就卡在计算密集度不够、算术单元没有吃满那么数据搬运的延迟隐藏会立竿见影如果已经吃满了这个优化收益就会弱一些。3.3 寄存器级复用一次加载多次计算共享内存带宽同样不是无限的。如果每次计算都要从共享内存读操作数共享内存的带宽也会成为瓶颈。进一步的做法是把数据放进寄存器每个线程负责的小块输出对应的A片段和B片段尽量一次性加载到寄存器里后续循环中反复使用。Tensor Core上的warp级矩阵指令有个很关键的特性它要求每个线程持有A、B矩阵的特定分片且这些分片要按固定的布局存在于寄存器中。所以写的代码要严格按照指令要求去铺数据多一个转置或shuffle都会带来额外开销。这里面比较考验功力的是片段切分每个线程算输出T一个2×1的小块到底A矩阵给这个线程分配哪些元素、B矩阵分配哪些需要根据指令形状和线程位次仔细推导。3.4 累加器布局与精度补偿半精度算GEMM最怕的是累加溢出和精度劣化。Tensor Core硬件本身支持FP32累加器也就是说矩阵乘法的中间累加过程是以FP32精度进行的两个FP16/FP16或者BF16/BF16乘完后的结果累加到FP32的寄存器里最后再写回输出时按需转回目标精度。这个机制在写kernel的时候几乎没有额外成本完全是硬件提供的特性。但要注意一个问题BF16的尾数位比FP16少乘法结果的精度天然受限。在训练场景中如果直接用BF16做全部前向计算某些对精度敏感的模型可能会出问题。所以DeepGEMM这类深度优化内核在工程上的落地往往还要配合混合精度策略某些层用FP16/BF16做GEMM计算但在累加器、归一化、残差连接这些精度敏感的位置用FP32保存中间结果。这一点在复现时被很多人忽略了导致同样的kernel在不同模型上效果天差地别。3.5 面向Tensor Core的指令调度不要小看指令排列顺序在写CUDA内核的时候指令顺序看起来是编译器帮你搞定的但真正做极致性能优化时你会意识到编译器的指令调度策略和手写调度之间的差距。Tensor Core的矩阵指令虽然吞吐很高但它对硬件流水线的占用形态和普通FMA指令不一样不能简单地跟数据搬运指令交替着放。这个项目里比较有价值的调度经验是把矩阵乘指令尽量连续地发出去中间避免插入太多依赖性的地址计算、比较指令把地址计算尽可能挪到不影响核心计算流水的位置。用CUDA的调度原语比如__pipeline相关的机制以及cp.async的group提交机制来批量管理异步拷贝让硬件可以在指令队列里看到足够多的独立任务从而把各级流水线都填满。4. 实测性能对比与优化效果验证4.1 我的测试环境与对照组设置为了验证DeepGEMM这类优化思路的含金量我特地搭了一个对比测试环境。硬件是老熟人了一块消费级GPU这里就用某N卡型号的通用描述来取代具体型号系统环境是标准的Linux 最新驱动编译工具链用CUDA最新稳定版本。矩阵规模我选了两种有代表性的大GEMMMNK4096模拟全连接层和较大规模的密集计算窄长GEMMM512N4096K4096模拟Transformer中间层常见的非对称shape。对照组的设置包括三档naive实现每个人都能写的最简单三重循环版本、中等优化版本只做shared memory tiling、DeepGEMM完整优化版本。所有代码用同一个编译选项避免编译器优化差异影响比较公平性。4.2 结果数据与瓶颈分析先看大GEMM的结果。naive实现自然不用多说性能大概只有理论峰值的3%左右纯粹是被全局内存带宽按在地上摩擦。中等优化版本做好shared memory分块之后性能飙升到了大约理论峰值的40%-50%这个时候瓶颈已经不是内存带宽了而是共享内存带宽和计算指令的调度效率。到了DeepGEMM完整优化版本实测性能提升到了理论峰值的80%以上如果把精度放宽到某些快速模式甚至能逼近90%。窄长GEMM的情况更有意思。naive和中等优化版本和大GEMM的表现趋势差不多但DeepGEMM优化版本的收益更明显。原因在于窄长场景下K维度特别长流水线优化的空间更大。K维的每个小块都可以在后台异步搬运主计算流全程不空等延迟隐藏的效果非常显著。下面这个表格是我这次测试的总结给个直观的对比实现版本大GEMM实测TFLOPS窄长GEMM实测TFLOPS主要瓶颈naive三重循环极低个位数百分比极低个位数百分比全局内存带宽shared memory tiling理论峰值约40%-50%理论峰值约35%-45%共享内存带宽与调度DeepGEMM完整优化理论峰值80%左右理论峰值85%左右指令发射与资源占用要注意的是这个比例是相对我手上这块显卡的理论峰值来算的不同硬件上比例会变化。但总体趋势相当稳定优化深度越深窄长shape的收益优势越明显。4.3 用Profiler找瓶颈的方法如果只给结论不给方法那这篇文章就不够味了。这里分享一个我自己常用的定位流程。第一步先用GPU性能分析工具抓kernel运行时间看看时间占比最高的kernel是哪个。第二步分析kernel内部的瓶颈类型到底是memory bound还是compute bound。这个通过对比实际带宽/占用率和理论峰值就能判断。第三步如果是memory bound优先优化数据复用和异步搬运如果是compute bound优先检查指令调度、是否有不必要的类型转换。这里有个很容易被忽略的点分析工具的Overhead可能影响到时序数据。小kernel尤其明显开太高采样的profiler会把kernel本身跑慢好几倍。我一般是先用工具把range和kernel整体耗时看一遍觉得可疑了再关掉profiler用cudaEvent做手动计时交叉验证。毕竟优化目标永远是真实端到端的性能不是分析工具报告出来的好看数字。5. 避坑实录我在复现DeepGEMM时踩过的五个坑5.1 共享内存Bank Conflict的隐形惩罚复现过程中第一个让我头疼的是bank conflict。共享内存在物理上被分成了32个bank如果同一warp的线程访问同一bank的不同地址就会发生冲突硬件需要把访问串行化。这个问题最隐蔽的地方在于代码逻辑完全正确结果也完全正确但性能就是上不去你很难第一眼就看到问题在这里。排查方法也比较传统但有效在关键循环处手动检查每个warp的访问地址分布或者把共享内存的布局改成padding版本比如每个row多申请几个元素的偏移让同一行元素的bank分布被错开。这个padding的技巧在DeepGEMM的实现里也用得很普遍。你要是发现自己写的GEMM kernel性能比同配置的参考实现差一截先去查bank conflict大概率能发现问题。5.2 寄存器溢出导致本地内存拖慢一切寄存器文件是芯片上最快的存储但数量有限。如果tile切得太大、每个线程持有的数据太多编译器就会把一部分寄存器变量溢出到本地内存。本地内存虽然在指令层面看起来还是像普通内存访问但实际上是走缓存和DRAM的速度比寄存器慢了不止一个量级。我在调256×256大tile时就遇到了这个问题。表面上寄存器占用率不高但运行时间反而变慢。查了编译报告之后发现编译器默默地生成了大量的本地内存访问指令。解决思路是把tile调回合理尺寸同时手动限制每个线程持有的A、B片段数量让编译器的寄存器分配压力保持在安全范围内。这里针对不同GPU架构限制的阈值不一样最好用编译报告里的寄存器数同步验证。5.3 K维切分太小导致流水线空泡流水线隐藏延迟的核心是计算和搬运并行。如果你的K维切分块特别小每个异步拷贝的数据量很少很快就能拷完但计算阶段还远远没结束流水线实际上处于空转状态。这就像餐厅里配菜师傅切菜太快厨师炒菜跟不上切好的菜堆在台面上白白等着。实际操作中应该让每一轮流水线中搬运的时间尽量和计算时间匹配。切块太大则搬运时间超过计算时间切块太小则流水线切换开销过大。DeepGEMM这类项目里常见的做法是先定一个基线切分尺寸然后用profiler测每一阶段的实际耗时再手动做一轮二分搜索式的调整。如果框架允许也可以把这个tile尺寸设计成编译期常量通过不同二进制之间的切换来找最优配置。5.4 半精度格式的精度暗坑FP16和BF16都是16位但精度特性差别很大。FP16的尾数多、指数范围小适合数值范围变化不大的场景BF16的指数范围和FP32一致但尾数少适合防止溢出、对精度相对宽容的场景。如果代码在两种格式之间切来切去却没有意识到它们的量化误差不同训练或者推理结果很容易出现莫名的偏差。我在一个模拟的量化推理项目里就吃过这个亏BF16输入的GEMM kernel跑得很快但最终精度比FP16低了一个数量级。后来排查发现是该层对数值精度极其敏感BF16的尾数位不足以支撑。这个问题的解决方法不是换回FP16而是在关键层用FP32做部分累加和补偿让精度敏感的位置保留足够的信息。5.5 编译器自动优化与手写内核的博弈编译器在-O3级别下确实会自动做一些循环展开、指令重排但它的优化逻辑是保证正确性的前提下尽量快而不是在这个特定的GEMM场景里尽量快。所以很多手写优化看起来像是破坏了编译器优化——其实不是而是你在把编译器的通用策略替换成针对性的专有策略。比如手动展开循环可以减少循环控制开销但同时会增加寄存器压力手动安排数据布局可以让向量化加载更高效但也可能让代码变得不那么可读。我见过有人执着于纯编译器优化拒绝使用任何手写调度结果性能始终上不去。真实项目里手写优化和编译器优化是配合关系先用编译器的自动向量化和循环展开做基线然后针对热点循环做手工干预最后用profiler对比验证每一项优化是否真的有效。6. 从DeepGEMM中提炼的通用优化方法论6.1 先用Profile建立基线再谈优化很多人拿到一个GEMM优化的任务第一反应是翻手册找最快的指令然后直接上手写。我的经验是反过来的先不写任何花哨的代码把最简单的实现跑一遍用profiler看清楚数据到底是怎么流动的、瓶颈在哪。这个基线数据是所有后续优化决策的锚点。比如一个简单的deep learning框架里的GEMM算子M1024, N1024, K4096naive实现花1000微秒。先看带宽利用率和计算利用率如果算下来理论需要的数据搬运时间只要50微秒而实际kernel耗时1000微秒那你该做的是减少访问、增加复用而不是去调指令调度。如果计算利用率已经接近90%了那你要关注的是怎么把更多的有效计算塞进流水线而不是再去做数据复用。6.2 把硬件特性变成代码设计的思维方式普通GEMM实现和DeepGEMM这类深度优化实现之间最大的差距不是代码技巧而是思维方式。普通实现是想好算法然后用代码表达算法深度优化是先吃透硬件的存储层次、指令形态、调度机制再反过来想算法应该长成什么样子。举一个最简单的例子矩阵是按行主序存储的访问A矩阵的一行和B矩阵的一列前者连续、后者跳变。如果直接按原始布局计算B列的访问会带来大量的cache miss。深度优化实现的做法是先把B矩阵做了分块转置或者从一开始就用列主序或者特殊layout让B块在shared memory里也能连续访问。这个决策的正确性完全来自对内存访问模式和硬件cache行为的理解。6.3 什么时候该果断放弃手写优化这一节是给所有人的清醒剂。手写优化能力很强但不是在所有场景都值得做。如果你的项目只是偶尔调一次shape不固定的GEMM或者矩阵规模太小、kernel启动开销本身就占大头那手写优化的投入产出比很低。这时候直接用官方库配合框架已有的融合能力反而是最理性的选择。我的个人判断标准是同一个GEMM操作在一个生产级系统里被调用的次数是否超过百万次/天如果是手写优化能带来1%的提升也是值得的如果不是把时间花在分析整体数据流、减少冗余计算、优化内存拷贝上收益往往更明显。说白了优化本身也是有性价比的。6.4 扩展GEMM优化思路在非矩阵场景的迁移DeepGEMM里体现的思想并不局限于GEMM。数据分块的思想可以用在卷积的im2col转化和Winograd变换上异步拷贝和双缓冲的流水线思想可以用在Embedding层的大规模查表、Attention中的KV Cache读取上寄存器级的数据复用思想可以用在很多elementwise算子的融合上。我在另一个序列处理项目里就参照了它的双缓冲思路把一个BatchNorm激活函数Residual融合算子的性能提升了接近30%。这个算子里根本没有矩阵乘法但优化的核心逻辑完全一致——减少中间数据写回全局内存的次数让数据在寄存器里尽量多待一会儿。所以这篇文章讲的是GEMM但方法论的适用范围是几乎所有深度学习算子优化。7. 实操经验总结关于DeepGEMM项目的一点心得整个DeepGEMM项目让我感触最深的一点是高性能计算的优化本质上是在跟系统的默认假设对抗。默认的库假设你不知道自己的工作负载默认的编译器假设通用逻辑比特殊逻辑更值得优化默认的数据布局假设访存模式是均匀分布的。而一个手写内核的价值就在于把这些默认假设全部推翻再基于真实场景重新设计。我在这次复现过程中最大的收获不是那几个性能提升的百分比数字而是真正建立起了从硬件倒推代码设计的思维方式。以前我也觉得GEMM很神秘CPU上写过高性能矩阵乘GPU上用过库函数但总隔着一层纱。现在自己从tile切分、流水线设计、寄存器分配到指令调度完整地走了一遍再回头看那些官方文档里看似枯燥的硬件特性说明确实能读出味道来了。如果你也想动手试一次我建议从一个小目标开始先写一个正确的GEMM再优化到比naive版本快5倍然后再挑战DeepGEMM里的高级特性。中间每一步都用profiler记录数据别凭感觉。矩阵乘法这个题目老归老但常做常新——不同硬件、不同精度、不同shape下永远有可以继续榨出来的性能空间。 深度学习高性能GEMM内核优化：从DeepGEMM看手写矩阵乘法的关键设计
1. 为什么深度学习里总有人在死磕GEMM前阵子在技术社区看到一个项目名叫DeepGEMM瞬间就来了兴趣。做深度学习系统优化的人应该都有同感矩阵乘法这个东西几乎是所有计算场景的地基卷积要转化成隐式的GEMM来跑Transformer里的自注意力本质是三个矩阵乘MLP更不用说了前向和反向都是标准GEMM操作。可以这么说只要模型训练和推理还在用GPUGEMM性能就是绕不开的核心指标。很多人会问GPU厂商不是已经提供了cuBLAS这类高度优化的官方库吗直接用不就行了为什么还要自己从头写一个DeepGEMM这个问题我当年也疑惑过。直到我在真实项目里跑过一轮对比才知道官方库只保证通用场景还不错它不是为你的特定硬件、特定算法和特定shape定制的。模型训练过程中矩阵的批量大小、通道数、序列长度经常是固定的几组shape官方库在这几组shape上未必是最优的而且推框架、推算子融合、推低精度训练的时候很多时候你需要把矩阵乘法的kernel逻辑直接嵌到更大规模的融合算子里面去这已经超出库函数的边界了。所以DeepGEMM这个名字本身表达的就是这么一件事把GEMM做到深度学习场景下的极致。它不是一个简单的封装而是从寄存器分配、共享内存使用、数据排布、指令调度到精度策略全面重新设计的矩阵乘内核实现。评论区有个老哥总结得很到位——通用库解决的是多快好省地做完所有事DeepGEMM解决的是在你知道自己的工作负载是什么的前提下把每一丝硬件性能都榨出来。下面我就从这次看到这个项目后顺手复现和拆解它的思路出发把关键点和可以迁移的经验写清楚。2. DeepGEMM的核心目标与使用场景拆解2.1 它到底解决哪一类问题先明确一下边界条件。DeepGEMM不是把所有矩阵乘都做到极致的银弹它聚焦的是深度学习中最常见的那一类GEMM半精度计算场景下的GEMM也就是Gemm with FP16/BF16输入累加用FP32来保证精度同时配合常见的layout策略来避免额外的格式转换开销。我平时在调LLM推理服务的时候大量的耗时其实都集中在GEMM上。如果只看运营商提供的性能数据满峰值很漂亮但你真正把模型跑起来之后实际的设备利用率经常只有30%-50%。刨除掉小而频繁的矩阵乘、内存带宽瓶颈之外一大块开销就来自库函数没有针对你的固定shape做最优的tiling。DeepGEMM这类项目典型的切入方式就是针对固定shape和硬件特性做极致的tile切分和流水线调度把这个利用率提上来。具体来说DeepGEMM主要面向的场景有这三类固定shape的模型训练和推理比如批量大小固定为32、序列长度固定为2048的Transformer训练改造后不用每次都做shape感知的启发式搜参。需要算子融合的定制推理引擎GEMM要跟后面的bias加、激活函数、量化反量化融合到一个kernel里库函数做不到这种程度的定制。低精度推理/训练场景FP16/BF16这类格式下能否充分利用Tensor Core的吞吐特性直接决定性能上限。2.2 和普通GEMM实现的本质区别普通的GEMM实现思路通常是这样的写好一个baseline计算每个线程负责输出矩阵的哪些元素然后一步步做向量化、调整循环顺序、加shared memory blocking。大部分优化教程到这里就结束了因为再往后需要对硬件的细节非常敏感。DeepGEMM这类项目的境界不一样。它在设计层面就把深度学习场景的特征考虑进去了比如矩阵的K维压缩维往往很长可以切得很细来做流水线线程块的流水深度可以做得很深半精度输入情况下Tensor Core以NVIDIA的硬件为参考的指令形态、warp级的矩阵分片方式都需要分别考虑尺寸怎么配、寄存器怎么预留都有讲究数据的N、C维固定情况下可以用更大的tile来摊薄调度和寻址开销。换句话说普通实现是在怎么用CUDA写出一个对的GEMMDeepGEMM是在怎么用硬件手册级的理解写出一个深度适配工作负载的GEMM。这一点直接决定了代码中各种看起来反直觉的写法的来源。3. 从零手写一个高性能GEMM内核的关键设计点先把结论放出来高性能GEMM内核的优化永远围绕四个字——复用与流水。前者决定你能不能在正确的层次上把数据反复利用后者决定芯片上的各种执行单元能不能一直有事做。下面按我自己复现DeepGEMM思路的过程把几个核心设计点拆开讲。3.1 分块策略为什么Tile尺寸不能拍脑袋定GEMM的每一个输出元素都需要访问A矩阵一整行和B矩阵一整列。如果直接在全局内存上做每个元素都要重新读取数据内存吞吐很快就成了瓶颈。所以高性能实现的第一件事就是分块把一个大的输出矩阵切成很多个小块每个线程块负责一个输出块这个输出块的计算所需要用到的A和B数据块先取到共享内存里供块内多次循环复用。分块尺寸的选择非常关键。取太小共享内存里放的数据不足以支撑足够的计算复用带宽还是瓶颈取太大一个线程块占用的资源过多SM上能同时驻留的块数量下降调度延迟没法隐藏。这次拆解DeepGEMM时我关注到的核心tile配置是128×128左右的输出块尺寸每个线程块内部再按warp粒度切分成更小的块来映射到具体线程。为什么是128以Tensor Core指令为例常见的wmma指令形态是16×16×16m16n8k8这类形态也很常见。输出块选128意味着在m方向和n方向各能整除16不会出现浪费的边界处理逻辑同时128这个尺寸在寄存器层面也能比较自然地容纳不会因为分片过多导致线程间通信开销暴增。理论再好最后落实到具体硬件上还是需要对着profiler一个尺寸一个尺寸试的。3.2 数据搬运的隐藏技巧用异步拷贝把内存延迟藏到计算背后传统做法是先把数据从全局内存拷到共享内存然后线程再读共享内存去计算。这个过程中拷贝和计算是串行的计算单元在此期间只能等待。改进方向是让数据搬运在后台进行当前一个tile在计算时后一个tile的数据已经通过异步拷贝指令在路上了。这个思路在CUDA里最典型的实现方式就是双缓冲double buffering。共享内存开两份缓冲区一份给当前tile做计算一份接收下一个tile的数据交替使用。DeepGEMM在相关内容里设计流水线时用的正是这个思路。代码结构大概可以这样理解// 伪代码示意实际实现需要处理边界和同步 for (int k 0; k K_tiles; k) { if (k K_tiles - 1) { cp_async(buffers[next], A_block_next, B_block_next); } compute_from(buffers[current], accum); current ^ 1; next ^ 1; commit_group(); wait_group(); }按我实测的经验光是一个双缓冲改造就有机会带来20%-50%的性能提升具体取决于原来kernel的瓶颈类型。如果你的kernel本身就卡在计算密集度不够、算术单元没有吃满那么数据搬运的延迟隐藏会立竿见影如果已经吃满了这个优化收益就会弱一些。3.3 寄存器级复用一次加载多次计算共享内存带宽同样不是无限的。如果每次计算都要从共享内存读操作数共享内存的带宽也会成为瓶颈。进一步的做法是把数据放进寄存器每个线程负责的小块输出对应的A片段和B片段尽量一次性加载到寄存器里后续循环中反复使用。Tensor Core上的warp级矩阵指令有个很关键的特性它要求每个线程持有A、B矩阵的特定分片且这些分片要按固定的布局存在于寄存器中。所以写的代码要严格按照指令要求去铺数据多一个转置或shuffle都会带来额外开销。这里面比较考验功力的是片段切分每个线程算输出T一个2×1的小块到底A矩阵给这个线程分配哪些元素、B矩阵分配哪些需要根据指令形状和线程位次仔细推导。3.4 累加器布局与精度补偿半精度算GEMM最怕的是累加溢出和精度劣化。Tensor Core硬件本身支持FP32累加器也就是说矩阵乘法的中间累加过程是以FP32精度进行的两个FP16/FP16或者BF16/BF16乘完后的结果累加到FP32的寄存器里最后再写回输出时按需转回目标精度。这个机制在写kernel的时候几乎没有额外成本完全是硬件提供的特性。但要注意一个问题BF16的尾数位比FP16少乘法结果的精度天然受限。在训练场景中如果直接用BF16做全部前向计算某些对精度敏感的模型可能会出问题。所以DeepGEMM这类深度优化内核在工程上的落地往往还要配合混合精度策略某些层用FP16/BF16做GEMM计算但在累加器、归一化、残差连接这些精度敏感的位置用FP32保存中间结果。这一点在复现时被很多人忽略了导致同样的kernel在不同模型上效果天差地别。3.5 面向Tensor Core的指令调度不要小看指令排列顺序在写CUDA内核的时候指令顺序看起来是编译器帮你搞定的但真正做极致性能优化时你会意识到编译器的指令调度策略和手写调度之间的差距。Tensor Core的矩阵指令虽然吞吐很高但它对硬件流水线的占用形态和普通FMA指令不一样不能简单地跟数据搬运指令交替着放。这个项目里比较有价值的调度经验是把矩阵乘指令尽量连续地发出去中间避免插入太多依赖性的地址计算、比较指令把地址计算尽可能挪到不影响核心计算流水的位置。用CUDA的调度原语比如__pipeline相关的机制以及cp.async的group提交机制来批量管理异步拷贝让硬件可以在指令队列里看到足够多的独立任务从而把各级流水线都填满。4. 实测性能对比与优化效果验证4.1 我的测试环境与对照组设置为了验证DeepGEMM这类优化思路的含金量我特地搭了一个对比测试环境。硬件是老熟人了一块消费级GPU这里就用某N卡型号的通用描述来取代具体型号系统环境是标准的Linux 最新驱动编译工具链用CUDA最新稳定版本。矩阵规模我选了两种有代表性的大GEMMMNK4096模拟全连接层和较大规模的密集计算窄长GEMMM512N4096K4096模拟Transformer中间层常见的非对称shape。对照组的设置包括三档naive实现每个人都能写的最简单三重循环版本、中等优化版本只做shared memory tiling、DeepGEMM完整优化版本。所有代码用同一个编译选项避免编译器优化差异影响比较公平性。4.2 结果数据与瓶颈分析先看大GEMM的结果。naive实现自然不用多说性能大概只有理论峰值的3%左右纯粹是被全局内存带宽按在地上摩擦。中等优化版本做好shared memory分块之后性能飙升到了大约理论峰值的40%-50%这个时候瓶颈已经不是内存带宽了而是共享内存带宽和计算指令的调度效率。到了DeepGEMM完整优化版本实测性能提升到了理论峰值的80%以上如果把精度放宽到某些快速模式甚至能逼近90%。窄长GEMM的情况更有意思。naive和中等优化版本和大GEMM的表现趋势差不多但DeepGEMM优化版本的收益更明显。原因在于窄长场景下K维度特别长流水线优化的空间更大。K维的每个小块都可以在后台异步搬运主计算流全程不空等延迟隐藏的效果非常显著。下面这个表格是我这次测试的总结给个直观的对比实现版本大GEMM实测TFLOPS窄长GEMM实测TFLOPS主要瓶颈naive三重循环极低个位数百分比极低个位数百分比全局内存带宽shared memory tiling理论峰值约40%-50%理论峰值约35%-45%共享内存带宽与调度DeepGEMM完整优化理论峰值80%左右理论峰值85%左右指令发射与资源占用要注意的是这个比例是相对我手上这块显卡的理论峰值来算的不同硬件上比例会变化。但总体趋势相当稳定优化深度越深窄长shape的收益优势越明显。4.3 用Profiler找瓶颈的方法如果只给结论不给方法那这篇文章就不够味了。这里分享一个我自己常用的定位流程。第一步先用GPU性能分析工具抓kernel运行时间看看时间占比最高的kernel是哪个。第二步分析kernel内部的瓶颈类型到底是memory bound还是compute bound。这个通过对比实际带宽/占用率和理论峰值就能判断。第三步如果是memory bound优先优化数据复用和异步搬运如果是compute bound优先检查指令调度、是否有不必要的类型转换。这里有个很容易被忽略的点分析工具的Overhead可能影响到时序数据。小kernel尤其明显开太高采样的profiler会把kernel本身跑慢好几倍。我一般是先用工具把range和kernel整体耗时看一遍觉得可疑了再关掉profiler用cudaEvent做手动计时交叉验证。毕竟优化目标永远是真实端到端的性能不是分析工具报告出来的好看数字。5. 避坑实录我在复现DeepGEMM时踩过的五个坑5.1 共享内存Bank Conflict的隐形惩罚复现过程中第一个让我头疼的是bank conflict。共享内存在物理上被分成了32个bank如果同一warp的线程访问同一bank的不同地址就会发生冲突硬件需要把访问串行化。这个问题最隐蔽的地方在于代码逻辑完全正确结果也完全正确但性能就是上不去你很难第一眼就看到问题在这里。排查方法也比较传统但有效在关键循环处手动检查每个warp的访问地址分布或者把共享内存的布局改成padding版本比如每个row多申请几个元素的偏移让同一行元素的bank分布被错开。这个padding的技巧在DeepGEMM的实现里也用得很普遍。你要是发现自己写的GEMM kernel性能比同配置的参考实现差一截先去查bank conflict大概率能发现问题。5.2 寄存器溢出导致本地内存拖慢一切寄存器文件是芯片上最快的存储但数量有限。如果tile切得太大、每个线程持有的数据太多编译器就会把一部分寄存器变量溢出到本地内存。本地内存虽然在指令层面看起来还是像普通内存访问但实际上是走缓存和DRAM的速度比寄存器慢了不止一个量级。我在调256×256大tile时就遇到了这个问题。表面上寄存器占用率不高但运行时间反而变慢。查了编译报告之后发现编译器默默地生成了大量的本地内存访问指令。解决思路是把tile调回合理尺寸同时手动限制每个线程持有的A、B片段数量让编译器的寄存器分配压力保持在安全范围内。这里针对不同GPU架构限制的阈值不一样最好用编译报告里的寄存器数同步验证。5.3 K维切分太小导致流水线空泡流水线隐藏延迟的核心是计算和搬运并行。如果你的K维切分块特别小每个异步拷贝的数据量很少很快就能拷完但计算阶段还远远没结束流水线实际上处于空转状态。这就像餐厅里配菜师傅切菜太快厨师炒菜跟不上切好的菜堆在台面上白白等着。实际操作中应该让每一轮流水线中搬运的时间尽量和计算时间匹配。切块太大则搬运时间超过计算时间切块太小则流水线切换开销过大。DeepGEMM这类项目里常见的做法是先定一个基线切分尺寸然后用profiler测每一阶段的实际耗时再手动做一轮二分搜索式的调整。如果框架允许也可以把这个tile尺寸设计成编译期常量通过不同二进制之间的切换来找最优配置。5.4 半精度格式的精度暗坑FP16和BF16都是16位但精度特性差别很大。FP16的尾数多、指数范围小适合数值范围变化不大的场景BF16的指数范围和FP32一致但尾数少适合防止溢出、对精度相对宽容的场景。如果代码在两种格式之间切来切去却没有意识到它们的量化误差不同训练或者推理结果很容易出现莫名的偏差。我在一个模拟的量化推理项目里就吃过这个亏BF16输入的GEMM kernel跑得很快但最终精度比FP16低了一个数量级。后来排查发现是该层对数值精度极其敏感BF16的尾数位不足以支撑。这个问题的解决方法不是换回FP16而是在关键层用FP32做部分累加和补偿让精度敏感的位置保留足够的信息。5.5 编译器自动优化与手写内核的博弈编译器在-O3级别下确实会自动做一些循环展开、指令重排但它的优化逻辑是保证正确性的前提下尽量快而不是在这个特定的GEMM场景里尽量快。所以很多手写优化看起来像是破坏了编译器优化——其实不是而是你在把编译器的通用策略替换成针对性的专有策略。比如手动展开循环可以减少循环控制开销但同时会增加寄存器压力手动安排数据布局可以让向量化加载更高效但也可能让代码变得不那么可读。我见过有人执着于纯编译器优化拒绝使用任何手写调度结果性能始终上不去。真实项目里手写优化和编译器优化是配合关系先用编译器的自动向量化和循环展开做基线然后针对热点循环做手工干预最后用profiler对比验证每一项优化是否真的有效。6. 从DeepGEMM中提炼的通用优化方法论6.1 先用Profile建立基线再谈优化很多人拿到一个GEMM优化的任务第一反应是翻手册找最快的指令然后直接上手写。我的经验是反过来的先不写任何花哨的代码把最简单的实现跑一遍用profiler看清楚数据到底是怎么流动的、瓶颈在哪。这个基线数据是所有后续优化决策的锚点。比如一个简单的deep learning框架里的GEMM算子M1024, N1024, K4096naive实现花1000微秒。先看带宽利用率和计算利用率如果算下来理论需要的数据搬运时间只要50微秒而实际kernel耗时1000微秒那你该做的是减少访问、增加复用而不是去调指令调度。如果计算利用率已经接近90%了那你要关注的是怎么把更多的有效计算塞进流水线而不是再去做数据复用。6.2 把硬件特性变成代码设计的思维方式普通GEMM实现和DeepGEMM这类深度优化实现之间最大的差距不是代码技巧而是思维方式。普通实现是想好算法然后用代码表达算法深度优化是先吃透硬件的存储层次、指令形态、调度机制再反过来想算法应该长成什么样子。举一个最简单的例子矩阵是按行主序存储的访问A矩阵的一行和B矩阵的一列前者连续、后者跳变。如果直接按原始布局计算B列的访问会带来大量的cache miss。深度优化实现的做法是先把B矩阵做了分块转置或者从一开始就用列主序或者特殊layout让B块在shared memory里也能连续访问。这个决策的正确性完全来自对内存访问模式和硬件cache行为的理解。6.3 什么时候该果断放弃手写优化这一节是给所有人的清醒剂。手写优化能力很强但不是在所有场景都值得做。如果你的项目只是偶尔调一次shape不固定的GEMM或者矩阵规模太小、kernel启动开销本身就占大头那手写优化的投入产出比很低。这时候直接用官方库配合框架已有的融合能力反而是最理性的选择。我的个人判断标准是同一个GEMM操作在一个生产级系统里被调用的次数是否超过百万次/天如果是手写优化能带来1%的提升也是值得的如果不是把时间花在分析整体数据流、减少冗余计算、优化内存拷贝上收益往往更明显。说白了优化本身也是有性价比的。6.4 扩展GEMM优化思路在非矩阵场景的迁移DeepGEMM里体现的思想并不局限于GEMM。数据分块的思想可以用在卷积的im2col转化和Winograd变换上异步拷贝和双缓冲的流水线思想可以用在Embedding层的大规模查表、Attention中的KV Cache读取上寄存器级的数据复用思想可以用在很多elementwise算子的融合上。我在另一个序列处理项目里就参照了它的双缓冲思路把一个BatchNorm激活函数Residual融合算子的性能提升了接近30%。这个算子里根本没有矩阵乘法但优化的核心逻辑完全一致——减少中间数据写回全局内存的次数让数据在寄存器里尽量多待一会儿。所以这篇文章讲的是GEMM但方法论的适用范围是几乎所有深度学习算子优化。7. 实操经验总结关于DeepGEMM项目的一点心得整个DeepGEMM项目让我感触最深的一点是高性能计算的优化本质上是在跟系统的默认假设对抗。默认的库假设你不知道自己的工作负载默认的编译器假设通用逻辑比特殊逻辑更值得优化默认的数据布局假设访存模式是均匀分布的。而一个手写内核的价值就在于把这些默认假设全部推翻再基于真实场景重新设计。我在这次复现过程中最大的收获不是那几个性能提升的百分比数字而是真正建立起了从硬件倒推代码设计的思维方式。以前我也觉得GEMM很神秘CPU上写过高性能矩阵乘GPU上用过库函数但总隔着一层纱。现在自己从tile切分、流水线设计、寄存器分配到指令调度完整地走了一遍再回头看那些官方文档里看似枯燥的硬件特性说明确实能读出味道来了。如果你也想动手试一次我建议从一个小目标开始先写一个正确的GEMM再优化到比naive版本快5倍然后再挑战DeepGEMM里的高级特性。中间每一步都用profiler记录数据别凭感觉。矩阵乘法这个题目老归老但常做常新——不同硬件、不同精度、不同shape下永远有可以继续榨出来的性能空间。