很多搞深度学习的朋友都有过这种体验GPU 利用率监控看着没毛病可训练吞吐就是上不去排查链路转了一圈才发现时间几乎全耗在同一类运算上——矩阵乘法。我之前做模型性能剖析时统计过从 embedding 到输出层GEMM 类算子能占到整体耗时的七成以上。这也是我后来动手写 DeepGEMM 这套 FP8 矩阵乘法内核的根本原因与其在框架层反复做图优化不如直接吃透最底层这块硬骨头。DeepGEMM 的目标很单纯——利用 Hopper 架构 GPU 上 Tensor Core 与 TMA 的能力实现一个针对 FP8 输入、FP32 累加的矩阵乘法内核。它不是一个要替代工业级内核库的通用方案更像是我对自己的一次摸底GEMM 到底怎么设计才能在特定 shape、特定精度下把硬件算力吃干净。这篇文章会从为什么自研、底层原理、实现路径、性能调优到一次完整的排错过程都摊开来讲适合正在优化自有算子、或者单纯想搞清楚 Tensor Core 内部逻辑的工程师和研究者。1. 为什么自己写 GEMM模型跑得慢不一定怪 GPU1.1 先算一笔账一次 Transformer 前向里的 GEMM 占比随便拿出一个中等规模的 Transformer 模型前向计算里最重头的部分就是四个映射矩阵乘法Q、K、V 的投影是一次大 GEMMattention score 的计算是一次 batched GEMMscore 对 V 的加权又是一次 batched GEMMFFN 两层则贡献了剩余两次超大矩阵乘。再加上 softmax、LayerNorm 这些逐元素算子和归一化操作整个 forward 的 GEMM 占比很容易超过 70%。训练场景还多一个反向传播矩阵乘的数量直接翻倍这个比例只会更高。所以当你在性能剖析工具里看到某个 GEMM 算子耗时异常别急着怀疑数据加载或者并行通信先检查算子本身的实现效率。这也是自研 GEMM 内核最直接的收益点你可以精确控制每一字节的访问路径和每条硬件指令的调度而不是被迫接受别人预设的策略。我见过好几次换库版本后推理提速的案例本质就是新版本对某些 shape 选择了更合适的 tile 切分而不是什么玄学优化。1.2 cuBLAS 和 CUTLASS 都很好为什么还要有 DeepGEMM官方库封装得越彻底对调用者来说黑盒程度就越高。cuBLAS 确实在绝大多数 shape 上表现优秀但它不会把内核实现的细节暴露给你一旦遇到非标准 shape 或特定精度需求你只能按它提供的预估接口去猜性能行为。CUTLASS 项目给了高度灵活的模板组合理论上可以拼出各种内核策略但它概念繁多、模板深度很重想在里面改一条 TMA 加载路径或调整 swizzle 策略学习成本相当高。DeepGEMM 的定位就是轻量定制内核。它不需要覆盖所有硬件平台和所有矩阵 shape只需要在自己的目标机器上做到足够快。这种取舍带来两个直接好处一是代码量小、可读性强每一条关键路径都能很快定位二是遇到性能问题时可以拿官方库做基准对照从 profiling 数据中反推差距在哪个环节真正把性能优化的能力握在自己手里。如果你的业务只是定期跑几类固定矩阵规模自研内核完全有理由比通用库更贴合你的计算模式。1.3 DeepGEMM 的目标边界FP8、Hopper 与工程妥协做任何内核项目都要先划清边界否则会因为什么都想支持而陷入无尽复杂度。DeepGEMM 的初始范围如下精度A、B 矩阵输入为 FP8E4M3/E5M2 可配累加器用 FP32硬件仅针对 Hopper 架构 GPU 做优化重点使用其 TMA、wgmma 和分布式共享内存能力Shape面向超大矩阵M、N 在 2048 以上做满波优化小矩阵则退回到简单路径不做过深调优扩展性后端预留了 BF16 路径主要用于精度对照实验不作为默认主线。这个边界意味着我不需要处理 Ampere 等旧架构的兼容性可以把精力全部投入到大矩阵、批量加载、深流水线这一条主路上。工程上做减法和划边界是同样的道理方案越单纯越容易分析出性能瓶颈范围一旦蔓延调试时会发现每个符号都可能是错的。2. GEMM 内核的底层逻辑从最朴素循环到 Tensor Core 的关键跃迁2.1 最简单的 GEMM 内核长什么样写一个在 CPU 上可以跑通的朴素版 GEMM只需要三重循环对每个输出元素 C[i][j]累加 A 的第 i 行与 B 的第 j 列的内积。拿到 GPU 上如果也照这个逻辑写结果就是灾难每个线程反复从全局内存读取同一段数据复用率趋近于零带宽瞬间成为瓶颈。一个 GPU 上有几千个并发线程如果它们都在串行读同一批 A 元素缓存会被完全冲垮。// 最朴素的 GEMM 思路每个线程算一个输出元素 // 不考虑任何复用、布局、向量化仅作为反面教材 __global__ void naive_gemm_kernel(const float* A, const float* B, float* C, int M, int N, int K) { int row blockIdx.y * blockDim.y threadIdx.y; int col blockIdx.x * blockDim.x threadIdx.x; if (row M || col N) return; float sum 0.0f; for (int k 0; k K; k) { sum A[row * K k] * B[k * N col]; } C[row * N col] sum; }这段代码每个线程要读 K 个 A 元素和 K 个 B 元素如果一个 warp 的 32 个线程各自负责相邻 row 但相同 colB 的读取在缓存层面还能有一定复用A 的读取则完全重复。数据量稍微一涨DRAM 吞吐就直接封顶。真正高性能的 GEMM 内核必须解决两个核心问题数据复用结构怎么设计计算指令怎么选择。2.2 计算强度FLOPs 再高也怕带宽卡脖子判断一个计算任务是否访存密集还是计算密集关键看计算强度——单位字节的数据能支撑多少次浮点运算。一个 MNKn 的方阵乘法总计算量是 2n^3 FLOPs输入数据量是 3n^2A、B、C 各一次读写。把两者相除计算强度等于 2n/3。n1024 时约 680 FLOPs/Byte远高于现代 GPU 的算访比比如某 Hopper GPU 的 FP8 算力若按 1900 TFLOPS、HBM 带宽按 3.35 TB/s算访比约 570意味着理论上是计算密集型。可现实中为什么很多 GEMM 跑不出峰值因为没有有效的数据复用结构实际从 DRAM 搬运的字节数是理论值的几十上百倍。所以优化的本质就是通过分块让数据尽可能留在片上。一个数据元素被读到寄存器后如果能参与多次乘累加DRAM 的压力就成倍下降。考虑 tiling 之后计算强度公式会变成一个关于 tile 尺寸的函数块越大K 维循环中每个加载元素被使用的次数越多。2.3 三级存储的 Tiling 协作寄存器、SMEM、GMEMGPU 芯片上的存储层级可以类比成厨房的操作台全局内存相当于仓库共享内存相当于灶台寄存器相当于手边已经洗净备用的配料。GEMM 内核算法的核心就是设计一条从仓库批量取货、放到灶台上、再在手里快速加工的流水线。传统 CUDA 手动调度路径里一个 CTA 通常负责一个输出 tile先在共享内存铺好 A 和 B 的块再由各线程从 SMEM 读取数据做乘累加。Tile 尺寸的选择直接决定复用率和硬件资源占用。常见策略是每个 CTA 负责 128x128 或 128x256 的输出 tileK 维每次处理 16 或 32。计算一下 SMEM 占用就知道A 和 B 两个 tile 若都是 128x32且 A、B 都是 1 字节的 FP8需要的共享内存约 8KB如果上双缓冲翻倍到 16KB。这个数字相对于一块 Hopper GPU 上 228KB 的 SMEM 来说非常充裕于是还可以进一步加宽 N 方向的 tile用更多的片上数据换取更低的 DRAM 访问频率。A tile 128 x 32 -- FP8: 128*32 4096 字节 B tile 32 x 256 -- FP8: 32*256 8192 字节 双缓冲总占用: 2 * (4096 8192) 24576 字节 Hopper SMEM 典型容量约 228KB留白充足这个算术题看起来很朴素但它决定了内核能否在寄存器、SMEM、GMEM 之间保持足够深的流水线。数据一旦能在片上多待一轮全局内存的带宽压力就小一轮。3. DeepGEMM 的实现路径FP8 格式、Tensor Core 与 warp 分工3.1 FP8 的 E4M3 与 E5M2两个尾数位的取舍FP8 不是一个单一的浮点格式它有 E4M3 和 E5M2 两种子格式。E4M3 用 4 位指数加 3 位尾数取值范围相对窄但精度更高E5M2 用 5 位指数加 2 位尾数动态范围大但尾数噪声更明显。做矩阵乘法时通常把激活值用 E4M3 存储权重则根据是否需要大动态范围在 E4M3 和 E5M2 之间选择。格式符号位指数位尾数位最大值范围适用场景E4M3143约 448激活值、权重的主选E5M2152约 57344梯度、中间动态范围大的数据FP8 本身不能直接用于精度要求高的累加。实际计算必须以 FP32 作为累加器同时给输入乘一个缩放因子scale。DeepGEMM 在加载数据后、送入 Tensor Core 之前会先检查数据的整体范围将 FP8 值反量化到实际浮点范围。这个过程不是每元素做一次除法而是提前算好 scale最后在输出端乘回来。这个设计在代码层看起来只是多了两个乘法但它决定了整个量化误差的分布形态。3.2 Tensor Core 指令形态从 mma.sync 到 wgmma要真正把 FP8 算力烧起来光靠普通的 FMA 指令是不够的。Tensor Core 本质上是一组专用的矩阵乘累加硬件单元需要通过特定指令驱动。在 Ampere 时代最直接的编程方式是 mma.sync 指令族比如 m16n8k16 这样的形状表示一次让一个 warp 协同完成 16x8x16 的矩阵乘片段。Hopper 时代则提供了更强的 wgmma 指令它不再以一个 warp 为单位而是以一个 warp group 为单位一次可以处理更大的矩阵块并且支持异步流水线。DeepGEMM 的核心计算路径选择 wgmma 作为主用路径。使用 wgmma 的好处显而易见一次指令能覆盖更大的输出范围硬件调度也允许与数据搬运输送完全重叠。但也带来了一个副作用——warp 之间的同步和描述符管理变得复杂。一条 wgmma 指令需要引用 A、B 矩阵在共享内存中的地址布局描述符这些描述符由 Tensor Core 的 TMA 硬件读取若描述符内部的 stride 或 swizzle 参数写错结果会非常离谱。3.3 生产者-消费者分工TMA 搬运与计算解耦Hopper 最值得利用的新特性之一就是 TMATensor Memory Accelerator。它允许你用一条指令把全局内存中一个 2D 或 3D tile 直接搬到共享内存还可以在搬运过程中完成 swizzle 布局变换。这意味着传统内核里线程算地址、加载、再同步的模式被彻底改变了数据搬运不再消耗计算线程的指令周期。DeepGEMM 把 warp 分成生产者和消费者两类。生产者的职责只有两个准备 TMA 描述符、发布搬运请求消费者的职责是等待数据到达共享内存然后启动 wgmma 计算。这个模型很像工厂里的采购员和装配工采购员只管把原料运进仓库装配工不管原料怎么来的只等仓库里有货就开工。// 概念性伪代码数据流水线主循环 // 仅用于说明生产者-消费者结构不表示完整内核实现 for (unsigned k_step 0; k_step K / K_BLOCK; k_step) { if (k_step 1 K / K_BLOCK) { // 生产者角色预取下一块 A、B 到另一个共享内存缓冲 tma_load_to_smem(A_smem[1 - cur], A_global_desc); tma_load_to_smem(B_smem[1 - cur], B_global_desc); } // 消费者角色等待当前缓冲数据就绪后启动 wgmma wgmma_async(acc_smem, A_smem[cur], B_smem[cur], ...); }这个示例里的双缓冲结构把搬下一块和算当前块完全重叠只要 TMA 的搬运速度大于 wgmma 的消费速度流水线就能一直保持满载。实际工程里缓冲深度还可以加到三份甚至四份换取更大的容忍波动能力。3.4 分布式共享内存跨 CTA 协作的进阶玩法Hopper 还有一个容易被忽视的能力集群内的多个 CTA 可以组成一个线程块集群通过分布式共享内存机制互相读取对方在共享内存中的数据。这意味着一个 CTA 计算输出 tile 时不需要把完整的 A 或 B tile 都装进自己那一份 SMEM 里而是可以让相邻 CTA 把自己的 A tile 分享出来。DeepGEMM 在部分 shape 上启用了这个功能。比如一个大矩阵的 N 方向被切成很多窄条相邻 CTA 所需的 B tile 完全一样此时让其中一个 CTA 加载并广播给其他 CTA能省掉重复的 DRAM 读取。这个特性让 tile 切分策略有了新的维度不再单纯为了减少全局访问而把 SMEM 塞满还可以通过集群通信实现数据共享。当然分布式共享内存也带来同步复杂度。跨 CTA 访问数据必须用集群级的 barrier 保证顺序稍有不慎就会出现读到了半成品数据的问题。我的建议是第一版内核先别碰这个特性等单 CTA 路径跑稳、性能基线确定后再作为进阶优化加上去否则出现问题时会很难分清是 TMA 描述符写错还是集群同步丢步。4. 性能调优在 Nsight Compute 面前逐项逼出峰值4.1 第一轮先跑通基线版本记录真实的短板任何优化都离不开可量化的起点。DeepGEMM 第一版是直接照搬常见 CUTLASS 思路写出来的基础双缓冲版本没有做 swizzle也没有分布式共享内存。在这种状态下绝大多数矩阵 shape 是能跑对结果的但性能大概只有硬件峰值的 30% 左右。此时最该做的事不是继续加特性而是用分析器看清瓶颈到底在哪个硬件单元。Nsight Compute 里有几个关键指标值得重点看sm__pipe_tensor_cycles_active反映张量核心忙碌程度l1tex__data_bank_conflicts反映共享内存 bank 冲突严重性dram__throughput反映显存压力。第一版实测数据中tensor 核心利用率不高同时 bank conflict 告警非常密集。这说明当前内核的问题在数据供给路径而不是计算单元本身。这个判断把优化方向从改计算指令引导到了改数据布局。4.2 第二轮SMEM 地址布局与 bank conflict 的博弈共享内存按 bank 组织同一地址周期只能让每个 bank 输出一份数据。如果多个线程同时访问落在同一个 bank 但不同地址的位置就会发生所谓的 bank conflict一次访存被拆成多次。GEMM 这种高并发读 SMEM 的负载对 bank 布局极度敏感。最简单的布局下A tile 按行连续排列一行 128 字节。如果一行正好跨越了 32 个 bank 的两个周期而 warp 里的线程又在读同一列的数据冲突就会出现。解决办法是给地址加上一层 swizzle让每个线程读到的逻辑位置映射到不同的物理 bank。以常见 128 字节行宽为例可以按 8 行一组做异或扰动// 常见 swizzle 策略按行、列做 XOR 扰动 // 实际使用哪种模式取决于 tile 形状和 SMEM 宽度 constexpr int SMEM_STRIDE 128; // 字节 int swizzled_offset (row * SMEM_STRIDE) ^ ((col % 32) * 4);这行代码看起来像是在乱改偏移但它把原本顺序访问的模式打散成了跨 bank 的均匀分布。实测在加了 8 行 XOR swizzle 后bank conflict 基本消失张量核心利用率直接提升了一截。这个阶段最值得留意的是不要凭直觉猜 swizzle 公式而是每次调整后都重新跑分析器看冲突计数是不是真的下降。4.3 第三轮双缓冲升级为深度流水线延迟隐藏更进一步双缓冲能解决搬一阵、停一阵的交替问题但还不能完全覆盖 TMA 的启动延迟。TMA 搬运一个大 tile 的延迟可能达几百个周期如果流水线只有两级消费者等数据到达时仍会有明显暂停。DeepGEMM 后续把缓冲数加到三份让 TMA 始终在搬运下一块或下下一块的数据消费者计算当前块。从代码结构上讲这需要维护一个下一块指针队列而不是简单地切换 cur 和 next。伪代码上表现为// 三级缓冲的索引管理环形队列 unsigned buf_idx (k_step % 3); unsigned next_idx ((k_step 1) % 3); unsigned after_next_idx ((k_step 2) % 3);三级缓冲换来的是更高的流水线余量让 wgmma 和 TMA 的时序抖动被缓冲吸收。实测三级缓冲相比二级在超大 K 的矩阵上提升了大约 8% 到 10%。这是因为 K 越大每个输出 tile 需要连续搬运的数据块越多流水线越深越不容易把空缺暴露出来。至于要不要继续加到四级、五级取决于 SMEM 容量预算毕竟还要给 wgmma 的累加结果留位置。4.4 实测数据不同 shape 下 DeepGEMM 与官方库对比以下是某块 Hopper 架构 GPU 上的实测对照数据。需要说明的是官方库的版本更新很快不同版本在同一 shape 上数据会浮动这里只记录当时测试环境下的结果作为策略参考而非排名结论。矩阵规模 (M x N x K)实现方案耗时 (ms)TFLOPS8192 x 8192 x 4096官方库 BF160.86约 6408192 x 8192 x 4096DeepGEMM FP8 双缓冲0.54约 10198192 x 8192 x 4096DeepGEMM FP8 三缓冲swizzle0.49约 11231024 x 4096 x 2048官方库 FP80.041约 4191024 x 4096 x 2048DeepGEMM FP80.037约 464大矩阵规模下DeepGEMM 通过三缓冲和 swizzle 达到了硬件峰值的 50% 以上相比双缓冲版本有明显优势。但中小矩阵下差距就小很多了原因有两个CTA 总数量不够无法把所有 SM 都占满TMA 的单次搬运块偏小延迟隐藏的空间有限。这也验证了之前定下的边界主打超大矩阵小矩阵场景用简单路径就够了。5. FP8 运算的精度陷阱两个尾数位省下来的成本5.1 为什么 FP32 累加器并不能解决所有问题FP8 的输入只有 3 位尾数E4M3相乘结果的尾数位更是有限。即使 Tensor Core 支持 FP32 累加累加时的舍入顺序依然会带来浮点误差。举个例子连续累加几千个数如果其中某个中间结果相对其他值大很多小值就会被大值完全淹没这在浮点运算中被称为吸收消除。FP32 累加器可以缓解尾数截断但无法修复输入本身的量化噪声。想进一步提高数值稳定性有两条常用手段。其一是随机舍入以数据低位为概率将结果向上下两个方向随机舍入这样多次计算的平均误差趋于零。其二是把累加顺序做树形规约在硬件允许范围内尽量让相同量级的数先相加减少大数对小数的影响。DeepGEMM 的默认路径使用普通顺序累加因为 Tensor Core 的硬件路径已经固定但在需要高精度的场景可以切换到随机舍入模式。5.2 scale 设计per-tensor、per-channel 与在线量化scale 的选择直接决定 FP8 的量化误差。最简单的 per-tensor 方案是拿整个矩阵的最大绝对值去除得到一个统一的 scale。实现简单但如果激活值分布不均比如某些通道整体极小某些通道极大统一 scale 会让小通道几乎退化成零。策略开销误差特点DeepGEMM 支持情况per-tensor最低分布不均时误差大默认per-channel中等对权重效果好激活仍需谨慎可配per-group (128元素一组)较高误差最平稳训练场景建议训练场景中在线量化需要在每次前向时动态统计 scale这个操作本身又增加了一轮归约计算。DeepGEMM 在推理路径上采用离线校准的方式把 scale 在部署前算好在训练路径上则按 128 元素分组做动态统计配合随机舍入让 FP8 的精度损失接近 BF16 结果。5.3 实测失真的对照实验我在一个中等规模的 Transformer 训练任务上对比了三种方案。BF16 基线训练到固定步数后 loss 为 0.32FP8 per-tensor 方案为 0.35训练曲线尾部有抖动FP8 per-group 加随机舍入方案回到 0.33 附近收敛曲线与 BF16 基本贴合。这说明 FP8 用在训练上完全可行前提是 scale 不能偷懒至少要做分组级的统计。推理阶段相对宽松。固定权重的前提下per-tensor scale 带来的精度损失可能只有几个百分点对于大部分应用可以接受。这里我给个实际建议如果你的部署模型对精度非常敏感先在验证集上同时跑 per-tensor 和 per-channel 两组推理比较输出分布的变化再决定是否需要更复杂的量化策略不要一上来就上最贵的方案。6. 排错实录一次 K 维不对齐引发的 48 小时定位6.1 现象某些 shape 完全正常另一些随机出错DeepGEMM 跑常规 shape 时一直很稳直到某次我测试 M6144、N6144、K1280 的矩阵乘法结果出现随机性错误。更奇怪的是M4096、N4096、K1024 时一切正常K1024 和 K1280 都是 32 的整数倍为什么一个对、一个错复现时还是概率性的有时连续跑五次错两次有时全部正确非常折磨人。这种概率性出错通常指向两种情况数据竞争或者越界写入。如果只是算术逻辑错误结果应该是稳定地错而不是随机错。所以我第一时间把怀疑对象锁定在 TMA 搬运和共享内存生命周期管理上而不是数学公式。6.2 排查链路从数据断言到 TMA descriptor 边界第一步是加数据断言。我在每个 k_step 之后用 CPU 复算同一块输出和 GPU 结果比对发现偏差出现在特定的 tile 组合上而且往往发生在后半段 K 循环。接着关闭双缓冲把流水线退化成搬一块、算一块、再搬一块错误出现的概率降低但没消失。这一步排除了纯粹的双缓冲竞争——若是生产者消费者同步问题单缓冲下应该完全消失。第二步是用内存检测工具跑同一 case。工具很快报出共享内存越界写入的警告指向一个 CTA 在某次循环里写到了自己 SMEM 范围之外。顺着这个线索我开始逐行检查 TMA 描述符的边界计算尤其是 K 维循环接近末尾时的 tile 尺寸处理。最后发现问题出在一个边界分支里当 K 的剩余长度不足一个 K_BLOCK 时代码仍然用了完整的 K_BLOCK 去构造描述符导致 TMA 读取越界数据同时把共享内存写坏。6.3 根因tile 边界的回卷写覆盖与同步缺失为什么 K1024 正常而 K1280 就出问题因为 1024 除以 32 得到 32 步循环干净利落1280 除以 32 也是 40 步但某些 CTA 因为输出 tile 在 M、N 方向靠边剩余的 A 或 B 数据不足 32走了退化路径这时候计算 K 块尺寸的偏移量出现了偏差。这个偏差在地址计算上表现为回卷某个索引在本应结束的地方继续向后延伸覆盖了相邻缓冲区的数据。修复方案有两步。第一步是让每个 K 步都用实际剩余长度裁剪 TMA 描述符而不是盲信 K_BLOCK 常量第二步是给退化路径单独加一道 barrier确保最后一个未满 K 块在写入 SMEM 后所有消费者 warp 都能看到完整数据才启动 wgmma。6.4 同类隐患L2 hint、原子累加与分配器对齐这次排错让我意外收获了一串同族隐患。第一个是 L2 缓存策略TMA 搬运时可以指定数据在 L2 中的偏好如果对输出 tile 设置了错误的淘汰优先级某些数据可能被过早逐出在后续重算时引发不一致。第二个是多个 CTA 对同一输出区域的原子累加FP32 原子加不保证顺序结果浮点误差会在不同 run 间波动。第三个是共享内存分配的对齐问题128 字节对齐和 4 字节对齐在 TMA 描述符眼里完全不同写错一个字段就可能导致跨边界读。疑似问题排查手段结论双缓冲竞争切换单缓冲复测排除数据取值逻辑CPU 复算分组比对定位到特定 tileSMEM 越界内存检测工具扫描确认越界TMA 描述符边界逐行检查 K 剩余长度根因K 块未裁剪如果你也遇到类似小矩阵没问题、大矩阵偶发错的 case我的建议是直接运行内存检测工具再配合单缓冲复测把手动同步的变量排除干净。这类问题最怕一上来就猜是浮点误差把方向带偏。拿到确凿的越界报告后再回头检查所有描述符、指针、步长参数往往比自己盯着代码发呆更高效。最后再分享一个实际观察GEMM 内核的调试和调优要分开对待。调优阶段可以大胆试各种 swizzle 和缓冲深度因为错误是好观测的但一旦进入排错阶段先复算简单路径、再上检测工具、最后才动代码这个顺序不要乱。跑大量边界 shape 的矩阵乘法作为回归用例比单点性能数据更值得长期维护——像 K 维不整除、N 方向剩余窄条这类隐蔽问题都是靠边界用例抓出来的。 DeepGEMM：基于 Hopper Tensor Core 的 FP8 矩阵乘法内核优化实践
很多搞深度学习的朋友都有过这种体验GPU 利用率监控看着没毛病可训练吞吐就是上不去排查链路转了一圈才发现时间几乎全耗在同一类运算上——矩阵乘法。我之前做模型性能剖析时统计过从 embedding 到输出层GEMM 类算子能占到整体耗时的七成以上。这也是我后来动手写 DeepGEMM 这套 FP8 矩阵乘法内核的根本原因与其在框架层反复做图优化不如直接吃透最底层这块硬骨头。DeepGEMM 的目标很单纯——利用 Hopper 架构 GPU 上 Tensor Core 与 TMA 的能力实现一个针对 FP8 输入、FP32 累加的矩阵乘法内核。它不是一个要替代工业级内核库的通用方案更像是我对自己的一次摸底GEMM 到底怎么设计才能在特定 shape、特定精度下把硬件算力吃干净。这篇文章会从为什么自研、底层原理、实现路径、性能调优到一次完整的排错过程都摊开来讲适合正在优化自有算子、或者单纯想搞清楚 Tensor Core 内部逻辑的工程师和研究者。1. 为什么自己写 GEMM模型跑得慢不一定怪 GPU1.1 先算一笔账一次 Transformer 前向里的 GEMM 占比随便拿出一个中等规模的 Transformer 模型前向计算里最重头的部分就是四个映射矩阵乘法Q、K、V 的投影是一次大 GEMMattention score 的计算是一次 batched GEMMscore 对 V 的加权又是一次 batched GEMMFFN 两层则贡献了剩余两次超大矩阵乘。再加上 softmax、LayerNorm 这些逐元素算子和归一化操作整个 forward 的 GEMM 占比很容易超过 70%。训练场景还多一个反向传播矩阵乘的数量直接翻倍这个比例只会更高。所以当你在性能剖析工具里看到某个 GEMM 算子耗时异常别急着怀疑数据加载或者并行通信先检查算子本身的实现效率。这也是自研 GEMM 内核最直接的收益点你可以精确控制每一字节的访问路径和每条硬件指令的调度而不是被迫接受别人预设的策略。我见过好几次换库版本后推理提速的案例本质就是新版本对某些 shape 选择了更合适的 tile 切分而不是什么玄学优化。1.2 cuBLAS 和 CUTLASS 都很好为什么还要有 DeepGEMM官方库封装得越彻底对调用者来说黑盒程度就越高。cuBLAS 确实在绝大多数 shape 上表现优秀但它不会把内核实现的细节暴露给你一旦遇到非标准 shape 或特定精度需求你只能按它提供的预估接口去猜性能行为。CUTLASS 项目给了高度灵活的模板组合理论上可以拼出各种内核策略但它概念繁多、模板深度很重想在里面改一条 TMA 加载路径或调整 swizzle 策略学习成本相当高。DeepGEMM 的定位就是轻量定制内核。它不需要覆盖所有硬件平台和所有矩阵 shape只需要在自己的目标机器上做到足够快。这种取舍带来两个直接好处一是代码量小、可读性强每一条关键路径都能很快定位二是遇到性能问题时可以拿官方库做基准对照从 profiling 数据中反推差距在哪个环节真正把性能优化的能力握在自己手里。如果你的业务只是定期跑几类固定矩阵规模自研内核完全有理由比通用库更贴合你的计算模式。1.3 DeepGEMM 的目标边界FP8、Hopper 与工程妥协做任何内核项目都要先划清边界否则会因为什么都想支持而陷入无尽复杂度。DeepGEMM 的初始范围如下精度A、B 矩阵输入为 FP8E4M3/E5M2 可配累加器用 FP32硬件仅针对 Hopper 架构 GPU 做优化重点使用其 TMA、wgmma 和分布式共享内存能力Shape面向超大矩阵M、N 在 2048 以上做满波优化小矩阵则退回到简单路径不做过深调优扩展性后端预留了 BF16 路径主要用于精度对照实验不作为默认主线。这个边界意味着我不需要处理 Ampere 等旧架构的兼容性可以把精力全部投入到大矩阵、批量加载、深流水线这一条主路上。工程上做减法和划边界是同样的道理方案越单纯越容易分析出性能瓶颈范围一旦蔓延调试时会发现每个符号都可能是错的。2. GEMM 内核的底层逻辑从最朴素循环到 Tensor Core 的关键跃迁2.1 最简单的 GEMM 内核长什么样写一个在 CPU 上可以跑通的朴素版 GEMM只需要三重循环对每个输出元素 C[i][j]累加 A 的第 i 行与 B 的第 j 列的内积。拿到 GPU 上如果也照这个逻辑写结果就是灾难每个线程反复从全局内存读取同一段数据复用率趋近于零带宽瞬间成为瓶颈。一个 GPU 上有几千个并发线程如果它们都在串行读同一批 A 元素缓存会被完全冲垮。// 最朴素的 GEMM 思路每个线程算一个输出元素 // 不考虑任何复用、布局、向量化仅作为反面教材 __global__ void naive_gemm_kernel(const float* A, const float* B, float* C, int M, int N, int K) { int row blockIdx.y * blockDim.y threadIdx.y; int col blockIdx.x * blockDim.x threadIdx.x; if (row M || col N) return; float sum 0.0f; for (int k 0; k K; k) { sum A[row * K k] * B[k * N col]; } C[row * N col] sum; }这段代码每个线程要读 K 个 A 元素和 K 个 B 元素如果一个 warp 的 32 个线程各自负责相邻 row 但相同 colB 的读取在缓存层面还能有一定复用A 的读取则完全重复。数据量稍微一涨DRAM 吞吐就直接封顶。真正高性能的 GEMM 内核必须解决两个核心问题数据复用结构怎么设计计算指令怎么选择。2.2 计算强度FLOPs 再高也怕带宽卡脖子判断一个计算任务是否访存密集还是计算密集关键看计算强度——单位字节的数据能支撑多少次浮点运算。一个 MNKn 的方阵乘法总计算量是 2n^3 FLOPs输入数据量是 3n^2A、B、C 各一次读写。把两者相除计算强度等于 2n/3。n1024 时约 680 FLOPs/Byte远高于现代 GPU 的算访比比如某 Hopper GPU 的 FP8 算力若按 1900 TFLOPS、HBM 带宽按 3.35 TB/s算访比约 570意味着理论上是计算密集型。可现实中为什么很多 GEMM 跑不出峰值因为没有有效的数据复用结构实际从 DRAM 搬运的字节数是理论值的几十上百倍。所以优化的本质就是通过分块让数据尽可能留在片上。一个数据元素被读到寄存器后如果能参与多次乘累加DRAM 的压力就成倍下降。考虑 tiling 之后计算强度公式会变成一个关于 tile 尺寸的函数块越大K 维循环中每个加载元素被使用的次数越多。2.3 三级存储的 Tiling 协作寄存器、SMEM、GMEMGPU 芯片上的存储层级可以类比成厨房的操作台全局内存相当于仓库共享内存相当于灶台寄存器相当于手边已经洗净备用的配料。GEMM 内核算法的核心就是设计一条从仓库批量取货、放到灶台上、再在手里快速加工的流水线。传统 CUDA 手动调度路径里一个 CTA 通常负责一个输出 tile先在共享内存铺好 A 和 B 的块再由各线程从 SMEM 读取数据做乘累加。Tile 尺寸的选择直接决定复用率和硬件资源占用。常见策略是每个 CTA 负责 128x128 或 128x256 的输出 tileK 维每次处理 16 或 32。计算一下 SMEM 占用就知道A 和 B 两个 tile 若都是 128x32且 A、B 都是 1 字节的 FP8需要的共享内存约 8KB如果上双缓冲翻倍到 16KB。这个数字相对于一块 Hopper GPU 上 228KB 的 SMEM 来说非常充裕于是还可以进一步加宽 N 方向的 tile用更多的片上数据换取更低的 DRAM 访问频率。A tile 128 x 32 -- FP8: 128*32 4096 字节 B tile 32 x 256 -- FP8: 32*256 8192 字节 双缓冲总占用: 2 * (4096 8192) 24576 字节 Hopper SMEM 典型容量约 228KB留白充足这个算术题看起来很朴素但它决定了内核能否在寄存器、SMEM、GMEM 之间保持足够深的流水线。数据一旦能在片上多待一轮全局内存的带宽压力就小一轮。3. DeepGEMM 的实现路径FP8 格式、Tensor Core 与 warp 分工3.1 FP8 的 E4M3 与 E5M2两个尾数位的取舍FP8 不是一个单一的浮点格式它有 E4M3 和 E5M2 两种子格式。E4M3 用 4 位指数加 3 位尾数取值范围相对窄但精度更高E5M2 用 5 位指数加 2 位尾数动态范围大但尾数噪声更明显。做矩阵乘法时通常把激活值用 E4M3 存储权重则根据是否需要大动态范围在 E4M3 和 E5M2 之间选择。格式符号位指数位尾数位最大值范围适用场景E4M3143约 448激活值、权重的主选E5M2152约 57344梯度、中间动态范围大的数据FP8 本身不能直接用于精度要求高的累加。实际计算必须以 FP32 作为累加器同时给输入乘一个缩放因子scale。DeepGEMM 在加载数据后、送入 Tensor Core 之前会先检查数据的整体范围将 FP8 值反量化到实际浮点范围。这个过程不是每元素做一次除法而是提前算好 scale最后在输出端乘回来。这个设计在代码层看起来只是多了两个乘法但它决定了整个量化误差的分布形态。3.2 Tensor Core 指令形态从 mma.sync 到 wgmma要真正把 FP8 算力烧起来光靠普通的 FMA 指令是不够的。Tensor Core 本质上是一组专用的矩阵乘累加硬件单元需要通过特定指令驱动。在 Ampere 时代最直接的编程方式是 mma.sync 指令族比如 m16n8k16 这样的形状表示一次让一个 warp 协同完成 16x8x16 的矩阵乘片段。Hopper 时代则提供了更强的 wgmma 指令它不再以一个 warp 为单位而是以一个 warp group 为单位一次可以处理更大的矩阵块并且支持异步流水线。DeepGEMM 的核心计算路径选择 wgmma 作为主用路径。使用 wgmma 的好处显而易见一次指令能覆盖更大的输出范围硬件调度也允许与数据搬运输送完全重叠。但也带来了一个副作用——warp 之间的同步和描述符管理变得复杂。一条 wgmma 指令需要引用 A、B 矩阵在共享内存中的地址布局描述符这些描述符由 Tensor Core 的 TMA 硬件读取若描述符内部的 stride 或 swizzle 参数写错结果会非常离谱。3.3 生产者-消费者分工TMA 搬运与计算解耦Hopper 最值得利用的新特性之一就是 TMATensor Memory Accelerator。它允许你用一条指令把全局内存中一个 2D 或 3D tile 直接搬到共享内存还可以在搬运过程中完成 swizzle 布局变换。这意味着传统内核里线程算地址、加载、再同步的模式被彻底改变了数据搬运不再消耗计算线程的指令周期。DeepGEMM 把 warp 分成生产者和消费者两类。生产者的职责只有两个准备 TMA 描述符、发布搬运请求消费者的职责是等待数据到达共享内存然后启动 wgmma 计算。这个模型很像工厂里的采购员和装配工采购员只管把原料运进仓库装配工不管原料怎么来的只等仓库里有货就开工。// 概念性伪代码数据流水线主循环 // 仅用于说明生产者-消费者结构不表示完整内核实现 for (unsigned k_step 0; k_step K / K_BLOCK; k_step) { if (k_step 1 K / K_BLOCK) { // 生产者角色预取下一块 A、B 到另一个共享内存缓冲 tma_load_to_smem(A_smem[1 - cur], A_global_desc); tma_load_to_smem(B_smem[1 - cur], B_global_desc); } // 消费者角色等待当前缓冲数据就绪后启动 wgmma wgmma_async(acc_smem, A_smem[cur], B_smem[cur], ...); }这个示例里的双缓冲结构把搬下一块和算当前块完全重叠只要 TMA 的搬运速度大于 wgmma 的消费速度流水线就能一直保持满载。实际工程里缓冲深度还可以加到三份甚至四份换取更大的容忍波动能力。3.4 分布式共享内存跨 CTA 协作的进阶玩法Hopper 还有一个容易被忽视的能力集群内的多个 CTA 可以组成一个线程块集群通过分布式共享内存机制互相读取对方在共享内存中的数据。这意味着一个 CTA 计算输出 tile 时不需要把完整的 A 或 B tile 都装进自己那一份 SMEM 里而是可以让相邻 CTA 把自己的 A tile 分享出来。DeepGEMM 在部分 shape 上启用了这个功能。比如一个大矩阵的 N 方向被切成很多窄条相邻 CTA 所需的 B tile 完全一样此时让其中一个 CTA 加载并广播给其他 CTA能省掉重复的 DRAM 读取。这个特性让 tile 切分策略有了新的维度不再单纯为了减少全局访问而把 SMEM 塞满还可以通过集群通信实现数据共享。当然分布式共享内存也带来同步复杂度。跨 CTA 访问数据必须用集群级的 barrier 保证顺序稍有不慎就会出现读到了半成品数据的问题。我的建议是第一版内核先别碰这个特性等单 CTA 路径跑稳、性能基线确定后再作为进阶优化加上去否则出现问题时会很难分清是 TMA 描述符写错还是集群同步丢步。4. 性能调优在 Nsight Compute 面前逐项逼出峰值4.1 第一轮先跑通基线版本记录真实的短板任何优化都离不开可量化的起点。DeepGEMM 第一版是直接照搬常见 CUTLASS 思路写出来的基础双缓冲版本没有做 swizzle也没有分布式共享内存。在这种状态下绝大多数矩阵 shape 是能跑对结果的但性能大概只有硬件峰值的 30% 左右。此时最该做的事不是继续加特性而是用分析器看清瓶颈到底在哪个硬件单元。Nsight Compute 里有几个关键指标值得重点看sm__pipe_tensor_cycles_active反映张量核心忙碌程度l1tex__data_bank_conflicts反映共享内存 bank 冲突严重性dram__throughput反映显存压力。第一版实测数据中tensor 核心利用率不高同时 bank conflict 告警非常密集。这说明当前内核的问题在数据供给路径而不是计算单元本身。这个判断把优化方向从改计算指令引导到了改数据布局。4.2 第二轮SMEM 地址布局与 bank conflict 的博弈共享内存按 bank 组织同一地址周期只能让每个 bank 输出一份数据。如果多个线程同时访问落在同一个 bank 但不同地址的位置就会发生所谓的 bank conflict一次访存被拆成多次。GEMM 这种高并发读 SMEM 的负载对 bank 布局极度敏感。最简单的布局下A tile 按行连续排列一行 128 字节。如果一行正好跨越了 32 个 bank 的两个周期而 warp 里的线程又在读同一列的数据冲突就会出现。解决办法是给地址加上一层 swizzle让每个线程读到的逻辑位置映射到不同的物理 bank。以常见 128 字节行宽为例可以按 8 行一组做异或扰动// 常见 swizzle 策略按行、列做 XOR 扰动 // 实际使用哪种模式取决于 tile 形状和 SMEM 宽度 constexpr int SMEM_STRIDE 128; // 字节 int swizzled_offset (row * SMEM_STRIDE) ^ ((col % 32) * 4);这行代码看起来像是在乱改偏移但它把原本顺序访问的模式打散成了跨 bank 的均匀分布。实测在加了 8 行 XOR swizzle 后bank conflict 基本消失张量核心利用率直接提升了一截。这个阶段最值得留意的是不要凭直觉猜 swizzle 公式而是每次调整后都重新跑分析器看冲突计数是不是真的下降。4.3 第三轮双缓冲升级为深度流水线延迟隐藏更进一步双缓冲能解决搬一阵、停一阵的交替问题但还不能完全覆盖 TMA 的启动延迟。TMA 搬运一个大 tile 的延迟可能达几百个周期如果流水线只有两级消费者等数据到达时仍会有明显暂停。DeepGEMM 后续把缓冲数加到三份让 TMA 始终在搬运下一块或下下一块的数据消费者计算当前块。从代码结构上讲这需要维护一个下一块指针队列而不是简单地切换 cur 和 next。伪代码上表现为// 三级缓冲的索引管理环形队列 unsigned buf_idx (k_step % 3); unsigned next_idx ((k_step 1) % 3); unsigned after_next_idx ((k_step 2) % 3);三级缓冲换来的是更高的流水线余量让 wgmma 和 TMA 的时序抖动被缓冲吸收。实测三级缓冲相比二级在超大 K 的矩阵上提升了大约 8% 到 10%。这是因为 K 越大每个输出 tile 需要连续搬运的数据块越多流水线越深越不容易把空缺暴露出来。至于要不要继续加到四级、五级取决于 SMEM 容量预算毕竟还要给 wgmma 的累加结果留位置。4.4 实测数据不同 shape 下 DeepGEMM 与官方库对比以下是某块 Hopper 架构 GPU 上的实测对照数据。需要说明的是官方库的版本更新很快不同版本在同一 shape 上数据会浮动这里只记录当时测试环境下的结果作为策略参考而非排名结论。矩阵规模 (M x N x K)实现方案耗时 (ms)TFLOPS8192 x 8192 x 4096官方库 BF160.86约 6408192 x 8192 x 4096DeepGEMM FP8 双缓冲0.54约 10198192 x 8192 x 4096DeepGEMM FP8 三缓冲swizzle0.49约 11231024 x 4096 x 2048官方库 FP80.041约 4191024 x 4096 x 2048DeepGEMM FP80.037约 464大矩阵规模下DeepGEMM 通过三缓冲和 swizzle 达到了硬件峰值的 50% 以上相比双缓冲版本有明显优势。但中小矩阵下差距就小很多了原因有两个CTA 总数量不够无法把所有 SM 都占满TMA 的单次搬运块偏小延迟隐藏的空间有限。这也验证了之前定下的边界主打超大矩阵小矩阵场景用简单路径就够了。5. FP8 运算的精度陷阱两个尾数位省下来的成本5.1 为什么 FP32 累加器并不能解决所有问题FP8 的输入只有 3 位尾数E4M3相乘结果的尾数位更是有限。即使 Tensor Core 支持 FP32 累加累加时的舍入顺序依然会带来浮点误差。举个例子连续累加几千个数如果其中某个中间结果相对其他值大很多小值就会被大值完全淹没这在浮点运算中被称为吸收消除。FP32 累加器可以缓解尾数截断但无法修复输入本身的量化噪声。想进一步提高数值稳定性有两条常用手段。其一是随机舍入以数据低位为概率将结果向上下两个方向随机舍入这样多次计算的平均误差趋于零。其二是把累加顺序做树形规约在硬件允许范围内尽量让相同量级的数先相加减少大数对小数的影响。DeepGEMM 的默认路径使用普通顺序累加因为 Tensor Core 的硬件路径已经固定但在需要高精度的场景可以切换到随机舍入模式。5.2 scale 设计per-tensor、per-channel 与在线量化scale 的选择直接决定 FP8 的量化误差。最简单的 per-tensor 方案是拿整个矩阵的最大绝对值去除得到一个统一的 scale。实现简单但如果激活值分布不均比如某些通道整体极小某些通道极大统一 scale 会让小通道几乎退化成零。策略开销误差特点DeepGEMM 支持情况per-tensor最低分布不均时误差大默认per-channel中等对权重效果好激活仍需谨慎可配per-group (128元素一组)较高误差最平稳训练场景建议训练场景中在线量化需要在每次前向时动态统计 scale这个操作本身又增加了一轮归约计算。DeepGEMM 在推理路径上采用离线校准的方式把 scale 在部署前算好在训练路径上则按 128 元素分组做动态统计配合随机舍入让 FP8 的精度损失接近 BF16 结果。5.3 实测失真的对照实验我在一个中等规模的 Transformer 训练任务上对比了三种方案。BF16 基线训练到固定步数后 loss 为 0.32FP8 per-tensor 方案为 0.35训练曲线尾部有抖动FP8 per-group 加随机舍入方案回到 0.33 附近收敛曲线与 BF16 基本贴合。这说明 FP8 用在训练上完全可行前提是 scale 不能偷懒至少要做分组级的统计。推理阶段相对宽松。固定权重的前提下per-tensor scale 带来的精度损失可能只有几个百分点对于大部分应用可以接受。这里我给个实际建议如果你的部署模型对精度非常敏感先在验证集上同时跑 per-tensor 和 per-channel 两组推理比较输出分布的变化再决定是否需要更复杂的量化策略不要一上来就上最贵的方案。6. 排错实录一次 K 维不对齐引发的 48 小时定位6.1 现象某些 shape 完全正常另一些随机出错DeepGEMM 跑常规 shape 时一直很稳直到某次我测试 M6144、N6144、K1280 的矩阵乘法结果出现随机性错误。更奇怪的是M4096、N4096、K1024 时一切正常K1024 和 K1280 都是 32 的整数倍为什么一个对、一个错复现时还是概率性的有时连续跑五次错两次有时全部正确非常折磨人。这种概率性出错通常指向两种情况数据竞争或者越界写入。如果只是算术逻辑错误结果应该是稳定地错而不是随机错。所以我第一时间把怀疑对象锁定在 TMA 搬运和共享内存生命周期管理上而不是数学公式。6.2 排查链路从数据断言到 TMA descriptor 边界第一步是加数据断言。我在每个 k_step 之后用 CPU 复算同一块输出和 GPU 结果比对发现偏差出现在特定的 tile 组合上而且往往发生在后半段 K 循环。接着关闭双缓冲把流水线退化成搬一块、算一块、再搬一块错误出现的概率降低但没消失。这一步排除了纯粹的双缓冲竞争——若是生产者消费者同步问题单缓冲下应该完全消失。第二步是用内存检测工具跑同一 case。工具很快报出共享内存越界写入的警告指向一个 CTA 在某次循环里写到了自己 SMEM 范围之外。顺着这个线索我开始逐行检查 TMA 描述符的边界计算尤其是 K 维循环接近末尾时的 tile 尺寸处理。最后发现问题出在一个边界分支里当 K 的剩余长度不足一个 K_BLOCK 时代码仍然用了完整的 K_BLOCK 去构造描述符导致 TMA 读取越界数据同时把共享内存写坏。6.3 根因tile 边界的回卷写覆盖与同步缺失为什么 K1024 正常而 K1280 就出问题因为 1024 除以 32 得到 32 步循环干净利落1280 除以 32 也是 40 步但某些 CTA 因为输出 tile 在 M、N 方向靠边剩余的 A 或 B 数据不足 32走了退化路径这时候计算 K 块尺寸的偏移量出现了偏差。这个偏差在地址计算上表现为回卷某个索引在本应结束的地方继续向后延伸覆盖了相邻缓冲区的数据。修复方案有两步。第一步是让每个 K 步都用实际剩余长度裁剪 TMA 描述符而不是盲信 K_BLOCK 常量第二步是给退化路径单独加一道 barrier确保最后一个未满 K 块在写入 SMEM 后所有消费者 warp 都能看到完整数据才启动 wgmma。6.4 同类隐患L2 hint、原子累加与分配器对齐这次排错让我意外收获了一串同族隐患。第一个是 L2 缓存策略TMA 搬运时可以指定数据在 L2 中的偏好如果对输出 tile 设置了错误的淘汰优先级某些数据可能被过早逐出在后续重算时引发不一致。第二个是多个 CTA 对同一输出区域的原子累加FP32 原子加不保证顺序结果浮点误差会在不同 run 间波动。第三个是共享内存分配的对齐问题128 字节对齐和 4 字节对齐在 TMA 描述符眼里完全不同写错一个字段就可能导致跨边界读。疑似问题排查手段结论双缓冲竞争切换单缓冲复测排除数据取值逻辑CPU 复算分组比对定位到特定 tileSMEM 越界内存检测工具扫描确认越界TMA 描述符边界逐行检查 K 剩余长度根因K 块未裁剪如果你也遇到类似小矩阵没问题、大矩阵偶发错的 case我的建议是直接运行内存检测工具再配合单缓冲复测把手动同步的变量排除干净。这类问题最怕一上来就猜是浮点误差把方向带偏。拿到确凿的越界报告后再回头检查所有描述符、指针、步长参数往往比自己盯着代码发呆更高效。最后再分享一个实际观察GEMM 内核的调试和调优要分开对待。调优阶段可以大胆试各种 swizzle 和缓冲深度因为错误是好观测的但一旦进入排错阶段先复算简单路径、再上检测工具、最后才动代码这个顺序不要乱。跑大量边界 shape 的矩阵乘法作为回归用例比单点性能数据更值得长期维护——像 K 维不整除、N 方向剩余窄条这类隐蔽问题都是靠边界用例抓出来的。