GEMM通用矩阵乘法是一张大模型算力账单里永远最大的一行。Transformer 从注意力投影到 FFN 升维降维再到输出层 logits凡是能写成矩阵乘的地方几乎都是 GEMM。所以当一个名叫 DeepGEMM 的项目出现时我对它的第一反应不是又一个 GEMM 库而是它到底凭什么敢做得这么小。真正细看之后我的判断变了这是一个把 Hopper/Blackwell 硬件特性吃得很透的 FP8 GEMM 内核库核心 CUDA 逻辑只有三百行左右TMA 相关代码三百行左右kernel 启动层又是三百行左右总量不到九百行却能在多种形状下与厂商官方库掰手腕。这篇文章我想从使用者和拆解者的双重视角把 DeepGEMM 为什么快、快在哪、什么样的场景适合它、以及它有哪些边界一次讲清楚。适合正在做推理引擎、算子优化或者单纯想搞懂现代 GPU GEMM 原理的读者。1. GEMM 是大模型算力账单里最大的一行DeepGEMM 的定位与价值1.1 一个公式撑起几乎整个 TransformerGEMM 的数学形式就是 D A × B C但落到深度学习里它无处不在多头注意力的 Q、K、V 三个投影是 GEMM注意力分数计算是 batched GEMMFFN 的 up-projection 和 down-projection 是 GEMM最终输出层的 logits 计算还是 GEMM。拿一个标准的 Transformer decoder 层拆开看除了 softmax、LayerNorm、激活函数这类 elementwise 或归约操作几乎所有计算量都包含在 GEMM 里。这也是为什么所有深度学习框架的性能测试都绕不开 GEMM矩阵乘法跑得快大模型训练和推理就跑得快。显存带宽再高、缓存再大如果矩阵乘法的核心 loop 没有把 Tensor Core 喂饱理论峰值就永远只是一个数字。对做推理引擎的人来说GEMM 性能几乎就等于整个引擎的上限。1.2 DeepGEMM 不是又一个全家桶框架把 DeepGEMM 和 CUTLASS 放在一起比较是最容易产生误解的地方。CUTLASS 是一套完整的、模板化的 GEMM 框架理论上任何形状、任何精度、任何架构特性都有对应的模板组合但代价是你需要理解它的层级架构、依赖关系、以及庞大的模板展开体系光是把编译时间压下去就已经够喝一壶的了。DeepGEMM 的定位是一把手术刀只做一件事在 Hopper 和 Blackwell 平台上把 FP8 GEMM 跑到足够快并且同时覆盖 dense 和 MoE 两种主流的调用方式。它没有厚重的头文件体系没有冗长的构建系统不依赖 cuBLAS、不依赖 CUTLASS、甚至不依赖任何第三方库。我把它理解为一份可以读完的高性能实现既能直接集成进生产推理引擎也可以当教材逐行研究。1.3 它最适合谁来用第一类是框架开发者比如在做大模型推理引擎、训练框架、或者自定义算子库的人需要在一个特定硬件平台上拿到接近峰值的 GEMM 性能又不想被框架的复杂度绑架。第二类是系统课程或自学 CUDA 的同学考研究芯片指令集、共享内存布局、异步流水线DeepGEMM 是极好的真实案例。第三类是正在做 FP8 量化推理落地的人想抄一份经过验证的缩放因子处理方案。如果你是这三类之一这篇文章应该能帮你省下不少自己摸索的时间。2. FP8 不是“精度减半”这么简单数据格式、缩放因子和误差控制2.1 e4m3 和 e5m2两种 FP8 是两套脾气FP8 不是一个单一格式而是两种格式的统称。e4m3 用 4 位指数、3 位尾数精度相对好但动态范围小最大值只到 448 附近e5m2 用 5 位指数、2 位尾数动态范围大得多但尾数只有 2 位精度明显吃亏。两种格式的分布还都和 FP16 一样靠近零的地方更密远离零的地方更疏。主流的做法是 A激活和 B权重都用 e4m3因为训练和推理对精度更敏感e5m2 通常留给梯度或者动态范围极不稳定的张量。如果两者混用比如 A 用 e4m3、B 用 e5m2精度会下降但有时能换取更大的数值覆盖范围需要根据实际量化后的数据分布决定没有绝对的对错。2.2 缩放因子FP8 量化真正的胜负手8 位能表达的范围太有限直接把激活或权重塞进 FP8一旦某个 channel 的数值超出最大值就会溢出如果某个 channel 整体数值偏小又会浪费尾数精度。于是缩放因子 scale 就成了 FP8 方案的核心问题用一个 scale 把原始张量映射到 FP8 最舒服的表达区间计算完成后再乘回去。最简单的方案是 per-tensor整个矩阵共享一个 scale。实现最容易但一旦张量内部不同 channel 的动态范围差异大量化误差立刻变难看。于是出现了 per-row / per-column每个 token 或每个输出通道一个 scale、per-block把矩阵切块每块一个 scale、以及更细的 per-group在 K 维上按 group size 分组。DeepGEMM 支持的正是 per-block 和 per-group 这类细粒度缩放这是它在精度上比早期 FP8 实现明显更有底气的关键。2.3 细粒度缩放在硬件上到底怎么算这里有一个真正考验内核功底的细节。假设 A 有 per-block 缩放因子 sA[m, kBlock]B 有 per-block 缩放因子 sB[n, kBlock]那么输出 D[m, n] 实际上是对每个 K 块分别做矩阵乘再乘上对应块的 sA × sB 之后累加。这里有一个真正考验内核功底的细节。假设 A 有 per-block 缩放因子 sA[m, kBlock]B 有 per-block 缩放因子 sB[n, kBlock]那么输出 D[m, n] 实际上是对每个 K 块分别做矩阵乘再乘上对应块的 sA × sB 之后累加D[m, n] Σ_k A_int[m, k] × B_int[n, k] × sA[m, kBlock] × sB[n, kBlock]问题在于 sA × sB 这个乘积是随 kBlock 变化的不能等到整个 K 循环结束再统一乘一个标量。如果你想用一条指令处理多个 K 块就必须要么在累加中途把累加器按 K 块切开要么把其中一个 scale 提前折进矩阵里。社区里常见的做法是把 B 的 scale 预先折进 B 矩阵本身。B 是权重离线量化时就能完成缩放kernel 内部只需要处理 A 的 scale省掉一次 scale 表的频繁访问。DeepGEMM 的接口允许你配置 scale 的布局方式这正是它在保持通用性的同时还能把额外开销压到极低的原因。如果你自己写 FP8 GEMM我强烈建议先考虑哪个 scale 可以离线折叠而不是一上来就在 kernel 里做动态 rescale。2.4 精度验证时最容易忽略的事集成 FP8 GEMM 之后很多人只看最终 loss 是否有变化这是不够的。用量化误差的视角看缩放因子本身的乘积 sA × sB 也可能发生浮点溢出尤其是在 group size 很小、scale 值很大时。我的习惯是在调试阶段单独跑一层线性层比较 FP8 kernel 输出与 FP32 参考实现的最大绝对误差和相对误差分布并且刻意把输入数据推向极端观察误差是否集中在某个 K 块。不要只看平均值平均值好看掩盖不了局部溢出的问题。3. Hopper 平台的三条加速支柱TMA、wgmma 与流水线编排3.1 TMA把数据搬运从指令流里摘出去在 Ampere 时代GEMM kernel 里最繁琐的部分是数据搬运。每个线程要自己发 LDG 指令把全局内存数据读到寄存器再通过 STS 写进共享内存循环里全是地址计算和同步。到了 Hopper 架构硬件提供了 Tensor Memory AcceleratorTMA这个专门做批量搬运的引擎你只需要在启动 kernel 前配置好一个 TMA 描述符指明源地址、目标共享内存地址、矩阵形状、步长、以及 swizzle 模式然后在 kernel 里发一条非常轻量的指令硬件就会自动完成全局内存到共享内存的搬运连边界填充和格式转换都一并处理。可以这么理解以前是每个工人自己一件一件搬砖现在是按下一个按钮一整条传送带开始工作。指令数量大幅减少线程不再被地址计算绑死可以提前去准备下一次计算。DeepGEMM 对 TMA 的依赖非常深基本上整个主循环就是TMA 搬运下一块数据 wgmma 计算当前块 mbarrier 等待结果的组合。3.2 wgmma一条指令吃下 128 行的大矩阵乘Hopper 引入的 wgmmawarpgroup MMA指令组彻底改变了 GEMM 的编程模型。在 Ampere 上一个 warp 用 mma.sync 指令做 M16N8K16 的小块矩阵乘操作数要从寄存器片段里手工拼接。wgmma 则是以 warpgroup4 个 warp128 线程为单位直接读取共享内存中的 A 和 B 操作数一次性完成 M128 × N256 × K16FP8 输入这样的大块矩阵乘累加器分布在 warpgroup 的寄存器里。这个变化意味着你不再需要关心每个线程具体持有哪些 fragment重点变成了共享内存里的操作数布局必须符合硬件要求的 swizzle 格式以及如何把一系列 wgmma 指令串成一个大的 K 循环。DeepGEMM 的核心逻辑本质上是围绕如何高效编排这些 wgmma 指令展开的这也是整个项目最值得精读的部分。3.3 mbarrier 与异步流水线的节拍wgmma 是异步指令发出去之后不会立刻得到结果。你怎么知道累加器就绪怎么知道 TMA 搬运完成答案都指向 Hopper 的 mbarrier 机制。mbarrier 相当于一块共享内存里的信号旗TMA 搬运完成后可以触发一次 arrivewgmma 完成时也能触发一次 arrive另一边的线程通过 wait 操作等待这些信号。DeepGEMM 的典型流水线是TMA 把第 i1 块数据搬进共享内存 stagei1% stages同时 wgmma 正在计算第 i 块前一块的结果则通过 mbarrier 确认后进入 epilogue。stage 数量就是这条流水线的深度stage 越多越能藏住搬运延迟但共享内存有限FP8 单位数据占用的字节本来就小共享内存很快就会被塞满。实际使用中 2 到 3 个 stage 是常见选择具体要看你矩阵的 block_K 大小和共享内存预算。有一个容易被新手忽略的点由于 wgmma 的 A、B 操作数都必须从共享内存读共享内存的 swizzle 模式必须和 TMA 描述符里设置的模式严格一致。如果你自己改动了数据布局却没同步修改描述符轻则性能大幅下降重则直接读到错误数据而且这种 bug 非常难查因为它不会立刻崩只会悄悄污染结果。4. 约九百行代码的内核库代码骨架、接口设计与设计取舍4.1 一个 Dense FP8 GEMM 内核的主循环长什么样把 DeepGEMM 的逻辑抽象出来一个典型的内核主循环结构大致是这样的// 伪代码展示 dense FP8 GEMM 的主循环骨架 __global__ void deepgemm_kernel(...) { // 1. 初始化 mbarrier、共享内存指针、TMA 描述符 init_smem_and_mbarriers(...); // 2. 主 K 循环每次处理 block_K 大小的数据 for (int k_step 0; k_step num_k_blocks; k_step) { int stage k_step % stages; // 预取下一阶段数据TMA 搬运 tma_load_a(smem_a[stage], ...); tma_load_b(smem_b[stage], ...); // 等待当前阶段数据到位并进行 wgmma 累加 mbarrier_wait(stage); wgmma_fp8_accumulate(acc_regs, smem_a[stage], smem_b[stage], ...); } // 3. epilogue按缩放因子 rescale 累加器写回全局内存 rescale_accumulator(acc_regs, scale_a, scale_b, ...); tma_store_c(acc_regs, ...); }实际代码当然比这复杂但骨架就是这样。第一部分初始化描述符和 mbarrier是整个 kernel 里最容易出错的地方第二部分用 TMA 和 wgmma 交替让搬运和计算重叠第三部分在 epilogue 里处理缩放、偏置和写出。整个循环几乎没有同步指令所有等待都通过 mbarrier 完成这就解释了为什么它能跑得又快又稳。4.2 描述符思想把配置从 kernel 内部剥离DeepGEMM 的接口设计有一个鲜明的特点它不要求用户传一大堆标量参数而是把矩阵形状、步长、block 大小、group size、缩放布局等打包成描述符结构。kernel 启动时代价极低因为所有昂贵的解析和校验都在 host 侧完成。这种设计带来的好处是 kernel 内部的逻辑变得非常直白拿到描述符解析出必要的字段剩下就是纯粹的循环。对比另一种做法——把几十个参数拆成单独的 kernel 参数传进去——描述符模式在代码阅读和维护上的优势很明显。如果你在写自己的高性能算子我建议也采用类似思路把配置和执行解耦后续调参只需要换描述符而不用改 kernel。4.3 为什么它敢不依赖 CUTLASS以及你什么时候该学它不依赖 CUTLASS 是一个有意的取舍。CUTLASS 提供的是完整的 GEMM 组合空间你几乎总能找到一个模板组合来匹配需求但代价是学习成本、编译时间、以及模板层次带来的调试困难。DeepGEMM 的取舍是只服务 FP8、只服务 Hopper 和 Blackwell用更少的代码换取同级别的性能。对大多数业务场景来说我不建议为了不依赖 CUTLASS而刻意不用 CUTLASS但如果你发现自己只用到两三种 kernel 变体、或者你需要把每一行代码都读懂甚至改掉那 CUTLASS 的模板膨胀就会变成一种负担。DeepGEMM 展示了一个可贵的中间路线在一个足够窄的范围内手写内核的代码量完全可控且性能不输给大框架。这也是为什么我特别推荐把它当作学习材料——它逼着你理解硬件而不是背模板。5. Dense 和 MoE 一套内核通吃掩码索引与形状适配5.1 MoE GEMM 的稀疏性到底稀疏在哪Mixture of ExpertsMoE是大模型提升参数量的常见手段每个 token 只被路由到少数几个专家模块之一。从 GEMM 的视角看专家 FFN 的权重可以组织成一个大的三维权重专家数 × 中间维度 × 输入维度而每个 token 只和其中一行专家做矩阵乘。朴素的做法是把它拆成若干个小 GEMM分别处理每个专家的 token但这样会导致不同专家之间的计算量严重不均衡而且小 GEMM 的并行度差白白浪费 Tensor Core。DeepGEMM 的方式是保留一个大 GEMM 的形状M 是所有 token 总数N 是专家数乘以专家中间维度然后通过掩码让每个 M 块只访问对应的 N 列区间。计算仍然在大矩阵上进行但空块被显式跳过几乎没有浪费。5.2 掩码、缩放与专家数的配合方式要在单个内核里支持这种路由稀疏性最核心的数据结构是一个路由掩码表它记录每个 M 块里有哪些 token 实际被路由到了某个 N 块。kernel 在进入某个 N 块的循环之前先查一下掩码如果该 M 块没有任何 token 落到这个 N 块就直接跳过整个块的搬运和计算。缩放因子的处理也和 dense 模式不太一样。MoE 场景下每个专家都有自己的统计分布scale 通常是按专家预计算的所以 group size 和缩放布局需要能精确对应到专家维度的边界。DeepGEMM 在接口层面把 dense 和 MoE 统一到一个入口kernel 内部通过掩码和缩放描述符切换行为这让上层调用方不需要维护两套代码路径。实测下来在专家数较多、token 路由较均匀时这套掩码方案的利用率非常高如果路由极度不均衡少数专家拥入大量 token则空块跳跃机制会尽力兜底但极端情况下还是建议路由均衡策略做配合。5.3 什么样的形状最适合 MoE 模式如果你的模型专家数 E 很大、每个专家的中间维度适中、token 分配大致均匀那 MoE GEMM 的掩码模式优势最明显。反过来如果专家数量很少、每个专家几乎承载所有 token那它就退化成普通 Dense GEMM掩码逻辑只是多了一层开销。判断标准很简单把路由信息画成一个稀疏度矩阵计算一下非零块占比。如果低于三成掩码模式就非常有价值如果高于七成直接走 Dense 路径或者干脆把 token 按专家分组重排可能更省事。6. 性能的边界在哪怎么解读跑分、什么时候不该用它6.1 公开性能数据的正确打开方式从 DeepGEMM 项目公开的性能数据看它在 Hopper 平台的多种形状下都能和厂商官方库的 FP8 路径打得有来有回部分形状还有明显领先。尤其值得注意的是FP8 GEMM 本质上就是把数据传输量对半砍而 DeepGEMM 又把核内调度做到了足够紧凑所以很多场景下看到它接近或超过官方库并不奇怪。公开材料提到它的峰值吞吐超过每秒千万亿次级这在 FP8 路径上已经是相当漂亮的数字。但读跑分资料要带着两个心眼。第一性能数据永远和形状绑定M、N、K 的比例一变最优 kernel 可能完全不同第二基准测试一般会选择对自己最有利的形状组合。所以我更关注的是它在各种形状下的稳定性而不是某一两个峰值点。整体感受是在 M 适中偏大、N 和 K 取常见模型维度的区域DeepGEMM 的综合表现非常能打。6.2 哪些场景我劝你别硬上先泼几盆冷水。DeepGEMM 的核心范围是 FP8如果你当前模型还在用 FP16 或 BF16它并不能直接替代厂商官方库你需要同时保留 FP16 的路径。其次它是相当底层的内核库不提供你开箱即用的 Python API也不负责把输入张量自动量化为 FP8这些前置工作得自己在框架层解决。再次如果你的矩阵形状特别不规则、块尺寸无法对齐掩码和 TMA 的边界处理会引入额外的填充开销可能反而拖慢速度。还有一个很现实的问题团队维护成本。接入一个不到千行的内核库意味着你要自己吃透它、自己维护它一旦硬件升级到下一代就需要跟进新指令集的适配。相比之下厂商官方库的维护风险要低得多。我的建议是如果你的项目对 GEMM 性能极其敏感、形状相对固定、并且团队里有足够强的 CUDA 工程师选 DeepGEMM 或者类似的自研内核是合理路径否则先用它来理解性能上限在哪里再决定是否值得投入。6.3 我实际集成中的几个经验最后分享几条实操层面踩过或看别人踩过的坑每一条都值得记进笔记第一对齐要求。TMA 描述符要求矩阵地址、leading dimension、共享内存地址满足一定的对齐约束通常是 16 字节或 128 字节对齐。集成之前先用一个小脚本把各种形状跑一遍检查是否触发 alignment 相关的错误而不是等跑真实模型时再排查。第二量化流程要独立可验证。FP8 GEMM 的前置量化步骤算出 scale、执行缩放、四舍五入要在 CPU 上有一个参考实现方便对拍 kernel 输出。常见问题是量化过程中用了向零取整而不是四舍五入导致系统性偏差这种偏差很难靠后期调参纠正。第三stage 数不是越大越好。增大 stage 数确实能提升延迟隐藏能力但会吃掉共享内存导致 SM 上能驻留的 block 数减少。如果你发现增大 stage 后占用率下降、性能反而变差不妨减少 stage 数换更多并发 block性能往往更优。第四开启编译器优化级别时注意保留调试符号。我自己遇到过 debug 版本完全正常、release 版本结果错误的情况最后定位到是共享内存读写冲突被编译器重排暴露了出来。排查这类问题最快的办法是用 CUDA 计算器工具检查共享内存访问冲突而不是靠肉眼盯代码。对我来说DeepGEMM 最有价值的不是那点性能数字而是它提供了一个简单到可以全部读懂又强到可以上生产的范本。如果你正在写自己的算子库我强烈建议找一段完整的时间把它的内核代码、描述符定义、TMA 初始化代码全部读一遍再对照这篇文章里的骨架自己动手改一个形状试一试。读十篇总结不如亲手把一个 block_K 改成两倍然后量一次性能变化来得深刻。