1. 从一块显卡说起:为什么我要折腾 DeepGEMM
去年年底,我在一台单卡工作站上跑一个中等规模的矩阵计算任务,模型不大,参数量也就几亿级别,但推理延迟始终卡在一个让人难受的数值上。GPU 利用率看着不低,可吞吐就是上不去。我一开始怀疑是数据加载的问题,折腾了几天数据管线,收效甚微。后来用性能分析工具一查,发现大量时间花在了通用矩阵乘法(GEMM)上——这是深度学习里最基础也最频繁的运算,全连接层、注意力机制、卷积展开,底层几乎都是它在撑着。
问题在于,我用的那套通用计算库是为“大而全”设计的,它要兼容各种硬件、各种精度、各种形状,代码里塞满了分支判断和兜底逻辑。对于我这种形状相对固定、精度明确的场景,这些通用性反而成了负担。于是我开始找更轻量、更专注的 GEMM 实现方案,就这样接触到了 DeepGEMM 这个方向。
DeepGEMM 不是一个具体的商业产品,而是一类针对深度学习场景深度优化的高性能矩阵乘法内核实现思路。它的核心目标很明确:在特定硬件上,针对深度学习常见的矩阵形状和精度需求,把 GEMM 的性能压榨到接近硬件理论峰值。它解决的问题是——当你不需要“万能”,只需要“够快够专”的时候,通用库给不了你最优解。
这篇文章适合谁看?如果你是在单卡或小规模集群上做模型训练、推理优化的工程师,如果你被 GEMM 性能卡过脖子,如果你对 CUDA 编程和 GPU 性能优化有基本了解但不知道从哪下手,那这篇内容应该能帮到你。我会从设计思路、核心细节、实操过程到踩坑经验,完整拆一遍。
2. 整体设计思路:为什么是“深度定制”而不是“拿来就用”
2.1 通用库的困境与专用内核的突围
先说清楚一个前提:通用 GEMM 库(比如各类 BLAS 实现)非常优秀,它们经过了几十年的打磨,在绝大多数场景下都能给出不错的性能。但“不错”和“极致”之间,往往隔着几倍的差距。通用库为了覆盖所有可能的矩阵形状、数据类型和硬件配置,必须做大量运行时判断和分支调度。这些逻辑在单次调用中开销不大,但在深度学习训练中,GEMM 会被调用成千上万次,累积起来就是可观的损耗。
更关键的是,深度学习里的 GEMM 有很强的规律性。比如 Transformer 结构里,注意力的 Q、K、V 投影矩阵形状通常是固定的几个值;前馈网络的两层线性变换也有稳定的维度关系。这些形状在训练开始前就确定了,不会中途变化。通用库不知道这一点,它每次都要重新做形状分析和内核选择。而专用内核可以针对这些固定形状,在编译期就把最优的线程块大小、共享内存布局、寄存器分配方案全部定死,运行时零开销。
DeepGEMM 的设计哲学就是“把已知信息用到极致”。既然形状固定,那就为这个形状写一个专属内核;既然精度明确(比如 FP16 输入 FP32 累加),那就砍掉所有不必要的精度转换逻辑;既然硬件型号已知,那就针对这块卡的 SM 数量、共享内存大小、寄存器文件容量做精确调优。这种思路牺牲的是通用性,换来的是性能的显著提升。
2.2 精度选择背后的权衡逻辑
深度学习里 GEMM 的精度选择是个绕不开的话题。常见的有 FP32、FP16、BF16、TF32,还有更激进的 FP8。每种精度背后都是数值稳定性和计算速度的权衡。
FP32 最稳,但吞吐最低。FP16 速度快,但动态范围窄,容易溢出。BF16 动态范围和 FP32 一样,但尾数精度低,适合对精度不那么敏感的场景。TF32 是折中方案,在保持 FP32 动态范围的同时用 19 位表示尾数,速度比 FP32 快不少。
DeepGEMM 的常见做法是:输入用 FP16 或 BF16,累加用 FP32。这样既保证了计算速度,又避免了累加过程中的精度损失。为什么累加必须用 FP32?因为 GEMM 的本质是大量乘加运算,如果累加也用 FP16,误差会随着累加次数增加而迅速放大。假设每次乘加引入的相对误差是 ε,累加 N 次后误差大约放大 √N 倍。对于 N=4096 的矩阵,误差会放大 64 倍,这在很多任务里是不可接受的。
注意:精度选择不是越高越好,也不是越低越好。关键看你的任务对数值误差的容忍度。分类任务通常对精度不敏感,FP16 甚至 FP8 都能跑;但涉及概率输出、归一化层、损失计算的部分,精度不够会导致训练不稳定。
2.3 内存层次结构的利用策略
GPU 的内存层次是:寄存器 → 共享内存 → L2 缓存 → 全局内存。速度依次递减,容量依次递增。GEMM 优化的核心就是让数据尽可能停留在快的层次里,减少对慢速内存的访问。
DeepGEMM 的典型策略是分块(Tiling)。把大矩阵切成小块,每个线程块负责计算一个输出块。输入块从全局内存加载到共享内存,然后每个线程从共享内存读取数据到寄存器进行计算。这样做的目的是让每个数据被多次复用——一个输入块会被多个输出块用到,放在共享内存里就能避免重复从全局内存读取。
分块大小的选择很讲究。块太小,数据复用率低,全局内存访问频繁;块太大,共享内存放不下,或者寄存器压力过大导致占用率下降。通常需要根据硬件的共享内存容量和寄存器文件大小来反推。比如某代 GPU 每个 SM 有 64K 个 32 位寄存器,共享内存可配置为 48KB 或 96KB。如果每个线程块用 256 个线程,每个线程最多能用 256 个寄存器,共享内存最多用 48KB。这些约束直接决定了分块的上限。
3. 核心细节解析:从线程块到寄存器的手工调优
3.1 线程块与线程的映射关系
GEMM 内核的线程组织方式直接影响性能。常见的有两种:一种是一个线程负责一个输出元素,另一种是一个线程负责多个输出元素。前者实现简单但数据复用率低,后者实现复杂但性能更好。
DeepGEMM 通常采用后者。具体来说,一个线程块负责计算一个 128×128 的输出块,块内有 256 个线程,组织成 8×32 的二维结构。每个线程负责计算 16×4 个输出元素。这样每个线程需要读取 16 行输入 A 和 4 列输入 B,做 64 次乘加。通过合理分配寄存器,这 64 次乘加可以流水线化,隐藏内存访问延迟。
为什么是 8×32 而不是 16×16?因为 GPU 的 warp 是 32 个线程,线程块的维度最好和 warp 大小对齐。8×32 意味着每个 warp 覆盖 4 行输出,这样在读取输入 A 时,同一个 warp 内的线程可以共享行数据,减少共享内存的 bank 冲突。bank 冲突是共享内存访问的常见性能杀手,当多个线程同时访问同一个 bank 的不同地址时,访问会被串行化。通过精心设计线程到数据的映射,可以尽量避免这种情况。
3.2 共享内存的布局与 bank 冲突规避
共享内存被分成 32 个 bank,每个 bank 宽度 4 字节。如果两个线程访问同一个 bank 的不同地址,就会发生冲突。对于 FP16 数据,每个元素 2 字节,两个元素才占一个 bank 宽度。这会导致更复杂的冲突模式。
常见的规避方法是在共享内存里给数据加 padding。比如原本一行 128 个 FP16 元素占 256 字节,正好是 64 个 bank 宽度。如果直接存储,同一列的不同行会映射到同一个 bank。加 8 个元素的 padding 后,每行占 272 字节,列方向的 bank 映射就错开了。这个 padding 大小不是随便定的,需要根据分块大小和访问模式计算。一般来说,padding 取 4 或 8 个元素比较常见。
实操心得:共享内存 padding 是 GEMM 优化里最容易见效的手段之一。我试过在一个内核里只加了 8 个元素的 padding,性能直接提升了 15%。但 padding 也不是越大越好,太大会浪费共享内存容量,反而限制分块大小。
3.3 寄存器分配与流水线设计
寄存器是 GPU 上最快的存储,但数量有限。每个线程能用的寄存器数量决定了能同时容纳多少中间结果。在 GEMM 内核里,寄存器主要用来存放输出块的累加值、输入 A 的片段、输入 B 的片段。
以 16×4 的输出块为例,需要 64 个寄存器存放累加值。输入 A 的片段假设是 16×4,需要 64 个寄存器;输入 B 的片段是 4×4,需要 16 个寄存器。加起来已经 144 个了,再加上地址计算、循环变量等,很容易超过 255 的上限。所以实际实现中,输入片段通常不会全部放在寄存器里,而是分批次从共享内存读取。
流水线设计的关键是让计算和内存访问重叠。当线程在计算当前批次的乘加时,下一批次的数据应该已经在从共享内存加载的路上了。这需要软件流水线技术:把循环展开,提前发出加载指令,让 GPU 的乱序执行引擎有机会隐藏延迟。具体展开几层,取决于共享内存的访问延迟和计算指令的吞吐。一般来说,展开 2 到 4 层比较常见。
4. 实操过程:从零搭建一个 DeepGEMM 内核
4.1 环境准备与基础框架
先确认硬件和软件环境。你需要一块支持 CUDA 的 GPU,计算能力 7.0 以上比较合适,因为新架构对 FP16 和 BF16 的支持更好。CUDA Toolkit 建议用 11.0 以上版本,编译器用 nvcc。操作系统不限,Linux 下开发体验更顺滑。
基础框架很简单:一个 host 端函数负责分配内存、拷贝数据、启动内核;一个 device 端内核函数负责实际计算。host 端用 cudaMalloc 分配全局内存,用 cudaMemcpy 传输数据。内核启动时指定网格维度和线程块维度。网格维度根据输出矩阵大小和分块大小计算,比如输出是 1024×1024,分块是 128×128,那网格就是 8×8。
// 内核启动配置示例 dim3 blockDim(256); // 256个线程 dim3 gridDim((N + 127) / 128, (M + 127) / 128); // 按128分块 gemm_kernel<<<gridDim, blockDim>>>(d_A, d_B, d_C, M, N, K);4.2 内核主体结构拆解
内核函数的第一件事是计算线程的全局索引,确定自己负责哪个输出块。然后声明共享内存数组,用于存放输入 A 和 B 的分块。接着进入主循环,沿着 K 维度迭代。每次迭代做三件事:协作加载输入分块到共享内存、同步线程、执行乘加计算。
加载阶段,256 个线程要协作把 128×128 的输入块从全局内存搬到共享内存。每个线程负责搬运 64 个元素。搬运时要注意合并访问——连续线程访问连续地址,这样内存控制器可以合并请求,提高带宽利用率。对于行优先存储的矩阵,让线程沿着行方向连续排列,每个线程负责一行中的连续元素,这样最容易实现合并访问。
计算阶段,每个线程从共享内存读取自己需要的输入片段,执行乘加。这里有个技巧:把 K 维度的循环展开,每次处理多个 K 值。这样可以在寄存器里保留多个累加值,减少对共享内存的访问次数。比如每次处理 8 个 K 值,那每个线程需要从共享内存读取 8 组 A 片段和 8 组 B 片段,但乘加次数是 8 倍,计算访存比提高了。
4.3 参数计算与选择过程
分块大小的选择需要算一笔账。假设目标 GPU 每个 SM 有 64K 个 32 位寄存器,共享内存 48KB。如果线程块用 256 个线程,每个线程最多 255 个寄存器,那总共需要 256×255=65280 个寄存器,刚好卡在上限。但实际不可能用满,因为还要留一些给系统。所以每个线程控制在 128 到 200 个寄存器比较现实。
共享内存方面,128×128 的 FP16 分块,A 和 B 各占 128×128×2=32KB,加起来 64KB,超过了 48KB 的限制。所以要么减小分块,要么用双缓冲。双缓冲的意思是分配两块共享内存,一块用于当前计算,一块用于加载下一批数据。这样虽然共享内存用量翻倍,但可以隐藏加载延迟。如果共享内存不够,就只能减小分块到 64×128 或 64×64。
我实际测试下来,对于计算能力 7.5 的 GPU,128×128 分块配合双缓冲需要 96KB 共享内存,刚好是上限。如果共享内存配置为 48KB,那就只能用 64×128 分块。性能差距大概在 10% 到 20% 之间,取决于矩阵形状。
| 分块大小 | 共享内存需求 | 寄存器压力 | 适用场景 |
|---|---|---|---|
| 64×64 | 16KB | 低 | 小矩阵,共享内存受限 |
| 64×128 | 32KB | 中 | 中等矩阵,平衡选择 |
| 128×128 | 64KB | 高 | 大矩阵,性能优先 |
| 128×256 | 128KB | 很高 | 需要双缓冲和大共享内存 |
4.4 性能测试与对比方法
写完内核后,怎么知道它好不好?最直接的方法是测吞吐,单位是 TFLOPS(每秒万亿次浮点运算)。计算公式是:2×M×N×K / 时间 / 1e12。为什么乘 2?因为每次乘加算两次浮点运算。
测试时要注意几点:先用小矩阵预热,让 GPU 频率稳定;然后跑多次取平均,排除波动;最后和通用库对比,看提升了多少。我一般会测三组形状:方形(如 4096×4096×4096)、瘦长形(如 1024×4096×1024)、扁宽形(如 4096×1024×4096)。不同形状对内核的压力不同,方形最均衡,瘦长形对 K 维度的循环优化要求高,扁宽形对输出块的写回效率要求高。
注意:测性能时一定要用 CUDA 事件计时,不要用 CPU 计时。CPU 计时会把内核启动开销算进去,对于短内核来说误差很大。CUDA 事件直接记录 GPU 时间线,准确得多。
5. 常见问题与排查技巧实录
5.1 性能不达预期的排查路径
内核跑通了但性能差,这是最常见的情况。排查要按顺序来,不要东一榔头西一棒子。
第一步,看占用率。用性能分析工具查每个 SM 上活跃的 warp 数量。占用率低说明寄存器或共享内存用太多,限制了并发。解决办法是减少分块大小或降低每线程寄存器用量。
第二步,看内存访问效率。检查全局内存加载是否合并,共享内存是否有 bank 冲突。全局内存的合并访问很容易判断:如果连续线程访问的地址是连续的,就是合并的。共享内存 bank 冲突需要看分析工具的报告,它会告诉你冲突发生在哪条指令。
第三步,看计算流水线是否饱和。如果计算单元利用率低,说明线程在等内存。这时候要增加流水线深度,或者调整分块大小提高计算访存比。
第四步,看指令混合。如果整数指令太多,会挤占浮点指令的发射槽。地址计算尽量用 64 位整数,但不要滥用,能用 32 位就用 32 位。
5.2 数值精度问题的定位与解决
精度问题比性能问题更难查,因为它不会报错,只会让结果慢慢变差。常见症状是训练损失不下降,或者推理输出异常。
首先确认累加精度。如果累加用 FP16,改成 FP32 试试。如果改了就好,说明是累加溢出。其次检查输入数据的范围。FP16 的最大值是 65504,如果输入数据接近这个值,乘法结果很容易溢出。解决办法是在加载输入时做缩放,把数据映射到安全范围。
还有一种情况是共享内存里的数据被覆盖了。这通常发生在双缓冲的同步没做好,加载线程和计算线程访问了同一块内存。检查 __syncthreads() 的位置,确保加载完成后才切换缓冲区。
实操心得:我遇到过一次精度问题,查了两天才发现是共享内存的 padding 没对齐。padding 导致数据布局错位,部分元素读到了错误的值。这种问题不会导致崩溃,但会让结果出现微小偏差,非常隐蔽。所以 padding 的计算一定要仔细,最好写个单元测试验证。
5.3 不同矩阵形状的适配策略
深度学习里的矩阵形状千变万化,一个内核不可能通吃。常见的做法是准备多个内核,根据形状选择。比如 M 和 N 都大于 128 时用大分块内核,M 或 N 小于 128 时用小分块内核,K 特别大时用 K 维度展开更多的内核。
选择逻辑可以放在 host 端,用一个简单的启发式规则。比如:
| 条件 | 推荐内核 |
|---|---|
| M≥128 且 N≥128 | 128×128 分块,双缓冲 |
| M<128 或 N<128 | 64×64 分块,单缓冲 |
| K>4096 | K 维度展开 8 层 |
| K≤4096 | K 维度展开 4 层 |
这个规则不是绝对的,需要根据实际硬件微调。但有了这个框架,至少不会出现某个形状性能特别差的情况。
5.4 常见错误速查表
| 错误现象 | 可能原因 | 解决方法 |
|---|---|---|
| 内核启动失败 | 共享内存超限 | 减小分块或减少缓冲 |
| 结果全零 | 内核没执行或写回失败 | 检查网格维度,检查写回逻辑 |
| 结果部分错误 | 边界处理不当 | 检查 M/N/K 不是分块整数倍时的处理 |
| 性能波动大 | 占用率不稳定 | 固定线程块大小,避免动态分配 |
| 精度逐渐变差 | 累加溢出 | 改用 FP32 累加,或缩放输入 |
| 共享内存冲突多 | 布局不合理 | 加 padding,调整线程映射 |
6. 工具链与调试手段:我常用的几把刷子
6.1 性能分析工具的选择与使用
NVIDIA 的 Nsight Compute 是首选。它可以给出每个内核的详细报告,包括占用率、内存吞吐、计算吞吐、指令统计、bank 冲突次数等。用起来也简单,命令行下ncu --set full ./your_program就能跑。报告里重点看几个指标:SM 占用率、DRAM 吞吐、共享内存吞吐、浮点指令发射率。如果 DRAM 吞吐接近峰值但计算吞吐很低,说明是内存瓶颈;反过来就是计算瓶颈。
Nsight Systems 适合看整体时间线,能告诉你内核启动间隔、数据传输耗时、CPU 和 GPU 的重叠情况。如果发现内核之间有空隙,说明启动开销大或者依赖没处理好。
还有一个轻量级手段:在代码里插 CUDA 事件,手动测每个阶段的耗时。虽然粗糙,但胜在灵活,可以精确到某一行代码。
6.2 单元测试与正确性验证
性能优化最怕的是改了半天,结果算错了。所以每改一版内核,都要跑正确性测试。测试方法很简单:用 CPU 实现一个朴素的三重循环 GEMM,结果作为基准;GPU 内核的结果和基准对比,误差在允许范围内就算通过。
误差范围怎么定?对于 FP16 输入 FP32 累加,相对误差在 1e-3 量级是正常的。如果超过 1e-2,说明有问题。测试用例要覆盖各种形状:方形、瘦长、扁宽、K 特别大、K 特别小、M/N/K 不是分块整数倍。边界情况最容易出 bug。
我习惯用一个小脚本自动生成随机矩阵,跑完对比后输出最大相对误差。这样每次改完内核,跑一遍脚本就知道有没有改坏。
6.3 版本管理与回退策略
GEMM 优化是个反复试错的过程,经常改了几十版才发现某一版性能最好。所以版本管理很重要。我一般用 git 管理代码,每改一个参数就提交一次,提交信息写清楚改了什么、性能变化多少。这样万一改坏了,可以快速回退到上一个好版本。
另外,性能数据也要记录。我会用一个简单的表格记录每次提交的配置和性能,方便对比。时间长了就能看出哪些参数影响大,哪些影响小,积累出经验。
7. 影响范围与延展思考
DeepGEMM 这类专用内核的影响范围其实比想象中广。最直接的是模型训练和推理的加速,尤其是 Transformer 类模型,GEMM 占了计算量的大头。间接的影响是,它让中小规模硬件也能跑动更大的模型——因为计算效率提高了,同样的硬件能做更多事。
从技术延展的角度看,DeepGEMM 的思路可以迁移到其他算子。比如卷积可以用类似的分块和共享内存优化,注意力机制里的 softmax 也可以做专用内核。核心思想是一样的:找到计算模式里的规律性,针对规律性做定制,把通用性带来的开销砍掉。
还有一个方向是自动调优。手工调参很累,而且换一块硬件就要重新调。现在有一些自动调优框架,可以自动搜索最优的分块大小、展开层数、padding 大小。虽然搜索过程耗时,但搜出来的配置往往比手工调的好。我试过用自动调优工具跑了一晚上,第二天得到的配置比我手工调的性能高了 8%。这个方向值得关注。
最后说个实际体会:GEMM 优化没有银弹,每个硬件、每种形状、每个精度组合都可能需要不同的最优配置。但掌握了基本原理和排查方法后,面对新场景就不会慌。先跑通,再测性能,然后按占用率、内存、计算的顺序排查,一步步逼近最优。这个过程很磨人,但看到性能数字往上跳的时候,还是挺有成就感的。