DeepGEMM：FP8精度下Tensor Core矩阵乘法的极致优化

我是在折腾某个大模型推理引擎的算子性能时第一次接触到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精度为什么值得专门写一个Kernel

FP16有约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比较小的细长矩阵。

我整理过一个简单的对照，方便根据场景选型：

线程块配置一般以256线程居多，每个线程要管理一小块寄存器累加器。如果每线程的累加器太多，寄存器不够用；太少，则MMA指令数量不够多，延迟隐藏不了。实践中可以从每线程4组累加器起步，配合8 warp的块大小，先用profiler看寄存器数和occupancy，再逐步调整。

4. 性能实录与对比测试

4.1 在旗舰GPU上的实测数据

我在一台配置某款主流数据中心GPU的机器上跑过一组对比测试，矩阵尺寸选M=8192、N=8192、K=8192的规则大矩阵，数据精度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的融合省掉了好几个额外kernel，GPU少了好几次启动和全局内存往返。

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，结果直接写回C，bias也一并融合掉。workspace用于存放per-token scale等临时数据，调用前需要查询一次所需大小，避免运行时动态分配。

编译时只要确保CUDA编译版本够新，支持目标GPU对应的架构即可。整个工程没有额外运行库依赖，放进推理框架的构建系统里很少遇到冲突。这一点对生产环境非常友好。

5.2 常见问题与排查技巧

我把实际调试中遇到过的坑整理成一张表，方便你快速定位：

其中最隐蔽的是第一条。我调试过一个模型的某几层，单测全过，端到端loss却异常，折腾了大半天才发现是scale描述符里per-token数组的stride和kernel内部假设不一致。这类问题不会报错，只会慢慢拉开精度差距，所以接入新模型时一定要先在单层上做全精度对比。

另外一个值得注意的点：如果batch很小，比如decode阶段batch=1，GEMM的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这一层真正吃透。