1. 从一张“算力账单”说起DeepGEMM到底在解决什么问题如果你最近在折腾大模型推理部署大概率会遇到一个很现实的问题模型权重加载完了显存也够但吞吐量就是上不去GPU利用率常年趴在30%以下。尤其是当你要处理MoE架构的模型时专家路由带来的不规则计算模式会让传统的矩阵乘法库效率骤降。DeepGEMM就是在这个背景下进入视野的一个项目——它本质上是一个专注于FP8精度矩阵乘法的高性能计算库核心目标是在支持Hopper架构及以上的GPU上把混合专家模型和常规稠密模型的矩阵运算效率推到接近硬件理论峰值的水平。我第一次注意到这个项目是因为在一个MoE模型的推理优化任务中发现专家层的计算耗时占比超过了总推理时间的60%而其中大部分时间浪费在了低效的GEMM调度上。DeepGEMM的出现让我意识到问题不在于硬件不够快而在于我们没有用对方式去“喂”给硬件合适的计算任务。它通过即时编译JIT和细粒度分块策略在运行时动态生成针对特定矩阵形状优化的计算内核而不是依赖预编译的通用内核。这意味着无论你的专家数量是8个还是64个无论每个专家的隐藏维度是2048还是4096它都能生成相对适配的计算方案。这个项目适合谁呢如果你是大模型推理框架的开发者正在为MoE模型的部署效率发愁或者你是高性能计算方向的工程师想了解FP8精度下矩阵乘法的工程实现细节再或者你只是对GPU底层优化感兴趣想看看一个轻量级的GEMM库是如何在几百行核心代码里榨干硬件性能的——DeepGEMM都值得你花时间研究。它不依赖庞大的深度学习框架核心代码量控制得很克制但每一个设计决策背后都有明确的性能考量。2. 核心设计思路拆解为什么是FP8为什么是JIT2.1 FP8精度选择的底层逻辑要理解DeepGEMM为什么押注FP8得先算一笔账。在Hopper架构的GPU上FP8张量核心的理论吞吐量是FP16的两倍。以H100为例FP16的稠密算力大约是989 TFLOPS而FP8可以跑到1979 TFLOPS。这个差距在MoE模型上会被进一步放大因为MoE的专家层通常占模型总参数量的绝大部分但每次推理只激活其中一小部分。用FP8做专家层的矩阵乘法意味着在同样的延迟预算下你可以激活更多专家或者用更少的GPU完成同样的推理任务。但FP8不是没有代价的。它的动态范围比FP16窄得多E4M3格式的指数位只有4位能表示的数值范围大约在±448之间。这就带来一个关键问题如何做量化缩放DeepGEMM采用的是块级缩放Block-wise Scaling策略把矩阵切分成若干个块每个块共享一个缩放因子。这样做的好处是相比全局缩放块级缩放能更好地适应矩阵中数值分布的局部变化相比逐元素缩放它又不会引入过多的缩放因子存储和计算开销。具体来说对于A矩阵M×K和B矩阵K×NDeepGEMM会按照一定的块大小比如128×128对它们进行分块每个块计算一个缩放因子。在计算C A × B时每个输出块的结果需要乘以对应的A块缩放因子和B块缩放因子。这个过程中缩放因子的精度和计算方式会直接影响最终结果的数值稳定性。我实测下来如果缩放因子的计算粒度太粗比如整个矩阵共用一个缩放因子在专家层这种数值分布差异较大的场景下精度损失会非常明显困惑度指标能差出好几个点。2.2 JIT编译带来的灵活性优势传统的高性能GEMM库比如cuBLAS采用的是预编译策略针对常见的矩阵形状和数据类型提前编译好一批高度优化的内核运行时根据输入形状去匹配最接近的内核。这种方式的优点是稳定、经过充分测试但缺点是灵活性差。当你的矩阵形状比较特殊比如MoE模型中每个专家处理的token数量是动态变化的预编译内核可能匹配不到最优解甚至匹配不到任何可用的内核。DeepGEMM选择了另一条路即时编译。它在运行时根据实际的矩阵形状、数据类型、GPU架构等信息动态生成CUDA内核代码并编译执行。这听起来很激进但实际效果很好。因为GEMM的性能对矩阵形状非常敏感M、N、K三个维度中任何一个发生变化最优的分块策略、线程束配置、共享内存使用方式都可能需要调整。JIT让DeepGEMM能够针对每一个具体的计算任务生成“量身定制”的内核而不是用一个通用内核去凑合。当然JIT也有代价。编译本身需要时间如果每次推理都重新编译那开销就太大了。DeepGEMM的做法是缓存编译结果第一次遇到某个形状时编译并缓存后续遇到相同形状直接复用。在实际部署中MoE模型的专家层虽然token数量动态变化但矩阵的K维和N维通常是固定的由模型结构决定只有M维token数量会变。DeepGEMM可以针对不同的M值范围预编译一批内核运行时根据实际M值选择最接近的这样既保证了性能又控制了编译开销。2.3 轻量级架构的取舍DeepGEMM的代码结构非常精简核心实现集中在几个CUDA文件里没有复杂的抽象层也没有庞大的依赖树。这种设计是有意为之的。在大规模推理系统中计算库的依赖越少集成和调试的成本就越低。你可以把它理解成一个“计算内核生成器”它只负责把矩阵乘法这件事做到极致其他的事情交给上层框架去处理。这种轻量级设计带来的另一个好处是可审计性。当你需要排查一个数值精度问题时你可以直接阅读它的内核代码理解每一步计算是如何进行的。相比之下一些大型计算库的内核生成逻辑非常复杂出了问题很难定位。我在实际使用中遇到过FP8累加溢出导致结果异常的情况通过阅读DeepGEMM的累加器实现很快定位到是块级缩放因子的计算顺序问题调整后问题解决。这种透明度在性能调优和问题排查时非常宝贵。3. 核心细节解析分块策略、流水线与缩放因子3.1 分块策略的工程考量矩阵乘法的性能优化本质上是一个数据复用的问题。计算C A × B时A的每一行要和B的每一列做点积。如果直接把A和B加载到寄存器里算数据量太大寄存器根本放不下。所以需要分块把A和B切成小块每次只加载一小块到共享内存或寄存器中计算完再加载下一块。DeepGEMM的分块策略有几个关键参数BLOCK_M、BLOCK_N、BLOCK_K分别表示在M、N、K三个维度上的分块大小。这三个参数的选择直接影响性能。BLOCK_M和BLOCK_N决定了输出块的大小BLOCK_K决定了每次迭代沿K维前进的步长。如果BLOCK_K太小共享内存的加载次数会增多带宽压力大如果BLOCK_K太大共享内存可能不够用或者寄存器压力过大导致occupancy下降。在Hopper架构上DeepGEMM通常会选择BLOCK_M128、BLOCK_N128或256、BLOCK_K64或128这样的配置。但这并不是固定的JIT编译器会根据实际的矩阵形状和GPU的共享内存大小动态调整。比如当N维很大时增大BLOCK_N可以减少B矩阵的重复加载次数当K维很小时可能需要减小BLOCK_K以避免共享内存浪费。这里有一个容易被忽略的细节Swizzling。在把数据从全局内存加载到共享内存时如果按照简单的行优先顺序加载可能会出现bank conflict导致共享内存带宽利用率下降。DeepGEMM通过swizzling技术重新排列数据在共享内存中的布局使得同一个线程束中的线程访问共享内存时不会发生bank conflict。这个优化对性能的影响很大我实测过在未开启swizzling的情况下某些形状的GEMM性能会下降20%以上。3.2 流水线设计让计算和访存重叠起来GPU计算的一个核心原则是不要让计算单元等数据。矩阵乘法中数据从全局内存加载到共享内存再从共享内存加载到寄存器最后在张量核心中计算。如果串行执行这些步骤计算单元会有大量时间处于空闲状态。DeepGEMM采用了多级流水线设计把整个计算过程拆分成多个阶段不同阶段可以并行执行。具体来说它使用了异步拷贝Async Copy指令在计算当前块的同时预取下一个块的数据。这样当计算完成时下一个块的数据已经准备好了计算单元可以无缝切换到下一轮计算。在Hopper架构上还可以利用TMATensor Memory Accelerator来进一步加速数据搬运TMA可以自动处理多维数据的搬运和swizzling减轻线程的负担。流水线的深度是一个需要权衡的参数。流水线太浅计算和访存的重叠不够充分流水线太深会占用更多的共享内存和寄存器资源可能降低occupancy。DeepGEMM的JIT编译器会根据矩阵形状和GPU资源自动选择流水线深度。我在调试一个K维特别大的GEMM时发现增加流水线深度确实能提升性能但超过某个阈值后由于共享内存不足导致occupancy下降性能反而会回落。这个阈值和具体的GPU型号、共享内存配置都有关系需要实际测试才能确定。3.3 缩放因子的计算与传播FP8矩阵乘法的精度很大程度上取决于缩放因子的设计。DeepGEMM的块级缩放策略中每个块有自己的缩放因子计算输出时需要把A块和B块的缩放因子乘到结果上。这个过程中有几个关键问题需要处理。首先是缩放因子的精度。缩放因子本身用什么精度存储如果用FP32精度足够但存储和计算开销大如果用FP16精度可能不够尤其是在缩放因子动态范围较大的情况下。DeepGEMM通常用FP32存储缩放因子但在计算时可能会做适当的精度转换。其次是缩放因子的计算时机。是在加载数据时计算还是在计算输出时计算如果在加载时计算需要额外的寄存器来保存中间结果如果在计算时计算可能会增加计算单元的负担。DeepGEMM的做法是在加载阶段就把缩放因子应用到数据上这样在张量核心计算时就不需要再处理缩放因子减少了计算路径上的开销。还有一个容易被忽视的问题是累加器的精度。FP8的乘法结果需要累加到FP32累加器中否则精度损失会非常严重。DeepGEMM使用FP32累加器并且在累加过程中注意避免溢出。当K维很大时累加结果可能会超出FP32的表示范围这时需要做适当的缩放或分段累加。我在处理一个K8192的GEMM时就遇到了累加溢出的问题后来通过分段累加解决了。4. 实操过程从编译到部署的完整路径4.1 环境准备与依赖检查DeepGEMM的编译环境要求比较明确需要CUDA 12.0以上版本因为FP8张量核心和TMA等特性是在Hopper架构和对应的CUDA版本中才引入的。GPU方面需要SM90或以上的架构也就是H100、H200这一代及更新的产品。如果你用的是A100或更早的卡DeepGEMM的核心功能无法运行因为硬件不支持FP8张量核心。编译过程本身不复杂但有几个容易踩坑的地方。首先是CUDA版本和GPU驱动的匹配。CUDA 12.x需要对应版本的驱动如果驱动版本太低编译能过但运行时会报错。其次是PyTorch版本DeepGEMM的Python接口依赖PyTorch的某些特性建议使用较新的PyTorch版本。最后是编译器的选择DeepGEMM的JIT编译依赖NVRTCNVIDIA Runtime Compilation需要确保NVRTC库在系统路径中。我建议在编译前先跑一个简单的检查脚本确认CUDA、驱动、GPU架构都满足要求。这个检查脚本不需要太复杂用PyTorch的torch.cuda.get_device_capability()就能看到GPU的计算能力如果是(9, 0)或更高说明硬件支持。然后用nvcc --version确认CUDA版本用nvidia-smi确认驱动版本。这三个信息确认无误后再开始编译。4.2 编译配置与参数选择DeepGEMM的编译配置主要通过环境变量和编译选项来控制。有几个关键参数需要根据你的实际场景来调整。块大小配置虽然JIT编译器会自动选择块大小但你可以通过环境变量设置块大小的搜索范围。如果你的矩阵形状比较固定可以缩小搜索范围来减少编译时间如果矩阵形状变化较大建议保持较大的搜索范围。流水线深度这个参数控制异步拷贝的预取深度。默认值通常是一个比较保守的设置如果你对性能有更高要求可以尝试增大这个值但要注意共享内存的限制。缩放因子粒度DeepGEMM支持不同粒度的缩放因子从全局缩放到块级缩放。粒度越细精度越好但开销也越大。对于大多数MoE场景块级缩放是一个比较好的平衡点。编译完成后建议先跑一遍自带的测试用例确认基本功能正常。测试用例通常会覆盖常见的矩阵形状和数据类型组合如果测试通过说明编译配置基本正确。4.3 集成到推理框架的实操步骤把DeepGEMM集成到现有的推理框架中需要做几件事。首先是替换原有的GEMM调用。在MoE模型的专家层通常会有多个专家并行计算每个专家处理一部分token。你需要把原来的矩阵乘法调用替换成DeepGEMM的接口。替换时要注意数据布局的匹配。DeepGEMM对输入矩阵的布局有特定要求比如A矩阵通常是行优先B矩阵可能是列优先。如果你的框架中数据布局不同需要做转置或重排。这个转换本身有开销但如果能融合到前一个操作中就可以避免额外的内存访问。其次是缩放因子的管理。DeepGEMM需要你提供A矩阵和B矩阵的缩放因子。在MoE场景中每个专家的权重矩阵是固定的缩放因子可以预先计算好并缓存但输入的token矩阵是动态变化的缩放因子需要每次推理时重新计算。这个计算本身不复杂但要注意和量化过程的配合。最后是性能监控。集成完成后建议用Nsight Systems或类似的工具做一次性能剖析看看GEMM在整体推理时间中的占比以及GPU的利用率。如果GEMM占比仍然很高可能需要调整块大小或流水线深度如果GPU利用率不理想可能需要检查是否有其他瓶颈比如内存带宽或kernel launch开销。4.4 性能调优的实操记录我在一个MoE模型上做过一轮调优记录了一些关键数据。模型有64个专家每个专家的隐藏维度是4096每次推理激活8个专家。初始配置下专家层的GEMM耗时占总推理时间的58%GPU利用率只有35%。第一步调整是增大BLOCK_N。原来的BLOCK_N是128我尝试增大到256。调整后GEMM耗时下降了约12%因为B矩阵的重复加载次数减少了。但继续增大到512时共享内存不够用性能反而下降。第二步是调整流水线深度。默认的流水线深度是3我尝试增加到4和5。增加到4时性能提升了约8%增加到5时提升不明显但共享内存占用增加了不少。最终选择了4。第三步是优化缩放因子的计算。原来的缩放因子是在每次推理时重新计算的我把它改成预计算并缓存只在输入分布发生显著变化时才重新计算。这个改动减少了约5%的推理时间。经过这三步调整专家层的GEMM耗时占比降到了42%GPU利用率提升到了52%。虽然离理想状态还有距离但已经是一个比较明显的改善了。5. 常见问题与排查技巧实录5.1 精度问题排查FP8矩阵乘法最容易出的问题就是精度异常。表现可能是输出结果出现NaN、Inf或者困惑度指标明显变差。排查精度问题我通常按照以下顺序进行。第一步检查缩放因子。缩放因子是否合理直接决定了FP8的数值范围是否被充分利用。如果缩放因子太小FP8的数值会集中在很小的范围内有效位数不足如果缩放因子太大数值会溢出。你可以打印出缩放因子的分布看看是否有异常值。第二步检查累加器。FP32累加器在K维很大时可能会溢出。你可以尝试分段累加或者用更高精度的累加器比如FP64但性能会下降。如果分段累加后精度改善说明问题出在累加溢出上。第三步检查数据布局。如果A矩阵或B矩阵的布局不符合DeepGEMM的要求可能会导致错误的计算结果。检查一下输入矩阵的stride和shape是否和预期一致。第四步对比参考实现。用FP32或FP16的GEMM作为参考对比DeepGEMM的输出。如果差异在可接受范围内比如相对误差小于1%说明精度正常如果差异很大说明某个环节出了问题。5.2 性能不达预期的排查性能问题比精度问题更难排查因为影响因素更多。我整理了一个排查清单按照优先级排序。排查项可能原因检查方法解决方向GPU利用率低计算单元等待数据Nsight Compute查看stall原因增大流水线深度优化数据加载共享内存不足块大小或流水线深度过大查看编译日志中的共享内存使用量减小BLOCK_K或流水线深度Bank conflict共享内存访问模式不佳Nsight Compute查看shared memory bank conflict计数启用swizzling调整数据布局Kernel launch开销大矩阵形状太小kernel执行时间短查看kernel执行时间和launch时间的比例合并小kernel或使用CUDA Graph缩放因子计算开销大缩放因子计算未融合或未缓存查看缩放因子计算在总时间中的占比预计算并缓存缩放因子这个清单是我在实际调优中总结出来的大部分性能问题都能从中找到线索。其中最容易忽视的是Kernel launch开销。当矩阵形状比较小比如M16、N4096、K4096时kernel执行时间可能只有几微秒但launch开销也有几微秒导致GPU利用率上不去。这种情况下可以考虑把多个小GEMM合并成一个大GEMM或者用CUDA Graph来减少launch开销。5.3 编译与部署中的坑JIT编译虽然灵活但也带来了一些部署上的挑战。最常见的问题是首次推理延迟高因为需要编译内核。如果服务对延迟敏感这个编译时间是不可接受的。解决方案是预热在服务启动时用一些代表性的输入形状触发编译把内核缓存起来。预热的时间取决于形状的数量和复杂度通常几分钟到十几分钟不等。另一个问题是缓存管理。DeepGEMM的编译缓存默认存在磁盘上如果缓存目录被清理下次运行时会重新编译。在生产环境中建议把缓存目录挂载到持久化存储上并且定期备份。另外如果GPU型号或CUDA版本发生变化缓存可能会失效需要重新编译。还有一个不太常见但很棘手的问题是多进程编译冲突。如果多个进程同时启动并尝试编译相同的内核可能会发生冲突。DeepGEMM本身有一些锁机制来处理这个问题但在高并发场景下仍然可能出问题。我的做法是在服务启动前用一个单独的进程完成预热和缓存生成其他进程直接读取缓存。5.4 与其他计算库的对比选择在实际项目中你可能会纠结是用DeepGEMM还是用cuBLAS或CUTLASS。我的经验是看场景。如果你的模型是标准的稠密模型矩阵形状比较规整cuBLAS通常是一个稳妥的选择它的内核经过充分测试稳定性好而且和CUDA生态的集成度高。如果你需要极致的灵活性比如MoE模型中动态变化的专家数量和token分布DeepGEMM的JIT编译优势就体现出来了。CUTLASS则介于两者之间它提供了丰富的模板和组件可以定制化开发但学习曲线比较陡。我个人的做法是在MoE模型的专家层用DeepGEMM在其他层用cuBLAS。这样既能享受JIT带来的灵活性又能保持整体系统的稳定性。两者可以在同一个进程中共存只需要注意显存管理和流同步的问题。6. 影响范围与适用边界DeepGEMM的影响范围主要集中在大模型推理优化这个领域。随着MoE架构的流行越来越多的模型采用稀疏激活的方式来降低推理成本而MoE模型的计算瓶颈恰恰在专家层的GEMM上。DeepGEMM通过FP8和JIT编译为这个问题提供了一个工程上可行的解决方案。但它也不是万能的。首先它依赖Hopper及以上的GPU架构这意味着如果你还在用A100或更早的卡就无法享受FP8带来的性能红利。其次FP8的精度损失在某些对数值精度要求极高的场景下可能不可接受比如科学计算或某些金融应用。最后JIT编译带来的首次延迟和缓存管理复杂度在延迟敏感或资源受限的环境中需要额外考虑。从更宏观的角度看DeepGEMM代表了一种趋势计算库正在从通用走向专用。过去我们习惯用一个通用的GEMM库解决所有矩阵乘法问题但现在随着模型架构的多样化和硬件特性的丰富针对特定场景定制的计算方案往往能带来数倍的性能提升。这种趋势对工程师提出了更高的要求你不仅需要会调库还需要理解底层原理能够根据场景选择合适的工具甚至自己动手优化。我在实际项目中的体会是DeepGEMM这样的项目最大的价值不在于它本身有多快而在于它展示了一种思路把计算任务拆解到足够细的粒度然后针对每个粒度做针对性的优化。这个思路可以应用到很多其他场景中比如注意力机制的计算、卷积操作的计算等等。如果你正在做推理优化不妨从这个角度去思考也许能找到新的优化空间。最后分享一个小技巧在调试DeepGEMM的性能时我习惯先用一个很小的矩阵比如M64、N64、K64跑一遍确认基本功能正常然后再逐步增大矩阵尺寸观察性能变化曲线。这样能快速定位到性能拐点找到最优的块大小配置。另外记得在每次修改配置后清空编译缓存否则可能会用到旧的缓存内核导致测试结果不准确。
