★ FEATURED ARTICLE 我最近在调一个 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 阶段一个请求就一个 tokenM 很小如果你还按传统方式一次只算一个请求启动开销就完全掩盖了计算。更好的方式是把你手头所有 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-scalePer-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 数据量远小于主数据量。举个例子假设输出矩阵是 MxNblock-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 到 4096K 取 4096 到 8192N 取 4096 到 12288精度统一用 E4M3。简单说几个典型的提升数据相比 FP16 基线DeepGEMM 在 M2048、K8192、N8192 的 shape 下能达到 2.6 倍左右的提速这个提升主要来自 FP8 本身的带宽减半和计算吞吐翻倍。相比通用库的 FP8 实现在 M2048 以上时提升约 1.4 倍在 M128 的小 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_qA_q, B_qB_q, scale_ascale_a, scale_bscale_b, out_dtypetorch.float16, )如果你要在 grouped GEMM 场景下使用需要先把每个请求的量化和 shape 信息整理到参数表里import deepgemm from deepgemm import GroupedGEMM gemm_op GroupedGEMM() plan gemm_op.prepare( num_problems8, shapes[(m1, n1, k1), (m2, n2, k2), ...], ) gemm_op.forward( planplan, 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 合并把这些碎片整理成对齐后的形状再丢给算子。虽然多花了一点显存和带宽但换来的稳定性非常值。
阅读完成 · 觉得有帮助？
