DeepGEMM：Hopper/Blackwell架构下的GEMM内核优化实战

DeepGEMM 这个项目我第一次点进去的时候第一反应是“就这么点代码”一个标榜撑起 Hopper 和 Blackwell 两代架构的 GEMM 内核库核心头文件居然只有几百行。但正是这种克制让它在今天一堆动辄几十个模板参数的 CUDA 库中间显得特别扎眼。简单说DeepGEMM 是一个聚焦 FP16/BF16 GEMM 的开源 GPU 内核参考实现专门针对英伟达 HopperH100 系列和 BlackwellB200 系列架构做优化。它不是 CUTLASS 那样的重型框架没有复杂的模板元编程而是把 Tensor Core 相关的核心技巧——TMA 异步拷贝、Warp Specialization、软件流水线、mbarrier 同步——全部压缩在一个可读性很强的内核里。如果你是做 AI 推理优化、算子开发或者单纯想把 Tensor Core 手里这些新硬件特性真正看懂这个项目比读 CUTLASS 源码要友好得多。我花了几个晚上把它跑通、拆开、改参数、看性能这篇文章把整个过程和思考整理出来。1. 这到底是个什么项目极简内核的设计取舍1.1 项目定位与适用场景先说清楚DeepGEMM 不是一个“部署到生产环境直接调用”的完整算子库。它更像一份写给 GPU 算子开发者看的、关于“如何用最前沿硬件特性写 GEMM”的参考实现。它的核心定位有三个作为高性能 GEMM 最小可行实现。代码量少但性能接近商业库的八成以上能让你看清楚哪些优化是决定性的“大头”。作为 Hopper/Blackwell 新特性的教学样本。TMA、tcgen05、warp specialization 这些概念光看指令手册太抽象看它怎么在真实 GEMM 里组合起来会直观很多。作为自研算子的起跑线。如果你想写一个针对特定 shape 优化过的 GEMM、GEMV 或者融合算子拿它做底子改比从零撸 CUTLASS 省事太多。我在本地一台 Hopper 架构的卡上跑过它的 benchmark128 到 8192 的方形矩阵范围内深水区的性能表现非常稳基本能贴着硬件峰值走一个平滑曲线。更重要的是整个过程没有遇到让人抓狂的模板爆炸问题编译时间短报错也友好。1.2 为什么“极简”反而是名副其实很多库号称“轻量”打开才发现只是把模板藏得更深。DeepGEMM 是真的很克制。作者有意不引入复杂的模板层次整个内核用一个统一的 kernel 函数承载不同的 tile 尺寸和流水线配置通过少量宏或者编译期常量控制行为。对比 CUTLASS 动辄typename Gemm...::CollectiveOps的嵌套类型DeepGEMM 更像是手写 CUDA 的“进阶版”——所有关键的优化手段都用原汁原味的 PTX/CUDA 方式表达没有任何魔法。这种设计带来的直接好处是单步调试变得可行。你可以在 Nsight Compute 里单指令单指令地看 TMA 什么时候发出去的、warp 什么时候在等待、mbarrier 的状态怎么翻转的。换成 CUTLASS这个追踪成本高得多得。不过极简也有代价。它不像 CUTLASS 那样支持任意 tile shape、任意数据类型组合、任意布局。DeepGEMM 专注在 FP16/BF16、常规 row-major / col-major 布局上走的是“少而精”的路线。这种取舍我认为是对的与其做一个泛而不精的框架不如把一条主路径磨到极致让后来者看得懂、学得会、改得动。2. Hopper 与 Blackwell 的硬件弹药库内核依赖的三板斧2.1 TMA 与异步内存拷贝Hopper 架构给 GEMM 带来的最革命性变化是 TMATensor Memory Accelerator。这个硬件单元可以把全局内存中的 tile 数据直接搬到共享内存而且完全不占用线程的寄存器、不打断计算流水线。传统 CUDA 内核里你通常要写__ldg或者cp.async指令把数据从全局内存拉到共享内存。cp.async已经是异步的但粒度比较小需要线程自己管理地址、自己发完所有请求再等待。TMA 则是整块 tile 一次性搬运它还自动处理 swizzle 和 padding硬件替你算好共享内存里的存储布局。对 GEMM 这种要反复搬运 A/B 矩阵 tile 的计算TMA 等于把内存搬运从“软件要管理的事”变成了“硬件顺手完成的事”。DeepGEMM 里producer warp 只需要调用一个 TMA 拷贝指令剩下的等待和同步交给 mbarrier。我从 Nsight Compute 的 trace 里看到TMA 拷贝的传输效率和延迟表现非常稳定基本不会被内存地址散乱拖垮。2.2 MMA、寄存器与 Tensor MemoryGEMM 的核心计算单元是 Tensor Core对应 PTX 指令mma.sync。在 Hopper 上最常用的指令是wgmmawarpgroup MMA它允许一个 warpgroup4 个 warp发起更大的矩阵乘操作A/B 矩阵直接从共享内存取输出累加到寄存器。这一步把传统mma.sync的“每个 warp 单独干活”升级成“一个 warpgroup 集体干活”调度效率提高寄存器压力也更容易控制。到了 Blackwell指令变成了tcgen05.mma玩法又有升级。Blackwell 引入了一个叫 Tensor MemoryTMEM的片上存储MMA 的结果可以累加到 TMEM 而不是寄存器。这意味着输出 tile 不再占用寄存器池寄存器可以更多留给操作数和中间值更大规模的 MMA 指令可以一次性完成减少指令发射次数TMEM 配合后续的tcgen05.ld可以高效地把结果从 TMEM 搬回寄存器或共享内存。DeepGEMM 对 Blackwell 的支持重点就是吃透tcgen05.mma和 TMEM 的使用方式。从代码里能明显看到为 Blackwell 写内核需要考虑的东西和 Hopper 很不一样它不再是 Hopper 内核的小修补而是真真切切的另一个架构。2.3 Warp Specialization 与 mbarrier 流水线Warp Specialization即 warp 专门化是 DeepGEMM 性能的另一根支柱。传统的 GEMM 内核里所有 warp 一视同仁每个 warp 都自己搬运数据、自己算乘法中间插一个__syncthreads()来保证数据就绪。这种模式下总有一些 warp 在搬运的时候其他 warp 在等算得快的时候搬运又会成为瓶颈。同步屏障越多流水线越容易气泡。Warp Specialization 的思路是“分工”。把 warp 分成两组Producer warp只负责用 TMA 把 A/B 矩阵的 tile 搬到共享内存Consumer warp只负责用 Tensor Core 做矩阵乘从共享内存取数。两组 warp 之间不是用__syncthreads()这种“全队停下等”的方式同步而是用mbarrier。Producer 完成一块数据的搬运后在屏障上 arriveConsumer 在屏障上 try_wait等到对应 stage 的数据 ready 之后立刻开始计算。这种生产者-消费者模型天然支持多 stage 流水线上一块数据还在算下一块数据已经在搬了硬件利用率自然高。DeepGEMM 里流水线的 stage 数可以配置。一般来说stage2 是起步stage3 和 stage4 在内存延迟较高的时候表现更好。但 stage 不是越多越好因为共享内存总量有限每多一个 stage就要给 A/B 数据各多预留一份缓冲区。这个 trade-off 需要在具体硬件和 shape 上实测。我实际跑下来stage3 在很多中等尺寸矩阵上是一个不错的甜点既能把延迟藏住也不会因为共享内存耗尽导致 occupancy 下降。3. 编译、运行与性能剖析手把手跑一遍3.1 环境要求与编译验证DeepGEMM 对环境的要求比较明确你需要一块 Hopper 或 Blackwell 架构的卡以及配套的 CUDA 12.8 的 Toolkit。原因是它用了较新的 PTX 指令和 mbarrier 语义老版本 CUDA 编译不过。拉代码、编译、跑 benchmark 的基本流程如下git clone 项目地址 cd DeepGEMM # 准备一个简单的 benchmark 入口项目内自带 nvcc -archsm_90a -O3 --use_fast_math -o deepgemm_bench deepgemm_bench.cu ./deepgemm_bench这里-archsm_90a对应 Hopper如果手头是 Blackwell 就要换成sm_100a。注意这个a后缀很重要它代表启用架构相关的特定功能比如 Hopper 的 wgmma、Blackwell 的 tcgen05不加的话很多核心指令编不出来。我第一次编译时就因为顺手写成了sm_90结果一堆 PTX 报错。换成sm_90a之后一马平川。跑 benchmark 的预期结果是打印不同矩阵尺寸下的 TFLOPS。我那次跑 4096x4096 的 BF16 GEMM大概在接近理论峰值的水平上下浮动。如果你看到明显低一截别急着怀疑库先检查自己是不是没加--use_fast_math或者是不是因为其他进程占了 GPU导致实际频率没跑满。3.2 关键参数理解与调优方向每个 GEMM 内核库都有几个决定性能的旋钮DeepGEMM 也不例外。与其盲调不如理解每个参数背后的理由。第一个是 tile 尺寸M_TILE * N_TILE。它决定每个 block 一次处理多大的输出矩阵块。tile 越大数据复用率越高共享内存和寄存器的吞吐越划算。但 tile 太大也会让一个 block 的线程数膨胀导致调度不够灵活小矩阵时浪费严重。DeepGEMM 里常见的配置是128x128或128x256。对于 H100128x128搭配 256 线程是个很稳的起点如果你要优化超大矩阵128x256往往能进一步压榨 Tensor Core 的吞吐。第二个是 warpgroup 的数量。Hopper 上每个 block 可以有多个 warpgroup每个 warpgroup 负责一部分输出 tile。通常你会让一个 block 用 1 或 2 个 warpgroup 做主计算再用 1 个 warp 做 producer。组数太多会导致共享内存的端口争抢加剧反而性能下降。第三个是 swizzle。共享内存里的数据布局不是简单的行列直排而是按一定 pattern 做交织目的是让 TMA 搬运和 Tensor Core 读取时避免 bank conflict。DeepGEMM 里默认的 swizzle 模式我已经验证过基本是当前硬件上的最优解之一不需要你额外发明新 pattern。真要做修改建议先在 Nsight Compute 里确认 bank conflict 的占比再决定要不要动。第四个是主循环的 K 维度循环展开和 preload。DeepGEMM 会在进入主循环前先把第一个 stage 的数据用 TMA 发出去让消费者 warp 一开始就有数据可算。这种”先灌满一枪再开枪“的预热策略对吞吐量的提升非常明显。3.3 用 Nsight Compute 验证瓶颈跑分只是表面成绩要知道自己调得对不对必须看 profile。我的习惯是用 Nsight Compute 的 full section 跑一次重点看几个指标Tensor Core 利用率最直接的指标高则说明计算密度好Memory Workload Analysis看 TMA 的吞吐、共享内存的 bank conflict 占比Warp State Statistics看 warp 停在什么指令上。如果大量时间停在mbarrier.try_wait说明流水线没填满stage 太少如果停在wgmma的等待上说明数据搬运跟不上stage 太多或者 TMA 调度有问题。我跑过一次 profiling 之后明显发现在某个中等尺寸下TMA 的排队延迟偏高。后来把 stage 从 2 调整到 3且给 producer warp 加了一个优先级提示实测性能提高了大约 7%。这种调整不做 profile 根本看不出来。精度验证也别忘了。GEMM 的精度问题主要来自 BF16 的尾数不足和--use_fast_math的影响。DeepGEMM 的 benchmark 通常会自己算一个 reference 结果并对比误差建议你把阈值设严一点毕竟内核优化再怎么花哨算错数就没有意义。4. 踩坑实录与排查路径4.1 编译与运行期常见报错我整理了自己和周围同事经常遇到的几个问题做成一张速查表。现象可能原因解决办法编译报错ptxas fatal: Unsupported .archsm_90a写成sm_90或没加a后缀确认架构代号带aHopper 用sm_90aBlackwell 用sm_100a运行时报no kernel image is available编译架构和实际 GPU 不匹配物理确认卡型号重新用对应-arch编译结果 NaN 或误差爆表开启了--use_fast_math且累加顺序不稳定或数据初始化越界关闭 fast math 验证一遍检查 A/B 的 stride 和 tile 边界性能比理论峰值低 30% 以上没有开多 stage或 producer/consumer 比例不对将流水线 stage 调高给 producer warp 单独做优先级管理还有一个比较隐蔽的坑shared memory 的dynamic_shared_memory需要显式设置大小。如果你把 DeepGEMM 移植到自己的工程里忘记在 launch 配置中传sharedMemSize内核会直接崩或者神秘地结果错误。这类错误不会在编译期报出来只能靠运行时检查cudaGetLastError和 Nsight Compute 的 Shared Memory 配置去定位。4.2 性能不达标的排查清单跑分没到预期的时候我一般按照下面的顺序排查。先检查是不是卡在内存搬运而不是计算。如果 TMA 的利用率已经很高但 Tensor Core 的空隙很大说明消费者在等数据。这时候优先加 stage或者把 A/B tile 的 swizzle 改成 128 字节对齐的 pattern减少 TMA 的传输粒度浪费。再检查是不是 warpgroup 的分配比例不对。一个 block 里如果只有 1 个 warp 做 producer但它要喂饱 4 个 consumer warp在 K 维度很短时很容易喂不饱。反过来producer 太多会导致共享内存端口争抢。我倾向于用 1 个 producer warp 最多 2 个 warpgroup consumer 的配置超过这个比例收益就开始递减。第三个要检查的是 L2 缓存行为。GEMM 的 A 矩阵和 B 矩阵的访问模式不同A 沿着 K 方向走B 沿着 K 方向也走但两者在 L2 中的驻留情况不一样。你可以用cudaCtxResetPersistingL2Cache或者修改访问顺序让核心循环里反复用到的 tile 更容易命中 L2。DeepGEMM 默认的主循环顺序已经比较友好但如果你改了 tile 尺寸这块就要回头重新审视。最后千万别忽略 CPU 端的 launch 开销。如果矩阵尺寸很小kernel 本身的执行时间可能只有几十微秒此时 CPU 侧的 launch 延迟反而会吃掉一半性能。对这种情况用 CUDA Graph 把多次 kernel 启动合并起来效果立竿见影。4.3 几个容易被忽略的边界条件GEMM 的边界处理是很多优化库翻车的地方。DeepGEMM 为了保证内核极简在边界条件上并不像完整库那样对所有 shape 都做精细处理。我补充几个实际会踩到的点。第一是 M、N、K 不是 16 的倍数时tile 可能越界。DeepGEMM 的优化路径主要针对对齐良好的尺寸如果你要跑 1000x1000 这种非对齐尺寸需要在 TMA 描述里设置合适的边界值或者用一个 fallback kernel 处理边界 tile。第二是 stride 不等于 leading dimension 的情况。实际业务里矩阵往往是某个更大 tensor 的一个切片stride 比逻辑维度大。TMA 对这种 strided 布局是支持的但描述符必须正确设置 stride 字段。漏了这个字段数据就会错位结果完全不对。第三是共享内存总量和 occupancy 的权衡。stage 加多了共享内存占用上涨每个 SM 同时驻留的 block 数减少。如果 SM 上的 block 数太少先TMA 和 MMA 之间的重叠能力反而被打折扣。我跑 512x512 的小矩阵时stage4 因为 occupancy 被压低性能反而比 stage2 差。这提醒我们不要无脑堆流水线深度得看 SM 实际能塞下多少 block。最后还有一个容易忽略的点mbarrier的初始化。mbarrier 不是普通变量它需要调用mbarrier.init.shared.b64进行初始化而且线程间对同一个 mbarrier 的 arrive 和 wait 数量必须严格匹配。如果你在循环里混用了多个 mbarrier稍微计数错误就会导致 producer 和 consumer 卡死。DeepGEMM 的源码里对 mbarrier 的 stage 管理写得相当清晰但你自己改代码时改动一个循环变量都可能引发这种隐蔽的死锁。我之前犯过一个很经典的错把主循环的for改成#pragma unroll之后编译器对 mbarrier 的操作顺序做了重排结果消费者 warp 在 stage 不对齐时提前 wait整个 kernel 直接 hang。后来加了__syncwarp()和正确的 memory fence 才解决。相信我这种问题 profile 都查不出来只能一行行对着 PTX 看指令顺序。结尾我个人在实际操作中的体会是DeepGEMM 的价值不只是一份跑分好看的代码更是一个把 Hopper 和 Blackwell 的硬件特性“如何被真正使用”讲清楚的范本。官方文档里关于 TMA、warp specialization、TMEM 的描述再详细也不如你亲手把内核跑起来、用 profiler 看一遍调度过程来得深刻。如果你准备深入研究它我给一个非常实际的建议先别急着改代码先把默认配置在 Nsight Compute 里完整跑一遍记下每个指标的数值然后只改一个参数比如 stage 数量再看指标怎么变。这样一轮下来你对 GEMM 优化的理解会比读十篇论文都管用。后续你完全可以在它的基础上加入 Split-K、Stream-K 甚至多精度融合把它扩展成真正适合自己场景的算子库。毕竟榨干硬件的路从来不止一条。