第一次在H100上碰FP8 GEMM的时候我差点被一整套奇怪的数值结果劝退矩阵乘出来不是全0就是NaN偶尔正常一次结果又对不上参考实现。后来把关注点从“FP8就是少了一半精度”转到“更底层的数据管理和流水线调度”上才慢慢理解像DeepGEMM这类库为什么值得拆开看。DeepGEMM是一个面向FP8精度的轻量级GEMM库专门处理稠密矩阵乘和MoE分段矩阵乘代码量不大却把从数值缩放到Hopper架构底层调度的关键点全部串了起来。如果你在做推理加速、训练框架底层开发或者只是想搞懂FP8内核背后的原理这篇东西应该能给你省不少时间。1. 从FP16到FP8它卡住了不少团队1.1 硬件算力翻倍但软件没跟上GPU的浮点演进其实和写代码的懒惰程度是矛盾的。老框架里到处是FP16的假设算子注册表、自动混合精度策略、自定义CUDA kernel的模板参数全都围着FP16转。换FP8不是把某个dtype字符串改掉就完事几乎每个环节都要重新验证。但硬件层面的诱惑实在太大。Hopper架构里FP8的峰值算力直接是FP16的两倍内存带宽消耗也少一半。对大模型推理来说这等同于在相同显存和功耗下塞进更多请求对训练来说则代表更大的batch和更短的迭代时间。于是问题变成既然硬件给得这么爽快为什么很多团队迟迟没切换答案很简单软件没有跟上。FP8 GEMM不是“把float换成8位”这么朴素。乘积结果的范围、累加过程中的精度、缩放因子的计算位置每一样都会改变内核的数据流。很多人上手之后发现要么精度崩掉要么性能还不如FP16最后只能退回老方案。DeepGEMM这类库的价值恰恰在于把这些问题在库内部一次性解决好让上层使用者不必重造轮子。1.2 大模型GEMM的两个“反标准”形状传统GEMM库的设计目标是规整的大矩阵乘法M、N、K都很大用分块策略把SM喂得饱饱的reduction维度尽可能长。但大模型场景里的GEMM形状往往并不规整。第一个典型场景是推理生成阶段。M经常只有1到16K和N可以到几千甚至上万。这种细长形状会让标准tile划分产生大量空转每个block算不了多少活就结束了SM的利用率被严重拉低。哪怕你有很高的FP8峰值算力实际吞吐也会非常惨。第二个典型场景是MoE。多个专家权重共享一份激活输入每个专家分到的token数量不固定如果像常规做法那样逐个专家单独调GEMMkernel launch次数会爆炸单个小GEMM又凑不满一个完整waveSM不断闲置。DeepGEMM把这一类问题抽象成“分组/分段GEMM”用单次内核调用处理多个专家的计算目的是让本来细碎的工作变得足够稠密。这个方向恰恰是很多通用库没有为FP8做精细优化的死角。1.3 “小而精”的设计哲学很关键我最初看到DeepGEMM源码时第一反应是核心逻辑怎么这么短它并没有把cuBLAS几千行内部逻辑重写一遍而是围绕FP8 GEMM的关键路径做了精准裁剪。对于库开发者来说代码短不代表功能弱反而说明每一行都有明确目的便于阅读、审计和集成。这也是我认为它值得学习的原因之一。很多高性能库动辄上万行内部优化堆叠得密密麻麻你根本分不清哪些是架构约束、哪些是历史包袱。DeepGEMM把优化集中在“FP8数据布局”“缩放因子处理”“Hopper异步流水线”“MoE分段调度”这几个真正影响性能的点上剩下的保持朴素。这种极致的取舍对想自己动手改内核的人来说是很好的教材。2. FP8数值细节决定整个库的设计方向2.1 e4m3、e5m2怎么选FP8并不是一种格式而是两种e4m3和e5m2。它们的差异集中在指数位和尾数位的分配上。我整理了一张常用的对照表格式指数位尾数位最大有限值适用场景FP16510约65504通用训练/推理BF1687约3.39e38训练中间计算FP8 e4m343约448激活、权重FP8 e5m252约57344梯度、误差较大数据e4m3因为尾数只有3位精度很粗糙但动态范围在4位指数下也能覆盖常见的激活分布e5m2的指数多一位能表达更大的数值但尾数更少精度进一步下降。工程里最常见的组合是激活和权重用e4m3梯度、误差或者需要较大动态范围的数据用e5m2。这里有个很容易踩的坑不要想当然地觉得所有张量都适合e4m3。如果数值分布本身很宽比如某些层的中间激活偶尔会冒出一个明显偏大的channele4m3直接就会溢出到NaN。真正可靠的方案都是配合缩放因子做范围压缩而不是裸着硬上。2.2 scaling不是乘个系数那么简单FP8最容易被低估的部分是scaling。刚接触时我也以为就是给每个张量乘一个全局系数把最大绝对值限制在448以内。实际做进去才发现per-tensor整张量一个scale在真实模型里经常不够用因为激活分布并不均匀一个离群特征会拉低全图的有效精度。工程上更常见的是per-token结合per-channel的组合策略对激活A按token维度分组对权重B按channel维度分组甚至把K维度切块后做更细粒度的分组量化。这样虽然会多出一些scale计算的额外开销但能把量化误差控制在可接受范围。用个生活化的类比per-tensor像用一把长尺量全班身高班里有个两米多的体育生你为了不爆尺只能把刻度拉得很宽结果普通同学的身高误差全变大per-token / per-channel相当于分组测量名正言顺地按人群特征用不同的尺整体误差自然更低。DeepGEMM在设计上支持这类复杂的scale组合把缩放因子的传递和矩阵乘本身的流水线融合在一起。也就是说scale不是前处理阶段单独算好再传进来而是内核内部就处理掉了一部分缩放逻辑避免额外读写和kernel launch带来的开销。2.3 累加精度与延迟缩放FP32寄存器里的故事FP8能承受的精度有限但GEMM有一个天然优势乘法可以用8位累加却不需要停在8位。标准的做法是内部累加器用FP32整个K维的乘积一路累加完最后写回时才落到目标输出精度。这样即使输入是8位累加误差也不会在K很大时失控。另一个关键策略是延迟缩放。你可以把它理解成“用上一次统计好的scale先跑今天的计算”。训练场景里激活分布短时间不会突变没必要每一步都重新算一遍scale适当沿用上一轮的统计值就能稳住精度同时省掉一次额外的数值统计kernel。在线缩放则适合动态范围变化很剧烈的场景两者各有取舍。把这些细节串在一起你就会发现FP8 GEMM的性能高低不只是看乘法和加法本身还得看scale如何参与流水线、累加器放在哪一级、写回时有没有多余的格式转换。DeepGEMM在数值层面的每一个设计最终都是为了让这些步骤尽量不打断内存与计算的重叠。3. 双路径实现Triton把功能跑通CUDA把性能吃透3.1 Triton路径先把功能跑通Triton是一种面向GPU的Python式DSL写起来简洁社区普及率高很多项目用它做快速原型。DeepGEMM保留了Triton实现意味着普通工程师也能在不太懂CUDA细节的情况下把核心流程读懂、改改参数、跑通自己的用例。Triton自动化处理了很多琐事tile划分、共享内存分配、基本的数据搬运。对多数常规形状来说Triton路径已经能拿到可观的性能因为编译器会帮你做不少调度决策。它的主要限制是不能像手写CUDA那样精细控制异步流水线的每个环节也很难针对某个具体GPU架构做非常极致的寄存器和共享内存布局优化。所以可以把Triton路径看作“大众版”覆盖广、上手快、维护成本低。如果你只是想把一个模型从FP16迁移到FP8先跑Triton路径验证数值是完全可行的一条路线。3.2 CUDA路径把Hopper的潜力吃干净真正要榨干Hopper架构的算力还是得回到CUDA。Triton路径解决的是“能跑”CUDA路径解决的是“跑满”。这部分的优化点我列一下TMA批量拷贝Hopper上的Tensor Memory Accelerator可以异步把全局内存里的矩阵块搬到共享内存搬运过程不占用计算单元和计算流水线重叠。Warp Speciaization把warp分成producer和consumer两类一部分只负责搬运数据另一部分只负责矩阵乘让smem始终有数据可用消除流水线空泡。wgmma指令warp group级别的矩阵乘指令一个指令能完成更大块的FP8乘法比早期mma指令更省指令数和寄存器压力。动态wave调度根据矩阵形状和SM数量计算需要多少个block尽量让每个SM都分配到合适的任务避免最后一个wave出现大量空闲。这些点单独任何一个都不算新鲜但组合起来以后效果是明显的GeMM的算术强度被拉高内存延迟被隐藏SM的利用率稳定在高位。DeepGEMM的CUDA路径就是把这些东西按FP8的特点重新组织了一遍并且把代码控制在紧凑范围内这是它最难得的地方。3.3 为什么保留两条路径我最初觉得保留两套实现是重复劳动后来在实际工程里才理解这样做的好处。Triton路径是维护和教学的最优解CUDA路径是极限性能的兜底选项。运行时可以根据矩阵形状、是否有Hopper特性可用、输入是否规整等条件自动选择走哪条路。这种“双轨”思路对很多团队都有参考价值不要一上来就追求手写CUDA先用高抽象路径验证功能再对热点场景用低层优化做加固。两条路径可以长期共存而不是二选一。DeepGEMM的调度机制本质上也是这种渐进式优化思路的一次实践。4. 复现与实测我看懂了为什么它敢叫高效4.1 环境准备和三个最容易失败的地方想复现DeepGEMM硬件门槛比较清楚Hopper架构的GPU也就是通常说的H100、H200这类SM90设备。软件方面需要较新的CUDA环境Triton也要跟得上版本PyTorch作为上层依赖基本跑不掉。我实际遇到过的三个失败场景分享在这里编译时CUDA arch信息没配对导致报错或者生成出来的kernel根本起不了作用。解决办法是先确认torch.cuda.get_device_capability()输出再设置对应架构不要偷懒用默认值。Triton版本太旧JIT编译内核时直接炸。这个最隐蔽因为安装依赖时不一定报错但运行到特定路径就会挂。建议按仓库README的版本要求再往上对齐一点别用过于保守的版本。环境混装导致第三方库冲突。最好单独建一个虚拟环境避免系统里多个项目的CUDA依赖互相污染。跑通环境后的第一步不要直接上benchmark先跑内置的smoke test。它能一次性验证数值正确性、基本性能和当前架构兼容性是最省时间的路径。4.2 实测大shape足够稳小shape才是亮点我复现时的观察分为两类形状。大shape比如M4096、N4096、K4096这种规整稠密GEMM整体吞吐已经能逼近硬件理论峰值稳定性也高反复跑几轮波动很小。这类场景本来就是通用库的强项DeepGEMM做得好的地方在于没有明显拖后腿。真正让我眼前一亮的是小shape。M只有16、甚至更小时传统库的利用率会掉得很厉害但DeepGEMM还能保持相对高的有效吞吐原因就在于它对wave调度和tile粒度做了针对性处理不会让SM闲着。对推理场景来说这就是实打实的收益因为实际线上请求很少能凑出那种厂商标杆测试里的大batch。场景传统通用库的典型问题DeepGEMM的表现大而规整利用率接近峰值但难以超越稳定、接近理论峰值小而细长SM大量空转吞吐惨淡通过调度保持较高利用率MoE分组kernel launch多单个小GEMM效率低分段GEMM一次处理明显改善4.3 性能对比的“公平陷阱”跑benchmark时最容易犯的错误是拿FP8和FP16直接比FLOPS。FP8的峰值算力本来就翻倍了比赢了也不代表你的模型效果就一样好。更合理的对比是看“最终推理吞吐”或者“训练迭代时间”同时记录数值误差是否在可接受范围内。具体的建议是固定GPU频率多轮warmup后取稳定值对比时用相同的内存布局和同样的输入数据不仅看FLOPS还要看L2命中、kernel launch次数和峰值内存带宽。只有把这几项放在一起才能判断一个实现是“算法层面的优势”还是“配置层面的运气”。关于误差我用BF16参考实现与FP8输出做比对相对误差在可接受范围内符合这一精度档次的预期。如果你发现误差异常大先检查scale初始化和e4m3、e5m2是否混用这两个问题占比最高。5. 迁移到自己项目前需要绕开的五个坑5.1 形状对齐与padding策略FP8 GEMM对形状对齐很敏感。M、N、K是16的倍数时性能最稳达不到时要么由库内部补齐对齐逻辑要么在外部对张量做padding。能外部pad就外部pad因为库内部对齐会消耗共享内存和寄存器的有效利用率。padding的时候注意输出形状也要跟着变最后要把多余部分切掉不要让它影响后续算子。如果你的业务场景里M经常是1、2、4这样的值建议在上层就预先做微小batch合并处理拉高单次GEMM的有效宽度。5.2 FP8缩放因子的初始化细节我最开始测试时偷懒把scale全部设为1结果输出要么溢出成NaN要么精度差到没法用。正确的初始化思路是根据输入张量的实际分布估算统计最大绝对值然后用目标格式上限除以这个值得到初始scale。e4m3的有限上限约是448所以初始scale可以是max_val / 448。如果清楚数据分布比较集中还可以留一点余量避免训练或推理过程中偶尔出现的离群值触发溢出。不同张量的scale要独立计算不能把所有权重共用一个scale。这里有个额外体会scale并不是越大越好。scale过大会把有效数值压得太小尾数位的潜力浪费精度反而下降。比例尺要贴合数据分布而不是一味追求“不溢出”。5.3 集成MoE时的分组约定MoE场景下每个专家分到的token数量不固定这是分段GEMM的主要复杂度来源。接DeepGEMM时要格外注意group id和token偏移信息的传递。传错的话轻则性能回退重则直接算错结果而且这类错误在数值上不一定立刻看得出来很可能隐藏到后续层才会放大量级。我建议先构造一个形状最简单的固定分组用例比如两个专家、每组各8个token逐步验证索引映射关系。跑通之后再上动态分组动态分组的关键是确保每个专家内部的数据是连续的、可以按块搬运的布局而不是松散排布。5.4 混精度使用建议FP8不是万灵药。对精度极其敏感的层比如输出概率分布附近的计算直接用BF16可能更稳妥。我的经验是模型较深层的激活经过多次归一化后分布更稳定用FP8不容易出问题而靠近输入和输出的位置数值变化幅度大混合使用BF16收益更好。实际接入时可以按层配置精度策略前几层用FP8加速中间敏感层保留BF16末尾再回FP8。这种方法能兼顾性能和稳定也方便逐步扩大FP8覆盖范围。再补充一个工程细节切换不同精度时不要只改dtype还要重新确认对应的scale机制是否配套。FP8独有的scaling在BF16里没有反之亦然很多团队就是在这类“只换一半”的地方跌了交椅。6. 一些最后想说的花了两周时间把DeepGEMM拆完并迁移到自己项目里之后我最深的感触是优化不是靠一句“换上8位浮点”就能完成的系统工程。它需要你对底层硬件的数据通路有足够理解也需要对模型本身的数值边界有耐心。把DeepGEMM的源码从头读一遍比翻十篇综述对你的帮助都大因为这里面每一个优化点几乎都能对应到你未来可能在自家框架里遇到的真实瓶颈。如果条件允许我建议你从Triton路径跑通开始再对照CUDA路径逐行去看那些“为什么多写一行”“为什么这样调度”的决策最后再迁移回自己的业务场景验证收益。这样一个流程走完你得到的不仅是性能还有一套可迁移的优化思路。