“GEMM”这个缩写凡是碰过深度学习框架源码或者写过推理引擎的人都不会陌生General Matrix Multiplication通用矩阵乘法。我在开发模拟项目“DeepGEMM”时想做的其实是一件很朴素的事不依赖任何第三方库手写一个能在主流GPU上把矩阵乘法性能压榨到接近硬件理论峰值的内核同时把整个优化思路、踩坑过程记录下来。这篇文章就是基于那次实践整理出来的完整复盘适合正在学习CUDA内核优化、或者工作中需要手工调GEMM性能的开发者阅读。我会从“为什么GEMM值得折腾”讲起一路拆到分块、双缓冲、张量指令这些底层细节最后给出可直接参考的代码骨架和排障清单。1. 为什么非要折腾GEMM不少朋友看到“手写GEMM”第一反应是调库不就行了吗现实中确实如此大多数场景下直接调用现成的高性能数学库就够了。但一旦你开始做推理引擎定制、模型算子融合或者要支持某些库不友好的形状事情就变了。1.1 GEMM在深度学习里的实际面孔先想一个问题一个Transformer模型从输入到输出计算量到底花在哪我做过一次粗糙的统计拿一个中等规模的GPT风格模型为例前向推理里超过90%的乘加运算都落在各种矩阵乘法上。注意力里的QKV投影是GEMM注意力输出投影是GEMMFFN的两个全连接层还是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:k0BK] 和 B 的 [k0:k0BK][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.cuCUDA内核实现。bench.cu基准测试对比朴素实现和一个参考实现。test.cu正确性校验随机生成矩阵与朴素版本对比。内核启动参数是关键。我针对“M8192N8192K8192”的典型大矩阵线程块尺寸设为128×128每个线程计算8×8的寄存器tileK方向每次切16。这样算下来每个线程块包含256个线程线程块网格grid.x N / 128grid.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 对比结果实现版本计算时间相对值达到理论峰值比例备注朴素三重循环100x0.5%几乎完全被访存卡死分块共享内存寄存器tile约8x7%缓解了全局访存压力加入张量指令双缓冲约1.2x88%接近硬件峰值从这个表能明显看到每个优化阶段消除的瓶颈都不同。朴素版本输在内存访问上分块版本赢在数据复用而最后的张量指令把计算吞吐拉满。如果你看到自己的内核性能百分比长期停在个位数大概率是某一层瓶颈没解决而不只是浮点运算不够快。5.3 不同形状的敏感度我又测试了一组形状差异很大的矩阵包括M1的GEMV形状、M64的小批形状、M4096的超大形状。结论比较典型超大矩阵主内核性能稳定利用率高。M64、N8192、K8192虽然也能用主内核但边角浪费严重性能约下降25%。M1、N8192、K8192主内核完全不合适应该走专门优化过的矩阵向量乘实现。即便换路径性能也比大矩阵GEMM低很多。这个敏感度不是bug它来自并行度不足。GEMM要快需要足够多的线程块填满GPU的所有计算单元。M太小意味着线程块数量太少大量硬件资源闲置。群里经常有人问“为什么我的GEMM在某形状下那么慢”九成是这个问题。通用库通过维护多个内核来缓解但总有一些形状怎么都难做好。6. 常见问题与排查清单好代码也写了性能也测了接下来是实践中最容易被卡住的部分。我把过程中遇到的高频问题整理成一份速查表方便你直接对照排查。问题现象可能原因排查与解决性能利用率只到5%上下没有分块或者共享内存压根没用到检查是否每次计算都直接读全局内存补上分块循环共享内存吞吐上不去没有做寄存器tile线程每次都在读共享内存把每个线程的累加器改成8×8的寄存器数组性能上到一半遇到平台期线程数太少/太多显存带宽饱和用ncu看单元活跃度尝试不同块尺寸128或256数值对不上累加顺序变化、FP16溢出改为FP32累加检查布局转换是否改写了数据性能突然大幅下降有bank conflict加入1个float的共享内存padding再次实测某些形状性能特别差并行度不足/形状不匹配换用fallback路径对小M做专门优化程序偶尔结果不对缺少同步双缓冲竞态检查__syncthreads位置预估缓冲使用时机张量指令没有提升数据布局不满足要求阅读硬件指令文档确认tile排布必要时预重排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 Computencu它给出的指标对定位瓶颈非常关键。我日常主要看几个指标计算单元利用率SM busy反映计算资源是否吃满。内存吞吐Memory Throughput反映显存带宽是否成为瓶颈。共享内存bank冲突数Shared Memory Bank Conflicts帮助定位共享内存访问踩坑。warp占用率Occupancy判断线程块是否足够多能否掩盖延迟。把这几个指标组合起来基本能定位90%的GEMM性能问题。6.4 终极调试技巧如果实在找不到性能瓶颈还有一个笨办法不断地“删功能”。先把双缓冲去掉看性能变化再把张量指令降级为普通乘加看性能变化再把分块改小一倍。通过这种二分法对比你能快速定位到底是哪一层优化没生效哪一层反而拖了后腿。这套方法我一直用到现在比肉眼猜代码靠谱得多。最后补一句这次写DeepGEMM最大的收获是我真正理解了为什么那些高性能库能跑得那么快它们本质上不是在算矩阵而是在管理数据流动。一开始我满脑子都是“怎么让乘加更快”后来才意识到真正的跷跷板在“怎么让数据在正确的时间、出现在正确的位置”。踩过几次坑之后我对GPU计算的整体认知有了很大提升也更愿意去读那些底层资料了。如果你也想试试手写GEMM不用贪多先在固定形状下跑通一遍分块共享内存寄存器tile再逐步加张量指令和双缓冲。等你自己把第一个接近硬件峰值的GEMM内核调出来时那种满足感绝对值得。
“GEMM”这个缩写凡是碰过深度学习框架源码或者写过推理引擎的人都不会陌生General Matrix Multiplication通用矩阵乘法。我在开发模拟项目“DeepGEMM”时想做的其实是一件很朴素的事不依赖任何第三方库手写一个能在主流GPU上把矩阵乘法性能压榨到接近硬件理论峰值的内核同时把整个优化思路、踩坑过程记录下来。这篇文章就是基于那次实践整理出来的完整复盘适合正在学习CUDA内核优化、或者工作中需要手工调GEMM性能的开发者阅读。我会从“为什么GEMM值得折腾”讲起一路拆到分块、双缓冲、张量指令这些底层细节最后给出可直接参考的代码骨架和排障清单。1. 为什么非要折腾GEMM不少朋友看到“手写GEMM”第一反应是调库不就行了吗现实中确实如此大多数场景下直接调用现成的高性能数学库就够了。但一旦你开始做推理引擎定制、模型算子融合或者要支持某些库不友好的形状事情就变了。1.1 GEMM在深度学习里的实际面孔先想一个问题一个Transformer模型从输入到输出计算量到底花在哪我做过一次粗糙的统计拿一个中等规模的GPT风格模型为例前向推理里超过90%的乘加运算都落在各种矩阵乘法上。注意力里的QKV投影是GEMM注意力输出投影是GEMMFFN的两个全连接层还是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:k0BK] 和 B 的 [k0:k0BK][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.cuCUDA内核实现。bench.cu基准测试对比朴素实现和一个参考实现。test.cu正确性校验随机生成矩阵与朴素版本对比。内核启动参数是关键。我针对“M8192N8192K8192”的典型大矩阵线程块尺寸设为128×128每个线程计算8×8的寄存器tileK方向每次切16。这样算下来每个线程块包含256个线程线程块网格grid.x N / 128grid.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 对比结果实现版本计算时间相对值达到理论峰值比例备注朴素三重循环100x0.5%几乎完全被访存卡死分块共享内存寄存器tile约8x7%缓解了全局访存压力加入张量指令双缓冲约1.2x88%接近硬件峰值从这个表能明显看到每个优化阶段消除的瓶颈都不同。朴素版本输在内存访问上分块版本赢在数据复用而最后的张量指令把计算吞吐拉满。如果你看到自己的内核性能百分比长期停在个位数大概率是某一层瓶颈没解决而不只是浮点运算不够快。5.3 不同形状的敏感度我又测试了一组形状差异很大的矩阵包括M1的GEMV形状、M64的小批形状、M4096的超大形状。结论比较典型超大矩阵主内核性能稳定利用率高。M64、N8192、K8192虽然也能用主内核但边角浪费严重性能约下降25%。M1、N8192、K8192主内核完全不合适应该走专门优化过的矩阵向量乘实现。即便换路径性能也比大矩阵GEMM低很多。这个敏感度不是bug它来自并行度不足。GEMM要快需要足够多的线程块填满GPU的所有计算单元。M太小意味着线程块数量太少大量硬件资源闲置。群里经常有人问“为什么我的GEMM在某形状下那么慢”九成是这个问题。通用库通过维护多个内核来缓解但总有一些形状怎么都难做好。6. 常见问题与排查清单好代码也写了性能也测了接下来是实践中最容易被卡住的部分。我把过程中遇到的高频问题整理成一份速查表方便你直接对照排查。问题现象可能原因排查与解决性能利用率只到5%上下没有分块或者共享内存压根没用到检查是否每次计算都直接读全局内存补上分块循环共享内存吞吐上不去没有做寄存器tile线程每次都在读共享内存把每个线程的累加器改成8×8的寄存器数组性能上到一半遇到平台期线程数太少/太多显存带宽饱和用ncu看单元活跃度尝试不同块尺寸128或256数值对不上累加顺序变化、FP16溢出改为FP32累加检查布局转换是否改写了数据性能突然大幅下降有bank conflict加入1个float的共享内存padding再次实测某些形状性能特别差并行度不足/形状不匹配换用fallback路径对小M做专门优化程序偶尔结果不对缺少同步双缓冲竞态检查__syncthreads位置预估缓冲使用时机张量指令没有提升数据布局不满足要求阅读硬件指令文档确认tile排布必要时预重排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 Computencu它给出的指标对定位瓶颈非常关键。我日常主要看几个指标计算单元利用率SM busy反映计算资源是否吃满。内存吞吐Memory Throughput反映显存带宽是否成为瓶颈。共享内存bank冲突数Shared Memory Bank Conflicts帮助定位共享内存访问踩坑。warp占用率Occupancy判断线程块是否足够多能否掩盖延迟。把这几个指标组合起来基本能定位90%的GEMM性能问题。6.4 终极调试技巧如果实在找不到性能瓶颈还有一个笨办法不断地“删功能”。先把双缓冲去掉看性能变化再把张量指令降级为普通乘加看性能变化再把分块改小一倍。通过这种二分法对比你能快速定位到底是哪一层优化没生效哪一层反而拖了后腿。这套方法我一直用到现在比肉眼猜代码靠谱得多。最后补一句这次写DeepGEMM最大的收获是我真正理解了为什么那些高性能库能跑得那么快它们本质上不是在算矩阵而是在管理数据流动。一开始我满脑子都是“怎么让乘加更快”后来才意识到真正的跷跷板在“怎么让数据在正确的时间、出现在正确的位置”。踩过几次坑之后我对GPU计算的整体认知有了很大提升也更愿意去读那些底层资料了。如果你也想试试手写GEMM不用贪多先在固定形状下跑通一遍分块共享内存寄存器tile再逐步加张量指令和双缓冲。等你自己把第一个接近硬件峰值的GEMM内核调出来时那种满足感绝对值得。
和strcasecmp()使用总结)
