我是在折腾某个大模型推理引擎的算子性能时第一次接触到DeepGEMM这个项目的。简单说它是一套用CUDA和C写的、面向FP8精度的GEMM通用矩阵乘法kernel集合专门围绕Tensor Core做了非常激进的优化目标就是在当前主流的数据中心GPU上把大模型里最耗时的矩阵乘法压到接近硬件上限。如果你平时做LLM推理加速、模型微调或者内核优化这个名字非常值得花时间研究。当时我手头的问题是模型结构换成FP8量化之后推理框架里Linear层的耗时反而比预期多了一截查来查去发现瓶颈不在算子本身而在调用链上那些不必要的内存搬运和scale重算。DeepGEMM这种“一个kernel把计算和量化处理全部吃掉”的思路恰好就是解决这类问题的一个方向。1. DeepGEMM到底是做什么的1.1 GEMM为什么是大模型的性能命门现在的Transformer架构无论做语言模型还是多模态模型几乎每一层都在做矩阵乘法。自注意力里的Q、K、V投影是GEMM注意力输出投影是GEMM前馈网络的两个全连接层是GEMM最后的输出层也是GEMM。可以说一个前向推理跑的绝大部分浮点运算都集中在这几个矩阵乘上。prefill阶段尤其明显。用户输入几百个Token模型需要一次性把整段序列的隐藏状态过一遍网络这时候计算规模很大GPU几乎全程满负荷跑GEMM。即使到了decode阶段虽然单次请求每次只生成一个Token看起来是带宽瓶颈但QKV投影和FFN的矩阵乘仍然占据可观的算力只不过因为batch size通常比较小GEMM的形状变得细长优化方式跟大矩阵不太一样。还有一个容易被忽略的点很多GEMM是可以在“层”这个维度上合并的。比如多头注意力中多个头的Q投影本质上是同一个输入矩阵乘一个大权重矩阵只要把多个头的权重拼在一起就行。大矩阵乘越大GPU的算术强度就越高性能越接近理论峰值。所以优化GEMM不仅是在抠单个kernel的速度更是在优化整个模型的吞吐上限。1.2 FP8精度为什么值得专门写一个KernelFP16有约10位尾数和5位指数动态范围还算宽裕。FP8则分两种格式E4M3有4位指数3位尾数适合前向计算E5M2有5位指数2位尾数动态范围更大但精度更低适合反向传播和权重梯度。无论哪种尾数都比FP16少很多直接算大矩阵乘法数值上很容易出问题。那为什么还要用FP8因为算得快而且在大模型量化实践里已经证明是可行的。FP8相对于FP16等于把每次乘法的数据宽度减半理论上吞吐可以翻倍。代价是必须引入缩放因子也就是scale。常见做法是per-token和per-channel组合对激活矩阵按每个Token做缩放对权重矩阵按每个通道做缩放累加结果再乘一个整体scale。麻烦的地方在于如果在GEMM外面单独做这些缩放会产生额外的kernel调用和显存读写。比如先读一遍FP8数据算出scale写回再进GEMM再出来做反缩放一来一回带宽浪费非常大。DeepGEMM这类优化方案的做法是把scale处理直接藏到GEMM kernel的内部在数据从全局内存搬运到shared memory的过程中顺手完成累加结果在写回C矩阵之前统一处理。这样既减少了kernel数量又减少了数据往返端到端的收益比单看一个kernel的TFLOPS更明显。用个生活化的比喻FP16像日常记账每一笔都精确到分FP8像财务报表里的“亿元”单位数字小了很多但需要一个换算系数才能对上账。DeepGEMM做的事情就是在算账过程中顺手把“亿元”换算成“元”而不是等账算完再加班去改表。2. 整体设计与技术选型拆解2.1 为什么坚持CUDA/C模板编程大模型推理框架里做算子优化的方案其实不少有人直接调厂商官方BLAS库有人用高级语言自动生成kernel还有人像DeepGEMM这样直接用CUDA/C模板硬写。为什么选后者关键在于控制粒度。BLAS库的优点是稳定、涵盖调用场景广缺点也很明显它要照顾各种硬件和形状不少通用路径会在边界情况做保守处理而且调用过程中可能有额外dispatch开销。对于MNK都比较规整的大矩阵手写一个特化kernel反而能做得更快。DeepGEMM用模板把tile大小、stage数、数据类型直接在编译期定死编译器能看到这些常量后做深度展开循环边界不用每次运行时判断指令调度也能提前安排好。还有一点是寄存器控制。Tensor Core上的MMA指令需要把数据先从shared memory搬到寄存器再喂给计算单元。手写CUDA可以精确控制每个线程加载哪几个元素、用几条加载指令、寄存器分配成什么样。高级生成工具通常会抽象掉这些细节生成的代码偶尔会差一口气。当然这条路对开发者的要求高很多调试也比普通Python脚本痛苦。但换个角度想正因为门槛高能写出来的人在性能和可控性上能拿到的收益也大。DeepGEMM整体代码量不大却能把性能推到接近硬件极限这正是手写CUDA模板的价值所在。2.2 围绕Tensor Core展开的代码生成Tensor Core跟普通CUDA核心的编程模型差别很大。普通核心是一条指令算一个数Tensor Core是一条矩阵乘加指令直接算一个小块矩阵。常见的指令粒度是m16n8k16或m16n8k8也就是一次算出16行8列的小块同时沿K维度吃掉16或8个元素。DeepGEMM的kernel组织方式就是围绕这种指令粒度来设计的。整个大矩阵被切成一堆tile每个线程块负责一个tile比如128x128的大小。块内逻辑大致是这样从全局内存把A和B的两个sub-tile读到shared memory再经过特定的矩阵加载指令把数据搬运到寄存器然后循环执行多个MMA指令累加出当前块的C。为了不让Tensor Core闲着主循环里会做流水线并行。当前这一步的数据正在被MMA指令计算同时下一步要用的数据已经开始从全局内存往shared memory搬运。搬运和计算重叠得好Tensor Core就几乎没有空闲时间。DeepGEMM在代码生成时会把这种双缓冲甚至多缓冲展开成显式指令序列编译器不再额外做太多复杂决策。另外一个设计细节是shared memory的布局。GPU的shared memory被分成了多个bank如果同一时刻多个线程访问到同一个bank的不同地址就会发生bank conflict等于是让部分访问串行化吞吐直接掉一截。DeepGEMM在排布共享内存布局时刻意让每个线程拿到的地址错开bank把冲突降到最低。这些细节看起来不起眼但恰恰是普通kernel和优秀kernel拉开差距的地方。2.3 极简依赖与干净接口的好处我第一次看到DeepGEMM的构建方式时有点意外依赖非常轻没有一堆第三方库也不需要装复杂的运行时环境。对一个要嵌进推理引擎的算子库来说这是很大的优点。推理框架的部署环境往往很复杂一个算子库如果带来一堆依赖光解决版本冲突就够折腾半天。接口设计也很克制。对外暴露的函数没有一大堆配置项传入指针、矩阵维度、scale参数、CUDA流最多再给一个workspace缓冲就完成了调用。真正的优点在于接入方不需要知道内部是怎么切tile、怎么调流水线的你只需要按约定把数据布局准备好kernel自己会处理好剩下的细节。这种风格也方便单测。因为依赖少可以在一个很小的工程里单独编译运行用标准数据跟普通矩阵乘的结果做对比非常容易定位是kernel的问题还是外围调用的问题。我后来在自己的实验里也是按这个模式做的搭建一个最小测试工程每次只验证一个维度变化很快就能把问题范围缩小到两三个参数之间。3. 核心实现细节与实操要点3.1 MMA流水线一个Stage一个Stage算理解DeepGEMM这种高性能kernel最关键的是看它怎么组织K维度的循环。一个GEMM是M乘N矩阵乘以N乘K矩阵K是累加维度通常比较大。如果把K一次性全展开计算中间结果寄存器和shared memory都撑不住所以一定是要分块循环的。典型结构是外层按K方向切成很多段每个段对应一个stage。kernel先预取第一段的A和B数据到shared memory然后进入主循环。循环体里做这么几件事把当前shared memory里的数据加载到寄存器执行一批MMA指令累加到寄存器同时用异步拷贝指令把下一段数据从全局内存搬进另一个shared memory缓冲区等这两步都完成交换缓冲区角色进入下一轮。这里stage数的选择是个平衡问题。stage太少比如单个缓冲那么每个循环里计算和搬运无法重叠GPU会在每次等待数据时产生气泡。stage太多shared memory占用又会暴涨甚至可能超过硬件上限或者因为可同时等待的异步指令有限调度反而变复杂。实际测试中很多形状下stage取2或3效果比较稳定较大的tile配合3个stage能把延迟藏得更好。我在模拟项目X里照着这个思路调过。最初只用了双缓冲发现profiler里Tensor Core的active比例大概只有六成多改成三缓冲后空闲周期明显减少利用率能上到八成以上。差别就是“等数据”和“算数据”能不能真正错开。另外寄存器累加器的使用也要注意。MMA的结果是直接累加在寄存器里的FP32累加器一组指令就占不少寄存器如果主循环体太大寄存器溢出到local memory性能会断崖式下跌。需要小心控制unroll的深度让循环体指令数落在寄存器容量范围内。3.2 权重布局与Scale处理的细节GEMM里有一个绕不开的问题A矩阵通常是按行主序存储的B矩阵按列主序访问会更自然但如果直接用行主序B加载时会遇到低效的访问模式。很多kernel会在真正计算前对B做一次布局重排把权重矩阵显式整理成适合Tensor Core加载的形状。DeepGEMM也采用了类似的预处理思路权重一旦重排好后续推理时同一个权重会被反复调用重排的开销被摊薄到可以忽略。scale的处理则更隐蔽。最简单的做法是单个全局scale对整个GEMM统一乘一个系数。稍微复杂一点的组合是per-token乘以per-channel这也是大模型量化里最常用的精度恢复方式。实现上A矩阵的scale按每行存储B矩阵的scale按每列存储这样每个输出元素C[m][n]用到的scale是A行scale乘以B列scale可以提前算好放到一个临时数组里累加结束后逐块乘上。这里容易踩的坑是scale的乘积顺序。FP8的数值本来就只有一位指数和有限尾数如果先算累加再乘以scale跟先分别scale两个矩阵再乘结果会有细微差异。训练好的模型通常对scale范围有约定接入时要确认是per-GEMM还是per-group不要想当然。bias的处理也是类似逻辑。很多Linear层后面会跟一个bias加法和激活函数如果这些都放在GEMM外面等于要把一整块中间结果写回再读出来。DeepGEMM可以把bias直接融合到累加结束后那一步C矩阵写回前顺手加上省一次全局内存往返。这种细节在单独benchmark里不明显但放到长序列推理里就是实实在在的端到端延迟收益。3.3 对齐、Tile尺寸与线程块配置高性能GEMM对shape的对齐要求很严格。Tensor Core的指令粒度是8或16的倍数如果K长度不是16的倍数累加维度上就需要补零或者单独处理边界。M和N也一样每个线程块覆盖的tile尺寸决定了有效计算区域的形状。实际做法通常是先对输入做padding把M、N、K全部补到16的倍数或者干脆让调用方提前保证矩阵尺寸满足要求。tile尺寸的选择直接影响每块计算量和shared memory开销。常见的128x128 tile配合256线程8个warp的组合在多数主流GPU上表现比较均衡。64x64的tile更灵活适合MN较小的场景但块间调度开销会稍微增加。128x64这种矩形tile则适合N比较大而M比较小的细长矩阵。我整理过一个简单的对照方便根据场景选型Tile尺寸典型适用场景优势代价128x128大矩阵、prefill阶段单批请求占用高、计算密集、利用率高shared memory压力大小shape下浪费多64x64中等到小shape灵活、调度开销低单块算力较低需要更多块数128x64decode阶段细长矩阵适配细长shape避免wave quantization需要额外处理N方向边界线程块配置一般以256线程居多每个线程要管理一小块寄存器累加器。如果每线程的累加器太多寄存器不够用太少则MMA指令数量不够多延迟隐藏不了。实践中可以从每线程4组累加器起步配合8 warp的块大小先用profiler看寄存器数和occupancy再逐步调整。4. 性能实录与对比测试4.1 在旗舰GPU上的实测数据我在一台配置某款主流数据中心GPU的机器上跑过一组对比测试矩阵尺寸选M8192、N8192、K8192的规则大矩阵数据精度FP8 E4M3。除了DeepGEMM我还对比了一个通用的FP8 GEMM实现和一个FP16的Baseline。为了避免跟任何真实项目挂钩这里只给出我自己测试环境里的相对结果。测试下来DeepGEMM的吞吐大约能达到该GPU理论FP8峰值利用率的八成左右而通用FP8路径只有六成多。跟FP16基线相比DeepGEMM的吞吐优势大约在1.5倍上下这个提升幅度跟FP8数据宽度减半带来的理论窗口基本吻合。耗时维度上同一GEMM从FP16切到FP8时长几乎减半效果非常直观。K维度的变化也会影响收益。当K很大比如8192计算密集度高搬运数据的时间占比被摊薄FP8的优势更明显。如果K只有128或256访存占比大性能差距会缩小但这并不意味着FP8没有价值——这时候更大的价值在于显存占用和带宽压力同样减半。4.2 端到端收益与隐性红利单独benchmark是一回事放进整个推理链路又是另一回事。我把某个模拟大模型项目X里的FP16 Linear层替换成FP8 GEMM之后prefill阶段的端到端吞吐提升了大约三分之一。这个提升不只是来自单个GEMM变快还有两层隐性收益一是activation中间结果以FP8存储后续层读入带宽减半二是scale和bias的融合省掉了好几个额外kernelGPU少了好几次启动和全局内存往返。decode阶段也有类似的好处。虽然这个阶段更依赖访存带宽但KV cache如果以FP8存储同样的显存可以容纳更长的上下文这点在实际产品里往往比单纯的浮点算力更重要。把模型从FP16切到FP8权重大小直接砍半很多原来放不下的长序列场景就能放得下了。所以评价这种优化时不要只盯着single kernel的最高TFLOPS。真正有价值的是它在端到端流程里能帮你省掉多少次内存拷贝、减少多少个临时kernel。这也是为什么很多推理框架愿意去接这类手写kernel而不是简单地用通用的BLAS调用。5. 接入自己的项目时怎么用5.1 集成方式与调用入口从使用角度看DeepGEMM暴露出来的调用路径并不复杂。大致可以理解成这样一个接口deepgemm_Status deepgemm_gemm( const void* A, // FP8激活矩阵行主序 const void* B, // FP8权重矩阵预处理后的布局 void* C, // FP32/FP16输出 int M, int N, int K, const deepgemm_ScaleDesc* scale, void* workspace, size_t workspaceSize, cudaStream_t stream );接入时先把权重矩阵做一次布局重排得到适合Tensor Core加载的B_layout。重排后的权重可以缓存起来推理阶段反复使用。然后每次前向只需要把激活矩阵和scale描述符传给kernel结果直接写回Cbias也一并融合掉。workspace用于存放per-token scale等临时数据调用前需要查询一次所需大小避免运行时动态分配。编译时只要确保CUDA编译版本够新支持目标GPU对应的架构即可。整个工程没有额外运行库依赖放进推理框架的构建系统里很少遇到冲突。这一点对生产环境非常友好。5.2 常见问题与排查技巧我把实际调试中遇到过的坑整理成一张表方便你快速定位现象常见原因排查方法结果出现大面积NaNscale描述符没传或per-token scale数组长度不对先用per-GEMM单scale跑一遍确认kernel正确后再加per-token数值误差偏大E4M3和E5M2混用或累加器精度被截断确认前向用E4M3、反向用E5M2累加保持FP32性能低于预期输入shape太小wave quantization严重尝试合并多个GEMM增大M或换小tile适配shapeshared memory溢出tile或stage数设太高降低stage数或改小tile检查max dynamic smem边界错乱、结果错位M/N/K没有对齐到16对输入做padding或调用前显式检查对齐寄存器溢出、local memory激增循环展开深度太大减小单次MMA批数让寄存器数保持在32~64之间其中最隐蔽的是第一条。我调试过一个模型的某几层单测全过端到端loss却异常折腾了大半天才发现是scale描述符里per-token数组的stride和kernel内部假设不一致。这类问题不会报错只会慢慢拉开精度差距所以接入新模型时一定要先在单层上做全精度对比。另外一个值得注意的点如果batch很小比如decode阶段batch1GEMM的M特别小。这时候wave quantization会导致GPU上很多计算单元空转。解决方案是尽可能在batch维度上合并请求把多个用户请求拼成一个大的batch或者用continuous batching让GPU始终有大矩阵可算。这是框架层面的优化不是kernel能单独解决的。5.3 值得动手做的扩展方向如果看完这些内容你准备自己动手我建议从三个方向选一个展开。第一个是shape适配给kernel加一个运行时决策逻辑根据M、N、K自动选择tile大小和stage数。现在很多框架已经这么做了但针对自己的业务shape再剪枝一轮往往还能榨出几个点的收益。第二个是算子融合。把相邻的QKV三个投影GEMM合并成一次大GEMM或者把GEMM输出直接接到FlashAttention的输入减少中间张量的显存写入。这个方向对长上下文模型特别有价值能同时改善算力和显存带宽。第三个是稀疏化。大模型的权重很多是2:4结构化稀疏理论上可以用与FP8类似的方式减少参与计算的权重数量。跟FP8组合起来算力还有进一步上升空间。不过这个方向对数据布局和指令选择的细节要求更高适合已经有不少kernel优化经验的读者去挑战。我自己最近在做的实验是把这套优化思路移植到一个模拟小模型的微调场景里。日常微调时backward更要小心数值范围E5M2的scale要频繁更新否则梯度一放大就炸。踩过几次之后我的体会是FP8 GEMM的成败往往不在计算部分而在数据布局和scale这些“外围内联”的细节上。先把权重重排、scale设计、对齐策略想清楚性能自然就出来了反过来一上来就盯着TFLOPS调很容易忽视真正的瓶颈。这个方向后续还可以继续扩展把注意力里的softmax和GEMM做更深层的融合、针对不同GPU架构再做一轮特化、或者把量化感知训练提前融入kernel接受FP8输入中心……每一项背后都有实实在在的加速空间。只要你手里有跑推理或训练的业务场景就值得花时间把GEMM这一层真正吃透。