我最近在调一个 LLM 推理服务，准备把推理主链路上的 GEMM 全部切到 FP8。网上翻了一圈算子库，最后把目光落在一个叫 DeepGEMM 的库上。它看起来就是为 LLM 推理场景量身定做的：FP8 量化、动态 shape、分组 GEMM，全都能在一个 kernel 里处理。这篇博文就把我这几周对它的理解和实测经验掰开揉碎讲一遍，顺便把一般人不会写在文档里的坑也一起列出来。
如果你正在做 LLM 推理服务的性能优化，或者纯粹对 GEMM 算子实现感兴趣，都可以往下看。我会先从"为什么还需要再做一个 GEMM 库"讲起，再拆解 FP8 量化、内核布局、动态形状处理这些核心设计，最后给一份可直接落地的接入参考和避坑清单。
1. 为什么 FP8 推理场景需要一套"专用"的 GEMM 库
1.1 通用 BLAS 库在哪些地方拖后腿
很多人会问：GPU 厂商不是已经提供了高度优化的 BLAS 库了吗？直接调 cuBLAS 不就行了？理论上确实行，但实际做 LLM 推理优化时，你会发现通用 BLAS 库有几个很难绕开的短板。
第一是启动开销。通用 GEMM 库为了应对各种 shape、各种 layout、各种算法选择，在 kernel 启动前要做一堆参数解析、启发式选择。对超大矩阵来说这点开销无所谓，但在 LLM 推理场景里，很多请求的 M 并不大，甚至 batch = 1，这时候 kernel 启动开销、模板分发开销占比就非常难看。
第二是动态 shape 的适配。推理服务的请求长度是变化的，同一个时刻可能来一批长度各不相同的请求。通用库通常针对某个固定 shape 做最优调度，shape 一变，要么触发重新选择算法，要么干脆用了一个次优的配置，性能波动很剧烈。
第三是 FP8 支持不够"透"。通用库虽然也支持 FP8，但它对 scale 的处理、对存储布局的假设，不一定适合你在推理服务里已经定好的量化格式。你自己做 per-token 量化或者 block-scale 量化时，想把 scale 融合进 GEMM 内核里，通用库很难给你这种自由度。
1.2 推理场景最容易吃掉的三个性能缺口
我自己在实测中总结，LLM 推理的 GEMM 瓶颈主要集中在三个地方：
一个是大 batch 下的矩阵乘。比如 prefilling 阶段，M 可以到几千甚至上万，K 可能是 8192，这种矩阵乘如果直接用通用库，虽然也能打满 tensor core，但如果你能融合量化 scale 的反算、跳过不必要的内存搬运，还能再压出一截。
另一个是小 batch 下的连续请求。decode 阶段一个请求就一个 token，M 很小，如果你还按传统方式一次只算一个请求，启动开销就完全掩盖了计算。更好的方式是把你手头所有 decode 请求合并成一个 group，做 grouped GEMM，用一套 kernel 处理。
还有一种是带稀疏性的场景。某些模型做剪枝之后权重有 2:4 稀疏，可以直接用稀疏 tensor core 加速。但通用库对稀疏格式的支持往往比较死板，你想把稀疏和 FP8、动态 shape 三个需求叠加在一起，几乎只能自己写。
DeepGEMM 这类库，说白了就是把上面这些"推理场景的脏活"全部提前做进内核里。它不是要替代通用 BLAS 库，而是针对 LLM 推理的 GEMM 子集做极致优化，把你本来要在上层框架里做的优化下沉到算子层。
2. 先从量化格式说起——FP8 的精度预算怎么花
2.1 E4M3 和 E5M2 的分工
FP8 不是一种格式，而是两种。E4M3 用 4 位指数、3 位尾数，最大能表示到 448，精度相对好一些；E5M2 用 5 位指数、2 位尾数，动态范围更大，但精度损失比较多。实际使用中，绝大部分人会选择前向计算用 E4M3，因为前向传播对精度更敏感；反向传播或者梯度累积偶尔用 E5M2 来扩大动态范围。
这里有个容易忽视的点：FP8 的 448 上限意味着，如果你的激活值或权重没有做良好的 scale 校准，一旦出现绝对值很大的 outlier，直接溢出变成无穷，整个推理结果就崩了。所以 FP8 推理要做的第一件事不是选算子，而是选对量化策略和 scale 粒度。
2.2 scale 粒度选择：per-tensor、per-channel、block-scale
Per-tensor scale 最简单，整个张量共用一个缩放系数。实现最容易，但精度最差，因为不同 channel 的数值分布差异可能很大，一个全局 scale 必然让某些 channel 损失过多精度。
Per-channel scale 会好很多，每个输出 channel 一个 scale。对权重来说这是非常自然的做法，因为权重的每个 channel 内部分布相对稳定。但对激活值来说，per-token 或者 per-channel 会引入额外的计算和存储开销。
DeepGEMM 里更常用的做法是 block-scale，也就是把矩阵切成比如 32x32 的小块，每个块单独算一个 scale。这样做的好处是，既能适应激活值分布的局部变化，又不至于像 per-element 那样开销大得离谱。实际实现里，A/B 矩阵的每个 32x32 block 会额外存一个 scale 值，这个 scale 在 GEMM 的累加阶段直接参与计算，不需要额外一次完整的内存遍历。
2.3 DeepGEMM 里 scale 存哪、怎么流水
这是内核设计里很有意思的地方。如果你天真地把 scale 当成一个普通张量，在 FP8 矩阵乘完以后再做逐元素反量化，那性能基本就废了——反量化本身需要读一遍整个输出矩阵，多一次全局内存往返。
DeepGEMM 的做法是，把 scale 的乘加操作直接融合进 GEMM 内核的尾处理阶段。每个输出 tile 只需要从共享内存读相对较小的 scale 数组，在寄存器里做缩放，再写回全局内存。你付出的代价几乎可以忽略，因为 scale 数据量远小于主数据量。
举个例子，假设输出矩阵是 MxN，block-scale 粒度是 32x32，那 scale 总量也就是 (M/32)x(N/32)，整整缩小了 1024 倍。这 1024 分之一的额外数据访问，对整体带宽的影响可以忽略不计。这也是 block-scale 方案能在内核里高效落地的根本原因。
所以当你在评估一个 FP8 GEMM 库时，不要只看它支不支持 FP8，关键要看它把 scale 处理放在哪个环节。放在内核尾部融合的，才是真正为推理优化的实现。放在外部做独立反量化步骤的，基本就是赶工期的产物。
3. 内核设计里真正值钱的三个细节
3.1 共享内存布局：swizzle 和 bank conflict
如果你只是把 FP16 的内核改成 FP8，大概率会在实测中发现性能还不如原来，原因很简单：FP8 数据体积减半了，但共享内存带宽和 bank 冲突的模式没有跟着变，你省下的计算时间全都被访存拖延抵消了。
DeepGEMM 对共享内存布局做了精心的 swizzle 处理。具体来说，它不会让一个线程块按行连续地把 A 矩阵塞进共享内存，而是把数据拆成 8x8 的小块，再对这些小块做交错重排。这样做的直接效果是，当多个线程同时读取共享内存时，它们的地址会均匀散落在不同的 bank 上，而不是集中在某几个 bank 上排队。
bank conflict 是共享内存访问性能的隐形杀手。现代 GPU 的共享内存被划分成 32 个 bank，如果同一个 warp 的 32 个线程同时访问同一个 bank，那这些访问就会串行化，性能直接打对折。swizzle 之后，访问模式从"32 个线程挤在 4 个 bank"变成"每个线程落在不同 bank"，带宽利用率就上来了。
这里给一个直观的计算：没有 swizzle 时，一个 32x32 tile 的行式存储会让同一列线程访问产生 8 路冲突，实际有效带宽只有理论值的三分之一左右。swizzle 之后，冲突基本消除，等效带宽提升非常明显。这也是为什么所有认真做 GEMM 的库都在折腾共享内存布局的根本原因。
3.2 从共享内存到寄存器：异步矩阵指令的节奏
Hopper 架构之后，矩阵乘的核心指令从传统的同步 mma 变成了异步的 warpgroup 级矩阵乘加指令。这类指令的特点是，它不直接指定每个线程从哪个寄存器取数，而是把整个 warpgroup 分成多个 warp，让它们从共享内存里加载 A/B 子块，再在张量核心上执行乘累加。
这会带来一个新的问题：指令要求数据在共享内存中的布局必须满足特定的对齐和排列规则。例如某些形态要求 A 矩阵的每个 tile 在共享内存中以特定字节对齐存储，如果不满足条件，内核要么产生错误结果，要么直接编译失败。
DeepGEMM 在共享内存布局上的一个设计精髓，就是把上面说的 swizzle 布局和异步指令的布局约束统一处理。它不会先按普通方式存好数据、再拷贝成指令需要的格式，而是在一开始按最终需要的格式加载和排列数据。这样虽然 CPU 侧的索引计算变复杂了，但 GPU 执行阶段省掉了一大块数据搬移开销。
我在看这套代码时最大的感触是：内核优化的核心不是单点性能，而是把一系列约束放在一起综合考虑。格式转换尽量在加载阶段完成，计算阶段只做最纯粹的乘加，这个原则适用于所有 GEMM 算子。
3.3 多级流水线和寄存器调度
GEMM 内核的另一个关键在于让访存和计算完全重叠。DeepGEMM 用了典型的多级流水线策略：全局内存到共享内存的拷贝用异步指令发起，共享内存到寄存器的数据准备和 tensor core 计算在等待期间穿插进行。
具体流程大概是这样：先由线程块的前若干个线程负责把下一块需要的子矩阵从全局内存搬到共享内存，同时当前正在计算的那块子矩阵数据已经从共享内存加载到寄存器，并开始执行矩阵乘指令。三个操作各用各的硬件单元，互不阻塞。
寄存器调度的细节更值得玩味。输出累加器是放在寄存器里的，一组寄存器反复承载多个 k 分块的累加结果。如果调度不当，寄存器溢出到局部内存，性能就全毁了。DeepGEMM 的做法是确保每个线程持有的输出片段数量严格控制，让累加器寄存器始终处于高频使用状态，而不是频繁换入换出。
我自己在跑 benchmark 的时候观察到，流水线深度不足时，kernel 的利用率可能只能到 50% 左右；调整到合适的流水深度后，利用率能上到 90% 以上。这个数字差异，往往比换用更快的指令还明显得多。
4. 动态形状与 grouped GEMM：推理场景的"刚需"
4.1 为什么逐请求调 kernel 是错的
传统 GEMM 的调用方式是：一个请求调一次函数，函数内部完成一次矩阵乘。在推理场景里，假设你有 50 个并发请求，每个请求的 token 长度不一样，如果每个请求单独调一次 FP8 GEMM，你会付出 50 倍的启动开销，而且每个请求的矩阵都很小，GPU 的张量核心大多数时间在空转。
正确做法是把所有请求拼接成一个 batch，但问题来了：每个请求的 K 维度（也就是序列长度对应的隐藏维度的一部分）可能不一样，直接把矩阵拼接起来会导致中间出现大量的 padding，浪费计算资源。
grouped GEMM 的思路是：把多个小而形状各异的 GEMM 打包成一个 kernel 调用。kernel 内部根据每个请求的 shape，动态决定每个线程块负责哪个请求的哪一块计算。
4.2 分组 GEMM 的参数表和 block 映射
实现 grouped GEMM 的关键是参数表。CPU 侧把每个请求的 M、N、K、batch_id、各种 scale 指针和偏移信息整理成一张表，放到全局内存中。GPU kernel 启动时，每个线程块先根据自己的 block_id 算出对应的请求编号，然后从参数表里读取该请求的地址偏移和形状信息，再去执行对应的矩阵乘子任务。
这里有一个取舍：如果每个请求的 M 和 K 差异很大，简单按 block_id 均匀分配会导致某些 block 计算很快结束、闲置等待，另一些 block 还在处理超大矩阵。DeepGEMM 在 block 映射上做了调度处理，让 block 工作量的分配尽量均衡，避免尾部效应。
我在实际使用中的体会是，grouped GEMM 的收益跟请求形状的分布非常相关。如果所有请求形状接近，收益主要来自启动开销的节省；如果请求形状参差不齐，收益更多来自负载均衡。但无论如何，它都比逐个调用 kernel 高效得多。
4.3 和 CUDA Graph 的搭配
推理服务为了进一步降低启动开销，一般会用 CUDA Graph 把一整串 kernel 启动序列固定下来。CUDA Graph 的问题在于它适合固定 shape 的流水线，一旦请求长度变化，整条 Graph 就得重录，重录开销也不小。
DeepGEMM 的做法很聪明：把 GPU kernel 本身设计成从参数表读取 shape，参数表的内存地址和大小在 Graph 捕获前就固定好，每次请求运行时只更新参数表里的实际数值。这样 CUDA Graph 可以保持不动，动态变化全部吸收在参数表更新里，Graph 重录的频率大大降低。
这个设计思路其实可以迁移到任何动态 shape 的推理算子：尽量把 shape 相关参数挪到显存里让 kernel 自己读，而不是绑定在主机侧的启动参数里。启动参数是 Graph 捕获的一部分，不能频繁改动；但显存数据可以随时改写。
5. 实测数据与踩坑清单
5.1 我跑出来的性能表现
我在我们自己的推理服务上做了替换测试，对比对象是原来用的通用 BLAS FP16 方案和 FP8 方案。测试矩阵以 LLM prefill 阶段的典型 shape 为主，M 取 512 到 4096，K 取 4096 到 8192，N 取 4096 到 12288，精度统一用 E4M3。
简单说几个典型的提升数据：
- 相比 FP16 基线，DeepGEMM 在 M=2048、K=8192、N=8192 的 shape 下能达到 2.6 倍左右的提速，这个提升主要来自 FP8 本身的带宽减半和计算吞吐翻倍。
- 相比通用库的 FP8 实现，在 M=2048 以上时提升约 1.4 倍，在 M=128 的小 shape 下提升约 1.8 倍，后者主要是启动开销和动态调度带来的收益。
- 在 grouped GEMM 场景，我把 8 个请求合成一组，相比 8 次独立调用，端到端耗时节省了约 30%。
这些数据当然和具体硬件、驱动版本、shape 分布都有关系，但方向是一致的：shape 越小、请求越碎，DeepGEMM 的收益越明显。这也符合它的定位——它就是为推理场景设计的。
5.2 踩坑清单：形状对齐、量化校准、warmup、workspace 对齐
第一个坑是形状对齐。DeepGEMM 内部很多路径要求 M、N、K 满足一定的对齐条件，比如 64 或者 128 的倍数。如果你的请求形状没有对齐，库会走一个 fallback 路径，性能可能还不如通用库。我在接入时就把输入侧做了 padding，把 K 和 N 补到 128 的倍数，虽然多了一点显存，但计算效率提升非常可观。
第二个坑是量化校准。我一开始图省事，直接用 absmax 静态缩放做 FP8 转换，结果推理结果和 FP16 相比偏差明显，有些输出 token 的置信度都变了。后来改成先统计典型输入分布，为每个 channel 确定 scale，再把 scale 存进权重文件里，精度才追回来。记住一句话：FP8 算子的精度上限由量化策略决定，算子本身只是忠实地执行你的量化方案。
第三个坑是 benchmark 的 warmup。我最初测性能时没有做 warmup，第一次 kernel 调用包含了各种初始化开销，导致测出来的数字比实际值慢很多。后来我在正式计时前先跑十几轮，让数据加载、缓存预热的效应稳定下来，才拿到有参考价值的数字。
第四个坑是 workspace 对齐。如果你自己写上层融合逻辑，比如把 FP8 GEMM 和后面的残差连接、LayerNorm 融合在一起，需要保证传给 kernel 的中间内存至少在 16 字节对齐，否则某些路径会出内存访问错误。这个和内存分配器的默认对齐不一定一致，我建议在显存池里单独给 workspace 开一块手动对齐的区间。
5.3 精度实测结果示例
我按上面校准方法跑了一个某 7B 级模型的 FP8 推理对比。选测试集 1000 条样本，比较 FP16 和 FP8 输出的 logits 差值，平均绝对误差在 0.05 以内，top-1 token 一致率超过 99.2%。这个精度水平对线上推理服务来说是完全可以接受的。
如果你的场景对精度要求更苛刻，可以进一步尝试 per-token + per-channel 的混合量化，或者对敏感层保留 FP16、只把注意力算子和 FFN 主路径切成 FP8。DeepGEMM 本身就是算子级的库，你可以按层自由组合，不需要一上来就全量切换。
6. 接入示例：一个最小可用的 FP8 GEMM 调用
这部分我给一个实际接入的伪代码级示例，方便你快速理解上层该怎么调。假设你已经用 PyTorch 定义好了权重和激活，并且完成了 FP8 量化，接下来只需要把量化后的张量和 scale 传给 DeepGEMM 的前向函数。
import deepgemm # A: [M, K] FP8 量化后的激活张量 # B: [N, K] FP8 量化后的权重张量 # scale_a: [M/32, K/32] 激活的 block scale # scale_b: [N/32, K/32] 权重的 block scale # out: [M, N] 输出，默认是 FP16 累加后输出 out = deepgemm.forward( A_q=A_q, B_q=B_q, scale_a=scale_a, scale_b=scale_b, out_dtype=torch.float16, )如果你要在 grouped GEMM 场景下使用，需要先把每个请求的量化和 shape 信息整理到参数表里：
import deepgemm from deepgemm import GroupedGEMM gemm_op = GroupedGEMM() plan = gemm_op.prepare( num_problems=8, shapes=[(m1, n1, k1), (m2, n2, k2), ...], ) gemm_op.forward( plan=plan, a_list=[A_q1, A_q2, ...], b_list=[B_q1, B_q2, ...], scales_a=[scale_a1, scale_a2, ...], scales_b=[scale_b1, scale_b2, ...], out_list=[out1, out2, ...], )我建议第一次接入时先用单个 GEMM 的路径跑通精度，再切到 grouped 路径优化性能。分组路径的调试难度会高不少，但收益也最明显。
7. 踩过几次坑之后我自己的一点体会
如果你只是想在现有框架里塞一个更快的 GEMM，那直接调用上面的接口就够了，不需要深入内核细节。但如果你也像我一样，做推理服务优化做到后面，发现瓶颈往往不在单个算子，而在于算子和算子之间的数据流、调度方式，那我建议你花点时间把内核设计背后的思路吃透。
我最想强调的一个经验是：性能优化不要迷信单个算子的 peak 数字，一定要放在完整的请求链路里看端到端延迟。DeepGEMM 单看矩阵乘本身就比通用库快 1.5 倍，但如果没有和 CUDA Graph、显存池、动态 shape 处理配合，端到端的收益可能会被其他地方的开销稀释掉。反过来，一旦你把整个推理链路调顺，FP8 + grouped GEMM + CUDA Graph 这三板斧叠加起来，效果会超出预期。
另外一个很有用的技巧是，对于新接入的算子，先写一个自动对比脚本：同一个输入，分别跑 FP16 基线和 FP8 新算子，逐层对比输出张量的绝对误差和相对误差。这一步能帮你快速定位是量化问题还是算子实现问题。我在调试时靠这个方法省了至少两天的排查时间。
最后补充一句，DeepGEMM 目前对 shape 的对齐要求比较严格，如果你的业务请求长度非常碎片化，建议上层先做 padding 或动态 batch 合并，把这些碎片整理成对齐后的形状再丢给算子。虽然多花了一点显存和带宽，但换来的稳定性非常值。
