很多时候一个项目的瓶颈并不是模型结构写不写得出来而是底层的矩阵乘法能不能跑满算力。DeepGEMM 这个项目起因很简单我手头有个推理场景模型的层数和宽度都谈不上夸张可实测延迟就是压不下去。逐个算子拆完以后发现几乎所有的等待时间都落在同一个东西上——通用矩阵乘法也就是 GEMMGeneric Matrix Multiply。Transformer 里的 attention 打分、FFN 的升维降维再远一点卷积如果用 im2col 的形式展开本质上也是一堆大矩阵相乘。所以我一直觉得想做好深度学习性能优化第一关就是把 GEMM 这件事弄明白。这篇内容会围绕我把 DeepGEMM 从零写出来、再一点点盘到能用的过程展开。适合谁看两种人一种是刚接触高性能计算想知道矩阵乘法怎么就变成了瓶颈的读者另一种是已经在用各种数学库但遇到自己不理解的性能上限想拆开看看底层发生了什么的人。我不会只给结论会把为什么这样设计、为什么这里慢、为什么换一种写法就快了一截的过程讲出来。1. GEMM 在深度学习里的出场频率说它不是瓶颈的人多半没认真数过1.1 一次前向计算里GEMM 到底被调用了多少次很多人一听到 GEMM第一反应是这不就是线性代数里的矩阵乘法吗我早就会了。会算和会调度是两回事。我先举个例子假设一个很普通的 Transformer decoder 层输入维度是 768注意力头数 12FFN 中间层维度是 3072。比较大的那几块开销基本全是 GEMM计算 Query、Key、Value输入矩阵乘以对应的投影矩阵三个形状分别是[序列长度, 768] × [768, 768]。注意力分数Query 乘以 Key 的转置形状大约是[序列长度, 768] × [768, 序列长度]。注意力加权求和注意力分数矩阵乘以 Value形状是[序列长度, 序列长度] × [序列长度, 768]。FFN 第一层[序列长度, 768] × [768, 3072]。FFN 第二层[序列长度, 3072] × [3072, 768]。这还只是一层里必须的 GEMM如果再算上位置编码、LayerNorm 这些小算子单层就有五六次大矩阵乘。模型一加深序列长度再拉长你就能算出来一个 batch 的前向时间里GEMM 占据的浮点运算量常常超过九成。这还没考虑反向传播反向的时候梯度也要通过同样的矩阵运算量再翻两倍左右。所以我说谁如果说深度学习性能瓶颈不是矩阵乘法那大概率是还没把一个真实模型的算子拆开数过一遍。1.2 并不存在笼统的矩阵乘法变快这回事既然 GEMM 这么重要那直接调一个现成的优化库不就行了现实没这么简单。深度学习里的 GEMM 形状非常不统一有的矩阵特别宽有的特别高还有的在推理时会因为 Padding 策略出现奇怪的维数。通用数学库为了覆盖各种场景通常会做多套算法分支但对某个具体的胖矩阵或瘦矩阵形状未必选了最适合你的那套策略。还有一个问题就是快的定义不一样。在训练阶段通常数据能完整塞进显存批处理规模又大这时候 GEMM 的核心竞争点是浮点运算峰值也就是算力有没有跑满。但在推理阶段batch 往往只有 1 到 4矩阵的列数和行数也比训练时小很多这时候拼的是内存带宽、启动延迟、算子之间的融合程度。你没有自己的实现就只能被标准库的默认策略牵着走。DeepGEMM 的目标就是针对我自己遇到的这类深度学习常见形状做一套能在不同条件下自适应切换策略的小型 GEMM 引擎。它未必能打赢所有情况但必须做到我明确知道每一步在做什么。1.3 我决定不做什么先忘掉那些玄学优化在动手之前我先把一些花哨的方向放到一边。比如 Winograd 这种通过减少乘法次数来加速卷积的算法比如基于 Strassen 的矩阵分解优化又比如针对极端稀疏结构定制的稀疏 GEMM。这些方法都有用但它们的适用前提很苛刻比如 Winograd 对卷积核尺寸和 padding 有约束稀疏 GEMM 则依赖稀疏模式乱套反而可能出现误差放大或访存局部性变差的问题。DeepGEMM 的第一版设计原则是先保证密集矩阵下的乘加语义完全一致误差可预测然后才谈性能。所谓最快不是靠奇技淫巧而是靠把数据流动、指令排列、并行划分这些基本功做到位。说白了GEMM 优化的本质是让计算单元在每一个时钟周期都有东西可算让内存系统在每一次请求里都搬运有用的数据。表面上是数学问题骨子里是硬件调度问题。抓住这条线后面所有的决策就有了统一的判断依据。2. DeepGEMM 的核心设计取舍分块大小、数据布局与指令调度2.1 分块大小为什么是 64 或者 128而不是越大越好矩阵乘法朴素写法是三重循环外层遍历输出矩阵的行中层遍历列内层做累加。这个写法的致命问题在于CPU 缓存也好GPU 的共享内存也好容量都很有限。拿一个[2048, 2048]的矩阵来说按单精度浮点存储光一个矩阵就要 16MB 空间这还没算另外两个矩阵。如果你用朴素遍历内层每次读 A 的一行和 B 的一列这两个访存流的空间距离很远缓存基本上一轮下来就被刷没了。解决思路和铺瓷砖很像与其把整面墙一次铺完不如切成一块一只寸刚好合适的小瓷砖贴完一块再贴下一块。GEMM 里的瓷砖就是分块tiling把大矩阵切成BM × BK、BK × BN、BM × BN的小块让参与计算的三个小块尽量同时留在缓存或共享内存里。我实际试下来分块大小一般取 64 或者 128 比较合适。太小比如 16×16分块数量太多每个分块都要重新加载数据循环开销和调度开销反而变大太大比如 512一块数据就把缓存或者共享内存塞满了下一块数据还得等前面的计算结果写回内存才能加载白白空转。举个直观例子假设某级缓存的容量是 256KB三个小块如果选择 64×64每个小块单精度是 16KB三个加起来 48KB很宽裕选 128×128每个小块 64KB三个加起来 192KB还能接受选 256×256每个小块 256KB三个 768KB已经超了。所以分块大小的上限由缓存容量反推下限由并行粒度决定并不是拍脑袋选的。2.2 数据布局为什么说连续内存比看起来整齐更关键分块解决的是缓存层的问题下一层问题是数据在内存里到底怎么摆放。绝大多数人默认的二维数组是按行连续存储的也就是A[i][j]的下一个元素是A[i][j1]。这在很多场景没问题但在 GEMM 里如果我们要反复读取 B 矩阵的一列那每一列的元素在内存里是跳着分布的一跳就是一行的大小。跳跃访问对缓存和预取都不友好明明读了一整块缓存行里面可能只有几个字节有用。我偏好的方式是提前做一次数据重排把参与运算的 B 矩阵改成按分块存放。说白了就是把 B 先拆成若干BK × BN的小块每个小块内部再按列变成按行连续存放。这样后面每一次内层循环读取 B 的小块时都是连续内存访问预取器也能识别出规律。重排的开销只是一次性的但对几千上万次小块读取来说收益非常可观。连续访存的好处一次缓存行读取可以被完整利用不会出现读 64 字节只用了 4 字节的浪费。减少 TLB 缺失连续大块访问的页表项数量更少不会因为反复横跳把快表打爆。数据对齐分块大小选 64 或 128 这种对齐友好的值能让向量加载指令直接一次性把数据搬进寄存器。2.3 指令调度别让计算单元闲着等数据分块和数据布局解决了数据在哪的问题指令调度解决的是数据到了以后计算单元能不能立刻用上的问题。现代处理器的乘加指令延迟通常有好几个周期比如一条fma从发射到出结果需要四到六个周期但它可以每周期发射一条甚至更多。这里的核心技巧就是通过手动展开循环把相互独立的乘加计算穿插排列让第一条 fma 还在算的时候第二条、第三条已经发射出去了。我用一个很朴素的生活类比来解释你去食堂打饭如果每打一份饭都站在原地等食堂阿姨把菜盛好再走队伍会排到门口。但如果你先刷卡然后赶紧去打汤、拿筷子等回来时饭刚好盛好你就能一直处于有事情做的状态。CPU 里的乱序执行和寄存器重命名就像打饭流程里的并行窗口我们要做的是让代码里同时存在足够多的打汤、拿筷子这类独立任务。这就是循环展开的意义。实际实现时我常把内层累加循环展开成 4 路或 8 路。每路独立维护一个累加寄存器组最后再合并。这样代码看起来啰嗦但吞吐量能拉高一大截。不过这里有个代价就是寄存器会被累加变量占掉很多。如果你同时有 8 个累加变量每个变量又对应一个B块里的一小段那留给数据加载的寄存器就少了。所以展开路数、分块大小、寄存器数量这三者必须一起调不能只看单项。3. 从能跑到跑得快三轮迭代实现的完整复盘3.1 第一轮先保证结果是正确的DeepGEMM 的第一版我写得很简单基本就是把三重循环翻译成代码甚至没有特别在意循环顺序。这一版的作用只有一个和标准数学库的结果对齐验证我对数据布局、边界填充、累加顺序的理解没有错。这一版的速度当然不行因为内层循环每次做C[i][j] A[i][k] * B[k][j]时B 的列访问在内存中是跳跃的。这个阶段我不看性能只看误差。我会用随机矩阵跑很多轮比较我自己实现的结果和标准库结果把绝对误差控制在浮点累加顺序不同所允许的范围之内。如果你的第一版连正确性都保证不了后面所有优化都白搭因为性能调优最容易出现的问题就是优化了半天算出来已经是错的了。3.2 第二轮引入分块和循环重排第二版做的第一件事是把外层循环顺序从i-j-k改成以分块为单位的bk-bi-bn-kj-ki之类的顺序让 B 和 A 的读取都回到连续空间上。写代码的时候我没直接用分块库而是自己动手做了三个临时缓冲区把 A 的一个大块切成BM × BK的子块把 B 的大块切成BK × BN的子块然后把它们塞到缓存友好的位置上再做内层乘法。用伪代码表示大概是这样for (int bi 0; bi M; bi BM) { for (int bj 0; bj N; bj BN) { // 把 C 的分块清零 for (int bk 0; bk K; bk BK) { // 加载 A 和 B 的分块到临时缓冲 for (int i 0; i BM; i) { for (int k 0; k BK; k) { for (int j 0; j BN; j) { C[bii][bjj] A_tmp[i][k] * B_tmp[k][j]; } } } } } }当然上面只是个逻辑示意真正的内层循环我会再改成指针累加的形式避免每层都做下标计算。这一轮做完以后性能已经比第一版好了一个量级主要原因是缓存命中率显著上升。但离目标还差不少因为内层i-k-j的顺序依然有一层循环在横跨行与行之间并且没有做到指令级别的并行。3.3 第三轮手写乘加核心把并行度抠出来第三版我做了两个比较重要的决定把最内层的k循环展开成 4 路或 8 路每一路同时推进连续不同的k值。让每个输出分块不再只由一个线程负责而是进一步拆成多个子块划分给不同线程或不同核心执行。大分块负责把数据留在缓存小分块负责让多个计算单元同时忙起来。展开以后的核心逻辑看起来会有点吓人但本质就是在同一个循环里同时算好几个独立的输出位置// 以 4x4 输出子块为例的累加片段 for (int kk 0; kk BK; kk) { float a0 A_reg[i0][kk]; float a1 A_reg[i1][kk]; float a2 A_reg[i2][kk]; float a3 A_reg[i3][kk]; c00 a0 * B_reg[kk][j0]; c01 a0 * B_reg[kk][j1]; c02 a0 * B_reg[kk][j2]; c03 a0 * B_reg[kk][j3]; c10 a1 * B_reg[kk][j0]; // ... 类似地累加 c10 到 c33 }这段代码有两个目的第一是减少循环控制指令的开销第二是为 CPU 提供一组互不依赖的乘加指令让它能在等待一条指令结果的周期里继续发射其他指令。编译器有时候能看懂这些循环并帮你展开但很多时候它因为担心指针别名、循环边界等问题不敢做激进的变换。手写展开的好处是稳定可控坏处是代码可读性下降所以我会把展开的核心单独放一个函数里外面配好注释。3.4 一个容易被忽略的正确性验证步骤第三版的性能上来以后我反而更小心了。因为浮点加法不满足结合律展开以后累加顺序变了结果和朴素版本会有细微差别。这是允许的但必须保证误差在一个很小的范围内。我一般会把两种实现的结果差值除以标准库结果的最大绝对值得到一个相对误差。如果相对误差超过1e-4量级说明代码里可能有 bug比如某个分块被重复累加了或者 B 的转置逻辑写错了。我自己在实际调优中遇到过一个典型案例某个分块的边界不是整块大小需要对 A 或 B 的尾部做补零处理。我图省事直接把补零循环跳过了结果矩阵尾部一小部分结果是错的而且因为只有特定尺寸才会触发一开始根本看不出来。后来我加了边界形状的自动测试把 M、N、K 分别设为64、65、127这种非对齐尺寸才把问题暴露出来。这种非对齐测试要比只测对齐尺寸重要得多因为真实深度学习形状里序列长度和隐藏维度经常不是 64 的整数倍。4. 性能上不去时我按这个顺序排查瓶颈4.1 第一个坑矩阵不大但内存带宽先满了我印象很深的一次性能异常是分块和数据布局都做完以后一个 512×512 的 GEMM 跑出来的时间比我预想的高很多。浮点运算量不大按理说计算单元应该很轻松但整个内核的运行时间几乎和访存量成正比这时我意识到这个512×512的场景其实是被内存带宽卡住了而不是被计算单元卡住。判断方法很简单我会算一下理论访存量。输出是512×512个浮点数也就是 1MB 写回A 和 B 各 1MB按理想情况只读一遍的话总共大约 3MB 的访存。如果内存带宽在 20GB/s 左右理论上最少要 0.15ms 左右。如果实测在 0.4ms 以上说明至少有大量重复读取或者访问不连续。遇到这种情况我不会急着加指令展开而是先检查是不是某个小块被反复加载或者 B 的重排没有生效。访存优化永远是先于计算优化的因为计算单元再快喂不进去数据就是白搭。4.2 第二个坑并行分块不均导致部分核心闲着另一个大坑出现在多核并行场景。我有一次把一个1024×1024的矩阵按行切成两半让两个线程各算一半看起来负载非常均匀。但由于机器上不止两个核在跑系统里其他进程也在抢核心实测两个线程的完成时间能差到两倍。更麻烦的是如果只按行切列方向上没有并行那么每个线程内部依然是串行处理长条状的矩阵访存局部性也不好。后来我改成二维分块并行把输出矩阵切成一个网格行方向和列方向都参与划分每个线程处理一个格子。这样不仅负载更均匀而且每个格子都能对应到一段较小且完整的内存区域。有一点很关键如果矩阵尺寸不能被线程数整除不要让某些线程多了一块而其他线程少了一块最好统一用任务队列方式动态领取分块。谁算完谁去拿下一块而不是静态分配这样就算某些核心被系统调度拖慢了其他核心也能自动补上整体耗时更稳定。4.3 排查工具和步骤不是靠猜而是靠对照现象可能原因排查手段时间随矩阵尺寸线性增长但和计算量不成比例内存带宽瓶颈统计实际访存量与理论最小访存量对比多核并行后加速比很低并行划分不均衡或锁竞争用计时器分别统计每个线程的耗时分块变大性能反而下降缓存或共享内存溢出检查分块占用容量是否超过缓存阈值非对齐尺寸性能断崖下跌边界处理分支导致流水线中断对尾部形状单独做补零避免在热路径加判断我还习惯在调优后生成一份汇编或反汇编看一眼内层循环里是否出现了不必要的分支跳转。热路径上任何带if判断的操作都可能让流水线停顿。如果发现边界判断在内部我会把它挪出去要么用补零让边界变得整齐要么对主循环和尾部循环分开处理。5. DeepGEMM 的实测数据与扩展路线5.1 测试环境与相对结果我测试用的机器并没有刻意追求顶级配置目的只是想看 DeepGEMM 在常见硬件上的表现。这是一个双路服务器内存带宽大概在 40GB/s 量级CPU 支持单指令多数据和乘加融合指令。所有测试都用单精度浮点矩阵尺寸覆盖了深度学习里常见的小批次形状。下表是我拿 DeepGEMM 和系统自带的通用数学库做的对比数字是多次运行后的中位时间为避免绝对性能受具体芯片影响我用相对倍数来表示标准化到通用库耗时矩阵形状M×N×K场景标准库耗时DeepGEMM 耗时相对差距128×128×128小规模推理1.00.72DeepGEMM 快约 1.4 倍512×512×512中等规模1.00.81快约 1.2 倍1024×1024×1024偏计算密集1.00.94基本持平3072×768×768FFN 常见形状1.00.89快约 1.1 倍我解释一下这个结果小尺寸矩阵比拼的完全是启动开销和数据搬运策略DeepGEMM 的专用分块能减少很多无效的初始化和检查所以优势明显。到了 1024 这种大尺寸通用库的自动调优已经选了不错的算法我这边能持平就已经很不容易了毕竟对方有更多的专家配置做支撑。FFN 形状属于比较高瘦的矩阵通用库可能为典型正方形做了过多优化反而 DeepGEMM 因为重排策略匹配拿到了一点点优势。5.2 从结果里读出的两层信息第一层信息很简单不存在一个完全通用的最优 GEMM 实现形状不同瓶颈不同策略就得跟着变。第二层信息更重要优化 GEMM 的收益有两个来源一是减少冗余访存二是提高计算并行度而这两者在不同尺寸下权重完全不同。小尺寸时哪怕只少一次边界判断节省的时间占比都很大大尺寸时计算单元能不能满负荷运转才是决定成败的主线。这也就意味着如果你要复刻这样一个项目最值得投入精力的不是背诵某个分块大小而是建立一套自己的测试基线。每次改一个参数就记一次时间和误差做成一个简单的回归测试表。这样你能够看到每一次调整到底带来了收益还是回退而不是凭着感觉大改。5.3 后续可以扩展的方向低精度、稀疏和动态形状DeepGEMM 目前完成的是单精度密集 GEMM 的优化但显然这部分工作还能延伸出去。第一个方向是低精度也就是 FP16 甚至 INT8。深度学习推理阶段用低精度是主流做法可以把数据带宽减半甚至减到四分之一同时利用新一代硬件里专门的低精度矩阵计算单元。这里要小心的是数据溢出和累加精度问题不能简单地把输入截断成低精度否则一个 batch 的误差会被放大。第二个方向是稀疏化。现实中很多模型的权重经过剪枝后存在大量零值如果我们能识别出固定的稀疏模式直接跳过这些零值理论上可以获得好几倍加速。难点在于稀疏模式有结构化比如按块剪枝和非结构化之分后者对硬件不友好性能收益很容易被索引开销吞掉。第三个方向是动态形状。推理场景里序列长度往往不固定这要求 GEMM 内核不能在每次遇到新形状时都重新自动调优。更好的做法是把内核编译成多个候选版本运行时根据当前形状做一次快速选择。这部分我已经在 DeepGEMM 的路标里排上了但目前还没有做到稳定可发布的程度。最后说一点我做这个项目最大的体会矩阵乘法优化不是那种看一两篇理论文章就能掌握的技能它需要在具体机器上反复做实验亲手把一段代码从能跑打磨到跑得快中间那些报错、异常、性能回退才是真正让你理解系统的地方。如果哪天你发现自己的分块调来调去都差不多不要沮丧那恰恰说明你对当前场景已经接近硬件真实上限了。换个形状、换组参数再试新的一轮瓶颈又会冒出来。保持这种对着计数器较真的劲头比任何一个现成库都值钱。写到这里DeepGEMM 的第一阶段总结就差不多了。如果你也想试试建议先从最朴素的版本开始然后依次加数据重排、分块、展开和并行每一步都记录时间和正确性。这个流程本身就是高性能计算里最值得练的基本功。 
长河君 · 2026/10/24 10:20:17 · 阅读 1.2k DeepGEMM：从零优化GEMM，破解深度学习矩阵乘法性能瓶颈
很多时候一个项目的瓶颈并不是模型结构写不写得出来而是底层的矩阵乘法能不能跑满算力。DeepGEMM 这个项目起因很简单我手头有个推理场景模型的层数和宽度都谈不上夸张可实测延迟就是压不下去。逐个算子拆完以后发现几乎所有的等待时间都落在同一个东西上——通用矩阵乘法也就是 GEMMGeneric Matrix Multiply。Transformer 里的 attention 打分、FFN 的升维降维再远一点卷积如果用 im2col 的形式展开本质上也是一堆大矩阵相乘。所以我一直觉得想做好深度学习性能优化第一关就是把 GEMM 这件事弄明白。这篇内容会围绕我把 DeepGEMM 从零写出来、再一点点盘到能用的过程展开。适合谁看两种人一种是刚接触高性能计算想知道矩阵乘法怎么就变成了瓶颈的读者另一种是已经在用各种数学库但遇到自己不理解的性能上限想拆开看看底层发生了什么的人。我不会只给结论会把为什么这样设计、为什么这里慢、为什么换一种写法就快了一截的过程讲出来。1. GEMM 在深度学习里的出场频率说它不是瓶颈的人多半没认真数过1.1 一次前向计算里GEMM 到底被调用了多少次很多人一听到 GEMM第一反应是这不就是线性代数里的矩阵乘法吗我早就会了。会算和会调度是两回事。我先举个例子假设一个很普通的 Transformer decoder 层输入维度是 768注意力头数 12FFN 中间层维度是 3072。比较大的那几块开销基本全是 GEMM计算 Query、Key、Value输入矩阵乘以对应的投影矩阵三个形状分别是[序列长度, 768] × [768, 768]。注意力分数Query 乘以 Key 的转置形状大约是[序列长度, 768] × [768, 序列长度]。注意力加权求和注意力分数矩阵乘以 Value形状是[序列长度, 序列长度] × [序列长度, 768]。FFN 第一层[序列长度, 768] × [768, 3072]。FFN 第二层[序列长度, 3072] × [3072, 768]。这还只是一层里必须的 GEMM如果再算上位置编码、LayerNorm 这些小算子单层就有五六次大矩阵乘。模型一加深序列长度再拉长你就能算出来一个 batch 的前向时间里GEMM 占据的浮点运算量常常超过九成。这还没考虑反向传播反向的时候梯度也要通过同样的矩阵运算量再翻两倍左右。所以我说谁如果说深度学习性能瓶颈不是矩阵乘法那大概率是还没把一个真实模型的算子拆开数过一遍。1.2 并不存在笼统的矩阵乘法变快这回事既然 GEMM 这么重要那直接调一个现成的优化库不就行了现实没这么简单。深度学习里的 GEMM 形状非常不统一有的矩阵特别宽有的特别高还有的在推理时会因为 Padding 策略出现奇怪的维数。通用数学库为了覆盖各种场景通常会做多套算法分支但对某个具体的胖矩阵或瘦矩阵形状未必选了最适合你的那套策略。还有一个问题就是快的定义不一样。在训练阶段通常数据能完整塞进显存批处理规模又大这时候 GEMM 的核心竞争点是浮点运算峰值也就是算力有没有跑满。但在推理阶段batch 往往只有 1 到 4矩阵的列数和行数也比训练时小很多这时候拼的是内存带宽、启动延迟、算子之间的融合程度。你没有自己的实现就只能被标准库的默认策略牵着走。DeepGEMM 的目标就是针对我自己遇到的这类深度学习常见形状做一套能在不同条件下自适应切换策略的小型 GEMM 引擎。它未必能打赢所有情况但必须做到我明确知道每一步在做什么。1.3 我决定不做什么先忘掉那些玄学优化在动手之前我先把一些花哨的方向放到一边。比如 Winograd 这种通过减少乘法次数来加速卷积的算法比如基于 Strassen 的矩阵分解优化又比如针对极端稀疏结构定制的稀疏 GEMM。这些方法都有用但它们的适用前提很苛刻比如 Winograd 对卷积核尺寸和 padding 有约束稀疏 GEMM 则依赖稀疏模式乱套反而可能出现误差放大或访存局部性变差的问题。DeepGEMM 的第一版设计原则是先保证密集矩阵下的乘加语义完全一致误差可预测然后才谈性能。所谓最快不是靠奇技淫巧而是靠把数据流动、指令排列、并行划分这些基本功做到位。说白了GEMM 优化的本质是让计算单元在每一个时钟周期都有东西可算让内存系统在每一次请求里都搬运有用的数据。表面上是数学问题骨子里是硬件调度问题。抓住这条线后面所有的决策就有了统一的判断依据。2. DeepGEMM 的核心设计取舍分块大小、数据布局与指令调度2.1 分块大小为什么是 64 或者 128而不是越大越好矩阵乘法朴素写法是三重循环外层遍历输出矩阵的行中层遍历列内层做累加。这个写法的致命问题在于CPU 缓存也好GPU 的共享内存也好容量都很有限。拿一个[2048, 2048]的矩阵来说按单精度浮点存储光一个矩阵就要 16MB 空间这还没算另外两个矩阵。如果你用朴素遍历内层每次读 A 的一行和 B 的一列这两个访存流的空间距离很远缓存基本上一轮下来就被刷没了。解决思路和铺瓷砖很像与其把整面墙一次铺完不如切成一块一只寸刚好合适的小瓷砖贴完一块再贴下一块。GEMM 里的瓷砖就是分块tiling把大矩阵切成BM × BK、BK × BN、BM × BN的小块让参与计算的三个小块尽量同时留在缓存或共享内存里。我实际试下来分块大小一般取 64 或者 128 比较合适。太小比如 16×16分块数量太多每个分块都要重新加载数据循环开销和调度开销反而变大太大比如 512一块数据就把缓存或者共享内存塞满了下一块数据还得等前面的计算结果写回内存才能加载白白空转。举个直观例子假设某级缓存的容量是 256KB三个小块如果选择 64×64每个小块单精度是 16KB三个加起来 48KB很宽裕选 128×128每个小块 64KB三个加起来 192KB还能接受选 256×256每个小块 256KB三个 768KB已经超了。所以分块大小的上限由缓存容量反推下限由并行粒度决定并不是拍脑袋选的。2.2 数据布局为什么说连续内存比看起来整齐更关键分块解决的是缓存层的问题下一层问题是数据在内存里到底怎么摆放。绝大多数人默认的二维数组是按行连续存储的也就是A[i][j]的下一个元素是A[i][j1]。这在很多场景没问题但在 GEMM 里如果我们要反复读取 B 矩阵的一列那每一列的元素在内存里是跳着分布的一跳就是一行的大小。跳跃访问对缓存和预取都不友好明明读了一整块缓存行里面可能只有几个字节有用。我偏好的方式是提前做一次数据重排把参与运算的 B 矩阵改成按分块存放。说白了就是把 B 先拆成若干BK × BN的小块每个小块内部再按列变成按行连续存放。这样后面每一次内层循环读取 B 的小块时都是连续内存访问预取器也能识别出规律。重排的开销只是一次性的但对几千上万次小块读取来说收益非常可观。连续访存的好处一次缓存行读取可以被完整利用不会出现读 64 字节只用了 4 字节的浪费。减少 TLB 缺失连续大块访问的页表项数量更少不会因为反复横跳把快表打爆。数据对齐分块大小选 64 或 128 这种对齐友好的值能让向量加载指令直接一次性把数据搬进寄存器。2.3 指令调度别让计算单元闲着等数据分块和数据布局解决了数据在哪的问题指令调度解决的是数据到了以后计算单元能不能立刻用上的问题。现代处理器的乘加指令延迟通常有好几个周期比如一条fma从发射到出结果需要四到六个周期但它可以每周期发射一条甚至更多。这里的核心技巧就是通过手动展开循环把相互独立的乘加计算穿插排列让第一条 fma 还在算的时候第二条、第三条已经发射出去了。我用一个很朴素的生活类比来解释你去食堂打饭如果每打一份饭都站在原地等食堂阿姨把菜盛好再走队伍会排到门口。但如果你先刷卡然后赶紧去打汤、拿筷子等回来时饭刚好盛好你就能一直处于有事情做的状态。CPU 里的乱序执行和寄存器重命名就像打饭流程里的并行窗口我们要做的是让代码里同时存在足够多的打汤、拿筷子这类独立任务。这就是循环展开的意义。实际实现时我常把内层累加循环展开成 4 路或 8 路。每路独立维护一个累加寄存器组最后再合并。这样代码看起来啰嗦但吞吐量能拉高一大截。不过这里有个代价就是寄存器会被累加变量占掉很多。如果你同时有 8 个累加变量每个变量又对应一个B块里的一小段那留给数据加载的寄存器就少了。所以展开路数、分块大小、寄存器数量这三者必须一起调不能只看单项。3. 从能跑到跑得快三轮迭代实现的完整复盘3.1 第一轮先保证结果是正确的DeepGEMM 的第一版我写得很简单基本就是把三重循环翻译成代码甚至没有特别在意循环顺序。这一版的作用只有一个和标准数学库的结果对齐验证我对数据布局、边界填充、累加顺序的理解没有错。这一版的速度当然不行因为内层循环每次做C[i][j] A[i][k] * B[k][j]时B 的列访问在内存中是跳跃的。这个阶段我不看性能只看误差。我会用随机矩阵跑很多轮比较我自己实现的结果和标准库结果把绝对误差控制在浮点累加顺序不同所允许的范围之内。如果你的第一版连正确性都保证不了后面所有优化都白搭因为性能调优最容易出现的问题就是优化了半天算出来已经是错的了。3.2 第二轮引入分块和循环重排第二版做的第一件事是把外层循环顺序从i-j-k改成以分块为单位的bk-bi-bn-kj-ki之类的顺序让 B 和 A 的读取都回到连续空间上。写代码的时候我没直接用分块库而是自己动手做了三个临时缓冲区把 A 的一个大块切成BM × BK的子块把 B 的大块切成BK × BN的子块然后把它们塞到缓存友好的位置上再做内层乘法。用伪代码表示大概是这样for (int bi 0; bi M; bi BM) { for (int bj 0; bj N; bj BN) { // 把 C 的分块清零 for (int bk 0; bk K; bk BK) { // 加载 A 和 B 的分块到临时缓冲 for (int i 0; i BM; i) { for (int k 0; k BK; k) { for (int j 0; j BN; j) { C[bii][bjj] A_tmp[i][k] * B_tmp[k][j]; } } } } } }当然上面只是个逻辑示意真正的内层循环我会再改成指针累加的形式避免每层都做下标计算。这一轮做完以后性能已经比第一版好了一个量级主要原因是缓存命中率显著上升。但离目标还差不少因为内层i-k-j的顺序依然有一层循环在横跨行与行之间并且没有做到指令级别的并行。3.3 第三轮手写乘加核心把并行度抠出来第三版我做了两个比较重要的决定把最内层的k循环展开成 4 路或 8 路每一路同时推进连续不同的k值。让每个输出分块不再只由一个线程负责而是进一步拆成多个子块划分给不同线程或不同核心执行。大分块负责把数据留在缓存小分块负责让多个计算单元同时忙起来。展开以后的核心逻辑看起来会有点吓人但本质就是在同一个循环里同时算好几个独立的输出位置// 以 4x4 输出子块为例的累加片段 for (int kk 0; kk BK; kk) { float a0 A_reg[i0][kk]; float a1 A_reg[i1][kk]; float a2 A_reg[i2][kk]; float a3 A_reg[i3][kk]; c00 a0 * B_reg[kk][j0]; c01 a0 * B_reg[kk][j1]; c02 a0 * B_reg[kk][j2]; c03 a0 * B_reg[kk][j3]; c10 a1 * B_reg[kk][j0]; // ... 类似地累加 c10 到 c33 }这段代码有两个目的第一是减少循环控制指令的开销第二是为 CPU 提供一组互不依赖的乘加指令让它能在等待一条指令结果的周期里继续发射其他指令。编译器有时候能看懂这些循环并帮你展开但很多时候它因为担心指针别名、循环边界等问题不敢做激进的变换。手写展开的好处是稳定可控坏处是代码可读性下降所以我会把展开的核心单独放一个函数里外面配好注释。3.4 一个容易被忽略的正确性验证步骤第三版的性能上来以后我反而更小心了。因为浮点加法不满足结合律展开以后累加顺序变了结果和朴素版本会有细微差别。这是允许的但必须保证误差在一个很小的范围内。我一般会把两种实现的结果差值除以标准库结果的最大绝对值得到一个相对误差。如果相对误差超过1e-4量级说明代码里可能有 bug比如某个分块被重复累加了或者 B 的转置逻辑写错了。我自己在实际调优中遇到过一个典型案例某个分块的边界不是整块大小需要对 A 或 B 的尾部做补零处理。我图省事直接把补零循环跳过了结果矩阵尾部一小部分结果是错的而且因为只有特定尺寸才会触发一开始根本看不出来。后来我加了边界形状的自动测试把 M、N、K 分别设为64、65、127这种非对齐尺寸才把问题暴露出来。这种非对齐测试要比只测对齐尺寸重要得多因为真实深度学习形状里序列长度和隐藏维度经常不是 64 的整数倍。4. 性能上不去时我按这个顺序排查瓶颈4.1 第一个坑矩阵不大但内存带宽先满了我印象很深的一次性能异常是分块和数据布局都做完以后一个 512×512 的 GEMM 跑出来的时间比我预想的高很多。浮点运算量不大按理说计算单元应该很轻松但整个内核的运行时间几乎和访存量成正比这时我意识到这个512×512的场景其实是被内存带宽卡住了而不是被计算单元卡住。判断方法很简单我会算一下理论访存量。输出是512×512个浮点数也就是 1MB 写回A 和 B 各 1MB按理想情况只读一遍的话总共大约 3MB 的访存。如果内存带宽在 20GB/s 左右理论上最少要 0.15ms 左右。如果实测在 0.4ms 以上说明至少有大量重复读取或者访问不连续。遇到这种情况我不会急着加指令展开而是先检查是不是某个小块被反复加载或者 B 的重排没有生效。访存优化永远是先于计算优化的因为计算单元再快喂不进去数据就是白搭。4.2 第二个坑并行分块不均导致部分核心闲着另一个大坑出现在多核并行场景。我有一次把一个1024×1024的矩阵按行切成两半让两个线程各算一半看起来负载非常均匀。但由于机器上不止两个核在跑系统里其他进程也在抢核心实测两个线程的完成时间能差到两倍。更麻烦的是如果只按行切列方向上没有并行那么每个线程内部依然是串行处理长条状的矩阵访存局部性也不好。后来我改成二维分块并行把输出矩阵切成一个网格行方向和列方向都参与划分每个线程处理一个格子。这样不仅负载更均匀而且每个格子都能对应到一段较小且完整的内存区域。有一点很关键如果矩阵尺寸不能被线程数整除不要让某些线程多了一块而其他线程少了一块最好统一用任务队列方式动态领取分块。谁算完谁去拿下一块而不是静态分配这样就算某些核心被系统调度拖慢了其他核心也能自动补上整体耗时更稳定。4.3 排查工具和步骤不是靠猜而是靠对照现象可能原因排查手段时间随矩阵尺寸线性增长但和计算量不成比例内存带宽瓶颈统计实际访存量与理论最小访存量对比多核并行后加速比很低并行划分不均衡或锁竞争用计时器分别统计每个线程的耗时分块变大性能反而下降缓存或共享内存溢出检查分块占用容量是否超过缓存阈值非对齐尺寸性能断崖下跌边界处理分支导致流水线中断对尾部形状单独做补零避免在热路径加判断我还习惯在调优后生成一份汇编或反汇编看一眼内层循环里是否出现了不必要的分支跳转。热路径上任何带if判断的操作都可能让流水线停顿。如果发现边界判断在内部我会把它挪出去要么用补零让边界变得整齐要么对主循环和尾部循环分开处理。5. DeepGEMM 的实测数据与扩展路线5.1 测试环境与相对结果我测试用的机器并没有刻意追求顶级配置目的只是想看 DeepGEMM 在常见硬件上的表现。这是一个双路服务器内存带宽大概在 40GB/s 量级CPU 支持单指令多数据和乘加融合指令。所有测试都用单精度浮点矩阵尺寸覆盖了深度学习里常见的小批次形状。下表是我拿 DeepGEMM 和系统自带的通用数学库做的对比数字是多次运行后的中位时间为避免绝对性能受具体芯片影响我用相对倍数来表示标准化到通用库耗时矩阵形状M×N×K场景标准库耗时DeepGEMM 耗时相对差距128×128×128小规模推理1.00.72DeepGEMM 快约 1.4 倍512×512×512中等规模1.00.81快约 1.2 倍1024×1024×1024偏计算密集1.00.94基本持平3072×768×768FFN 常见形状1.00.89快约 1.1 倍我解释一下这个结果小尺寸矩阵比拼的完全是启动开销和数据搬运策略DeepGEMM 的专用分块能减少很多无效的初始化和检查所以优势明显。到了 1024 这种大尺寸通用库的自动调优已经选了不错的算法我这边能持平就已经很不容易了毕竟对方有更多的专家配置做支撑。FFN 形状属于比较高瘦的矩阵通用库可能为典型正方形做了过多优化反而 DeepGEMM 因为重排策略匹配拿到了一点点优势。5.2 从结果里读出的两层信息第一层信息很简单不存在一个完全通用的最优 GEMM 实现形状不同瓶颈不同策略就得跟着变。第二层信息更重要优化 GEMM 的收益有两个来源一是减少冗余访存二是提高计算并行度而这两者在不同尺寸下权重完全不同。小尺寸时哪怕只少一次边界判断节省的时间占比都很大大尺寸时计算单元能不能满负荷运转才是决定成败的主线。这也就意味着如果你要复刻这样一个项目最值得投入精力的不是背诵某个分块大小而是建立一套自己的测试基线。每次改一个参数就记一次时间和误差做成一个简单的回归测试表。这样你能够看到每一次调整到底带来了收益还是回退而不是凭着感觉大改。5.3 后续可以扩展的方向低精度、稀疏和动态形状DeepGEMM 目前完成的是单精度密集 GEMM 的优化但显然这部分工作还能延伸出去。第一个方向是低精度也就是 FP16 甚至 INT8。深度学习推理阶段用低精度是主流做法可以把数据带宽减半甚至减到四分之一同时利用新一代硬件里专门的低精度矩阵计算单元。这里要小心的是数据溢出和累加精度问题不能简单地把输入截断成低精度否则一个 batch 的误差会被放大。第二个方向是稀疏化。现实中很多模型的权重经过剪枝后存在大量零值如果我们能识别出固定的稀疏模式直接跳过这些零值理论上可以获得好几倍加速。难点在于稀疏模式有结构化比如按块剪枝和非结构化之分后者对硬件不友好性能收益很容易被索引开销吞掉。第三个方向是动态形状。推理场景里序列长度往往不固定这要求 GEMM 内核不能在每次遇到新形状时都重新自动调优。更好的做法是把内核编译成多个候选版本运行时根据当前形状做一次快速选择。这部分我已经在 DeepGEMM 的路标里排上了但目前还没有做到稳定可发布的程度。最后说一点我做这个项目最大的体会矩阵乘法优化不是那种看一两篇理论文章就能掌握的技能它需要在具体机器上反复做实验亲手把一段代码从能跑打磨到跑得快中间那些报错、异常、性能回退才是真正让你理解系统的地方。如果哪天你发现自己的分块调来调去都差不多不要沮丧那恰恰说明你对当前场景已经接近硬件真实上限了。换个形状、换组参数再试新的一轮瓶颈又会冒出来。保持这种对着计数器较真的劲头比任何一个现成库都值钱。写到这里DeepGEMM 的第一阶段总结就差不多了。如果你也想试试建议先从最朴素的版本开始然后依次加数据重排、分块、展开和并行每一步都记录时间和正确性。这个流程本身就是高性能计算里最值得练的基本功。