很多时候一个项目的瓶颈，并不是模型结构写不写得出来，而是底层的矩阵乘法能不能跑满算力。DeepGEMM 这个项目，起因很简单：我手头有个推理场景，模型的层数和宽度都谈不上夸张，可实测延迟就是压不下去。逐个算子拆完以后发现，几乎所有的等待时间都落在同一个东西上——通用矩阵乘法，也就是 GEMM（Generic Matrix Multiply）。Transformer 里的 attention 打分、FFN 的升维降维，再远一点，卷积如果用 im2col 的形式展开，本质上也是一堆大矩阵相乘。所以我一直觉得，想做好深度学习性能优化，第一关就是把 GEMM 这件事弄明白。

这篇内容会围绕我把 DeepGEMM 从零写出来、再一点点盘到能用的过程展开。适合谁看？两种人：一种是刚接触高性能计算，想知道矩阵乘法怎么就变成了瓶颈的读者；另一种是已经在用各种数学库，但遇到自己不理解的性能上限，想拆开看看底层发生了什么的人。我不会只给结论，会把为什么这样设计、为什么这里慢、为什么换一种写法就快了一截的过程讲出来。

1. GEMM 在深度学习里的出场频率：说它不是瓶颈的人，多半没认真数过

1.1 一次前向计算里，GEMM 到底被调用了多少次

很多人一听到 GEMM，第一反应是"这不就是线性代数里的矩阵乘法吗，我早就会了"。会算和会调度是两回事。我先举个例子：假设一个很普通的 Transformer decoder 层，输入维度是 768，注意力头数 12，FFN 中间层维度是 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][j+1]。这在很多场景没问题，但在 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[bi+i][bj+j] += 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 和系统自带的通用数学库做的对比，数字是多次运行后的中位时间，为避免绝对性能受具体芯片影响，我用相对倍数来表示，标准化到通用库耗时：

我解释一下这个结果：小尺寸矩阵比拼的完全是启动开销和数据搬运策略，DeepGEMM 的专用分块能减少很多无效的初始化和检查，所以优势明显。到了 1024 这种大尺寸，通用库的自动调优已经选了不错的算法，我这边能持平就已经很不容易了，毕竟对方有更多的专家配置做支撑。FFN 形状属于"比较高瘦"的矩阵，通用库可能为典型正方形做了过多优化，反而 DeepGEMM 因为重排策略匹配，拿到了一点点优势。

5.2 从结果里读出的两层信息

第一层信息很简单：不存在一个完全通用的最优 GEMM 实现，形状不同，瓶颈不同，策略就得跟着变。第二层信息更重要：优化 GEMM 的收益有两个来源，一是减少冗余访存，二是提高计算并行度，而这两者在不同尺寸下权重完全不同。小尺寸时，哪怕只少一次边界判断，节省的时间占比都很大；大尺寸时，计算单元能不能满负荷运转，才是决定成败的主线。

这也就意味着，如果你要复刻这样一个项目，最值得投入精力的不是背诵某个分块大小，而是建立一套自己的测试基线。每次改一个参数，就记一次时间和误差，做成一个简单的回归测试表。这样你能够看到每一次调整到底带来了收益还是回退，而不是凭着感觉大改。

5.3 后续可以扩展的方向：低精度、稀疏和动态形状

DeepGEMM 目前完成的是单精度密集 GEMM 的优化，但显然这部分工作还能延伸出去。第一个方向是低精度，也就是 FP16 甚至 INT8。深度学习推理阶段用低精度是主流做法，可以把数据带宽减半甚至减到四分之一，同时利用新一代硬件里专门的低精度矩阵计算单元。这里要小心的是数据溢出和累加精度问题，不能简单地把输入截断成低精度，否则一个 batch 的误差会被放大。

第二个方向是稀疏化。现实中很多模型的权重经过剪枝后存在大量零值，如果我们能识别出固定的稀疏模式，直接跳过这些零值，理论上可以获得好几倍加速。难点在于稀疏模式有结构化（比如按块剪枝）和非结构化之分，后者对硬件不友好，性能收益很容易被索引开销吞掉。

第三个方向是动态形状。推理场景里序列长度往往不固定，这要求 GEMM 内核不能在每次遇到新形状时都重新自动调优。更好的做法是把内核编译成多个候选版本，运行时根据当前形状做一次快速选择。这部分我已经在 DeepGEMM 的路标里排上了，但目前还没有做到稳定可发布的程度。

最后说一点我做这个项目最大的体会：矩阵乘法优化不是那种看一两篇理论文章就能掌握的技能，它需要在具体机器上反复做实验，亲手把一段代码从"能跑"打磨到"跑得快"，中间那些报错、异常、性能回退，才是真正让你理解系统的地方。如果哪天你发现自己的分块调来调去都差不多，不要沮丧，那恰恰说明你对当前场景已经接近硬件真实上限了。换个形状、换组参数再试，新的一轮瓶颈又会冒出来。保持这种对着计数器较真的劲头，比任何一个现成库都值钱。

写到这里，DeepGEMM 的第一阶段总结就差不多了。如果你也想试试，建议先从最朴素的版本开始，然后依次加数据重排、分块、展开和并行，每一步都记录时间和正确性。这个流程本身，就是高性能计算里最值得练的基本功。