DeepGEMM高效FP8矩阵乘法内核:MoE与Dense场景大模型推理加速实践 DeepGEMM 这个项目开源之后我把手头的推理加速计划整个推翻重来了一遍。矩阵乘法GEMM是深度学习里承担计算最重的原语大模型训练和推理的每一层几乎都压在它身上。DeepGEMM 用极简的内核实现了 FP8 精度的矩阵乘法尤其针对 MoE混合专家场景做了专门优化实测性能非常有竞争力。对正在做 LLM 推理加速、训练调优或者想啃 GPU 高性能内核源码的人来说这个项目是最近最值得研究的参考实现之一。下面按我的拆解顺序来讲先理清 FP8 和大模型性能的关系再看 DeepGEMM 的代码思路然后给出一套可以直接抄的调用流程最后把我实际踩过的坑整理成清单。1. 核心思路FP8为什么是现代大模型性能的关键大模型的计算密度越来越高大家总说“显存不够”但更准确的原因是矩阵乘法的次数和体量都在涨。FP16 已经成为很多团队的默认精度可权重和中间激活都占在显存里每次矩阵乘法都要读一遍。如果能把精度降到 FP8权重少一半显存一次乘法少一半数据搬运这个收益是实打实的不是只在纸面上成立的“精度减半”。1.1 从FP16到FP8省下的不只是显存FP8 并不是一个统一的格式常见的有两种E4M34位指数、3位尾数动态范围比较大但小数位数少相对精度会低一点。E5M25位指数、2位尾数动态范围更大精度更粗主要用于梯度和误差累加。所以在实际使用 FP8 时通常会有个约定权重和激活用 E4M3梯度或误差用 E5M2。这样能把“动态范围”和“计算精度”分开照顾。可以做一个生活化类比同一块画布既要画远山又要画近景FP8 就像强行压缩灰度级但缩放因子scale能让山和树都落在自己合适的亮度区间里。量化过程可以写成下面这样X_fp8 clamp(round(X_fp32 / scale)) 输出 A_fp8 * B_fp8 * (scale_a * scale_b)这里最关键的一点是GEMM 内部的累加仍然用 FP32不是把中间结果也压在 FP8 里。很多新手以为 FP8 矩阵乘法就是所有数都用 FP8 算真这么做误差会爆炸。DeepGEMM 的设计就是让输入数据以 FP8 进入计算单元累加过程保持 FP32输出前再配合 scale 做反量化这也是高质量 FP8 内核的基本底线。1.2 DeepGEMM的定位Dense与MoE双模式一个 GEMM 内核想被广泛使用至少要覆盖两类实践场景。一类是标准 Dense 模式所有 token 都经过同一批参数常见于普通 Transformer 的线性层。另一类是 MoE 模式每个 token 只会被路由到若干个专家网络矩阵形状不规则。DeepGEMM 同时实现了这两种模式。Dense 模式很好理解输入是规整的 M×N×K内核按普通分块方案跑。MoE 模式就相对复杂不同专家分到的 token 数不同如果为每个专家单独启动一次 GEMM启动开销、显存碎片和调度不均都会冒出来。DeepGEMM 的做法是把多个小 GEMM 按 token 分组合并成一个大 GEMM再用 offsets 信息切分计算。简单说它把同一个 kernel 里的计算单位切分成若干段每个段处理一个专家的小 GEMM一次启动处理所有专家。场景Dense 模式MoE 模式适用情况每个 token 计算全部参数token 路由到少数专家矩阵形态单个规则大 GEMM多个形状不同的小 GEMM调优重点数据复用、scale 粒度分组调度、offsets 维护1.3 极简代码背后的取舍高性能 GEMM 内核在成熟计算库里通常有上万行因为要适配各种 shape、各种指令版本、各种缓存策略。DeepGEMM 只保留了 FP8 这一条主线代码量被压得非常少。它并不是“简单”而是把很多通用性砍掉了。固定 tile 尺寸、只考虑 FP8 输入、专门处理 MoE这些限制换来了更短的代码路径也让手写调优变得可行。这一点对刚上手的人特别友好。如果你之前想读计算库源码却被各种复杂模板和抽象层劝退DeepGEMM 是一个很好的入口。同时也要注意它意味着如果你需要 FP16 或 INT4不能直接拿它就用得看清楚源码结构再改造。2. 关键设计解析如何做到比闭源库更快闭源商业库的 GEMM 已经优化了很多年DeepGEMM 还能有很强的竞争力靠的不是魔法而是几个针对性设计。我把它拆成三层来说数据复用、量化路径、MoE 去碎片化。2.1 分块调度与数据流性能从哪来GPU 的算力远远高于显存带宽所以高效 GEMM 的核心就是最大化数据复用。DeepGEMM 沿用了分块矩阵乘法的思路把大矩阵切成小块让一个小块的数据留在片上寄存器或共享内存里反复使用而不是频繁回全局内存读。每个计算单元负责一个大 tile组内线程再分工到不同的子块。访存顺序和计算顺序是绑在一起的目的是保证加载一次数据后能够算足够多次乘法。这种数据流设计是整个 GEMM 性能的根基。一句话同样是矩阵乘法不同的调度方式可能差出好几倍性能。我自己看内核时会重点关注三个位置从全局内存到共享内存的搬运是否只做一次。累加变量是否长期留在寄存器里。一个 warp 负责哪几个子 tile边界怎么对齐。如果这三个位置都处理好了性能基本不会太差。2.2 FP8量化与反量化精度陷阱如何绕开FP8 的计算单元只支持“FP8 输入、FP32 累加”的路径。但实际流程里上游给过来的往往是 FP32 权重或激活下游也需要 FP32 输出。如果中间单独做量化再单独做反量化就会多出两组额外访存和计算。DeepGEMM 把这套流程直接融进 GEMM 内核从全局内存加载 FP8 后立刻配合 scale 参与计算输出写回前再乘回缩放因子。一个容易踩的坑是 scale 的粒度。如果所有数据共用一个 scale量化误差可能非常大。DeepGEMM 支持更细的按行、按列、按 token 分组 scale代价是要额外搬一点 scale 数据。实际项目里我建议先按行做 scale跑一轮看精度不满意再细分。不要一开始就上最细粒度的方案否则带宽开销会把性能优势抵消掉。从工程角度看量化本身也分“训练前静态量化”和“推理时动态量化”。静态量化需要先统计一批真实数据动态量化每步都要算 scale开销更高。DeepGEMM 主要解决的是计算内核问题但你怎么确定 scale、怎么更新 scale仍然是自己的工程工作。2.3 针对MoE的分组去碎片化MoE 模型推理时路由结果每个 step 都可能变。专家权重可以预先量化好但 token 到专家的分配是动态的。最朴素的做法是循环每个专家每次调用一次 GEMM。当专家数量多、单个 GEMM 又很小时调用开销甚至能超过计算本身。DeepGEMM 的 MoE 模式会用一次 kernel 调用处理多个专家通过预排序把属于同一个专家的 token 聚在一起再在 kernel 内部按分段计算。这样既减少了启动次数也减少了因为 padding 浪费的显存和算力。这里可以算一笔账假设有 64 个专家每个专家平均分到 64 个 token逐个启动 64 次小型 GEMM光 launch 开销就可能占据几微秒甚至十几微秒。如果合并成一个 grouped GEMM启动次数变成一次计算效率会明显改善。这也是 MoE 推理和 Dense 推理在性能调优上最大的差异。3. 实操把DeepGEMM集成进你的推理流程这部分我按自己的使用顺序来写。不同版本的 API 可能有差异但整体思路是一致的。你最好把它当成“参考路线”而不是一字不差的官方文档。3.1 环境准备与编译首先确认自己的 GPU 计算能力是否支持 FP8 矩阵运算。一般来说要选择支持 FP8 数据格式的数据中心级 GPU。开发环境需要 GPU 编译工具链和一个支持较新 C 标准的编译器。把源码拉下来之后按仓库说明同步子模块然后编译git clone DeepGEMM 项目地址 cd DeepGEMM mkdir build cd build cmake .. -DCMAKE_BUILD_TYPERelease make -j编译时如果遇到找不到头文件通常不是 CUDA 版本问题而是依赖子模块没有拉全。先同步子模块重新配置再编译。这里有个小经验编译用 Release 模式不要用 Debug否则性能测试的数据没有任何参考价值。3.2 用Dense模式做一次FP8矩阵乘法写一个最小程序把 FP32 的矩阵量化到 FP8再调用 DeepGEMM#include deep_gemm/deep_gemm.h #include vector #include cmath int main() { constexpr int M 4096; constexpr int N 4096; constexpr int K 4096; // 模拟输入实际场景里这里可能是权重和激活 std::vectorfloat A_fp32(M * K); std::vectorfloat B_fp32(N * K); // 按行量化 std::vectorunsigned char A_fp8(M * K); std::vectorunsigned char B_fp8(N * K); std::vectorfloat scale_a(M); std::vectorfloat scale_b(N); for (int i 0; i M; i) { float max_abs 0.0f; for (int k 0; k K; k) { max_abs std::max(max_abs, std::fabs(A_fp32[i * K k])); } // E4M3 的最大值约 448用它做分母可以充分利用表示范围 scale_a[i] max_abs / 448.0f; for (int k 0; k K; k) { float val A_fp32[i * K k] / scale_a[i]; // 实际实现中要 clamp 到 [-448, 448] A_fp8[i * K k] static_castunsigned char(val); } } auto C deep_gemm::fp8_gemmDense( A_fp8, B_fp8, scale_a, scale_b, M, N, K); // C 是 FP32 结果建议和普通 FP32 GEMM 对拍 return 0; }代码里我特意省略了 B 的量化过程因为逻辑和 A 一样。跑通之后一定要先做一步验证随机初始化一组小矩阵同时用 FP32 实现和 DeepGEMM 各算一遍比较最大绝对误差。误差在可接受范围内再谈性能。3.3 用MoE模式跑混合专家路由MoE 模式比 Dense 麻烦一些核心是要先把 token 按 expert_id 排序并且生成 offsets 数组。// 假设有 num_tokens 个 tokennum_experts 个专家 std::vectorint expert_ids(num_tokens); // 1. 根据 expert_ids 排序得到 token_order std::vectorint token_order sort_tokens_by_expert(expert_ids); // 2. 统计每个专家分到的 token 数生成 offsets std::vectorsize_t offsets(num_experts 1, 0); for (int tid : token_order) { offsets[expert_ids[tid] 1]; } for (int e 0; e num_experts; e) { offsets[e 1] offsets[e]; } // 3. 重新排列 token得到一个按专家连续排列的大矩阵 std::vectorunsigned char tokens_fp8(num_tokens * hidden_size); // ... 按 token_order 填入 std::vectorunsigned char experts_fp8(num_experts * hidden_size * expert_hidden); // ... 专家权重通常可以预先量化 auto out deep_gemm::fp8_grouped_gemmMoE( tokens_fp8, experts_fp8, offsets, scale_token, scale_expert);注意 offsets 的形状是[0, count_0, count_0 count_1, ...]。它的作用是告诉 kernel每个专家对应的 token 范围从哪开始、到哪结束。排序这一步虽然多了一次数据重排的代价但和反复启动 kernel 相比通常还是划算的。如果你不想自己重排也可以在 kernel 内部处理但那样内核的复杂度会明显上升我更推荐在外部做好。3.4 性能对比与调参建议集成之后除了看跑通还要用 profiling 工具确认 kernel 实际耗时。切忌只盯全局端到端时间因为如果数据拷贝和启动开销占比过大内核本身的优化会被掩盖。我一般会做一组 batch size 从 1 到 512 的曲线记录两个指标kernel 时间单位毫秒。相对 FP32 baseline 的加速比。调参优先级上先看 tile size再看是否开启 MoE 分组最后才调 scale 粒度。如果一个固定 shape 下性能始终不理想大概率是 tile 尺寸不匹配。你也可以试着改共享内存分配策略但这些需要重新编译反复试才能找到最优组合。4. 常见问题与排查实录这部分记录的坑很多都是我自己踩过的。有些问题看起来是 DeepGEMM 的 bug实际上往往是我们自己的数据组织方式有误。4.1 输出NaN/Inf的两种最常见原因我第一次跑 Dense 模式时输出一大片 NaN。后来检查发现量化代码里某一行数据的 max_abs 是 0除以 0 导致 scale 异常。另一种常见原因是 E4M3 的范围太小输入里有个特别大的离群值量化后直接溢出反量化之后就变成 Inf。排查思路比较简单打印每一行 scale 的统计信息确认没有 0。检查输入数据的绝对值范围是不是远大于 FP8 可表示的范围。在量化代码里加 clamp强制把超出范围的值截断。如果解决不了就把 scale 粒度从 per-tensor 改成 per-row。通常 per-row 已经很稳per-token 更多用于激活值分布特别不均匀的场景。4.2 小矩阵场景性能反而变差有人实际测试以后发现小尺寸 GEMM 用 FP8 并没有比 FP16 快多少。问题往往不在内核本身而在于 kernel launch 和从主机到设备的拷贝开销占了大部分时间。DeepGEMM 的 MoE 模式优势正好体现在这里把多个小 GEMM 合并成一次调用启动开销被摊薄。如果你的场景就是单个小矩阵建议考虑批处理打包或者干脆用 Dense 模式跑不要寄希望于换成 FP8 就自动加速。性能调优时还有一个容易忽视的点端到端时间里的 CPU 同步。如果每次调用都强制 GPU 和 CPU 同步流水线会被打断性能数据完全失真。正确做法是让多个 kernel 在同一个 stream 上排队执行只在最终验证结果时才做同步。4.3 硬件不支持FP8的备选方案不是所有 GPU 都支持 FP8 矩阵运算。如果硬件不支持DeepGEMM 的 FP8 路径就不成立。这时候要么换硬件要么退回 FP16 或 INT8。FP8 的优势主要体现在显存减半和计算密度提升如果硬件不支持强行模拟 FP8 反而更慢。我的建议是在项目早期就做一次硬件能力自查。如果目标是线上推理还要看推理服务器的 GPU 是否统一支持 FP8。否则代码调好了部署到旧设备上就只能降级前期的性能收益就没了。4.4 问题速查表现象可能原因排查思路输出 NaNscale 为 0 / 量化除零检查每行 max_abs 统计输出 InfFP8 溢出收缩 scale 或加 clamp精度不达标scale 粒度太粗改用 per-row / per-token性能低于预期小矩阵或 launch 开销占比高profiling 看 kernel 时间增大 batch编译失败子模块缺失同步子模块后重新 cmake只支持特定硬件GPU 架构不支持 FP8换硬件或改精度5. 把DeepGEMM用进生产之前我建议你先想清楚这几件事内核跑通只是一小步真正有挑战的是让它在自己的项目里稳定落地。我见过不少团队把新算子集成进去以后测试集上指标没问题一到线上就出现偶发异常。这类问题一般不在 GEMM 内核本身而在工程链路。5.1 精度验收标准不能只看错误率FP8 带来的精度损失在大多数情况下是可接受的但“大多数”不代表“所有”。如果你的模型对数值特别敏感比如用了某些归一化层或落到了长尾分布上FP8 的量化噪声可能被放大。上线前不要只看整体 loss 或准确率最好按模块单独对拍。如果某一层的输出相对误差超过预期优先调整该层的 scale 粒度而不是全链路一起改。我自己的经验是先跑一个 100 条样本的验证集比较 FP32 和 FP8 的中间激活分布。如果看到离群点再决定是否对特定层做混合精度。混合精度不是复杂的事情就是在关键层保留 FP16其余层用 FP8成本很低但能解决很多偶发问题。5.2 版本锁定与回归测试要提前做DeepGEMM 还在快速迭代API 和性能特征都可能在版本之间变化。引入之前把当前版本的源码、编译参数、测试结果全部固定下来。我通常会写一个简单的回归脚本跑一组固定 shape 和参考输出确保每次升级都能快速发现问题。性能回归也一样。上次性能好不代表这次性能好因为编译器版本、GPU 驱动、环境变量都可能影响结果。性能对比必须在同一套环境下做记录完整环境信息而不是只拿一个 TFLOPS 数字四处宣传。5.3 不要为了追性能丢了可维护性DeepGEMM 的极简实现很优雅但它不是银弹。如果项目里现有的推理框架已经封装好了一套动态 shape 和自动调参逻辑直接替换内核可能会破坏原有优雅性。你可以先以“插拔式”的方式接入对外保留原来的接口内部根据 shape 选择走 FP8 还是旧路径。这样风险最小也方便灰度。我最后想说的是读高性能内核源码比单纯用黑盒库有营养得多。你不需要把 FP8 当成什么玄学理解量化、缩放、数据流之后它在生产环境里很可控。所有 benchmark 都只是特定环境下的快照别人的快慢不代表你的快慢一定要在自己的模型和数据上复测再决定要不要引入。