FP8 GEMM这个主题放在CDNA4 GPU上写很多人第一反应是“直接调库不就行了”。但对真正做算子优化的人来说库函数只是起点。尤其是面对大模型推理的定制化shape、量化scale融合、以及各种花式矩阵布局时手写内核带来的性能空间远比想象中要大。这篇文章就围绕“在CDNA4 GPU实现高性能FP8 GEMM”这件事从格式选择、硬件指令细节、调度策略到实际调优踩坑完整走一遍。适合正在做ROCm算子开发、推理引擎适配或者想从CUDA生态迁移过来、想把FP8吞吐真正吃满的人参考。老规矩先说结论FP8 GEMM是当前大模型推理最依赖的核心算子没有之一。模型里的QKV投影、MLP、注意力合并几乎全是GEMM。把这块的性能榨干整个引擎的吞吐就上去了。CDNA4这一代在FP8的矩阵计算、L2调度、以及指令并发能力上都做了不少调整堆砌得很厚重但如果你只是调一个通用库大部分红利其实都吃不到。所以这篇文章的核心是带着各位从硬件的数值格式和指令原语出发自己把GEMM内核写出来、调好并给出可以直接复用的实验参数和排查方法。1. 项目背景与整体思路拆解1.1 为什么偏偏是FP8 GEMM先说背景。这两年我在做的大模型推理部署里FP8量化已经从一个可选项变成了必选项。原因很直白显存占用直接砍半带宽需求直接砍半而矩阵计算的峰值吞吐在FP8路径上是FP16的两倍以上。换句话说同样的卡、同样的模型FP8量化之后能塞进更大的batch单token延迟更低整体成本曲线也完全不同。但这只是理论收益。真正落地的时候你会发现想要把FP8算力真正跑满绕不开GEMM这个算子。而GEMM从FP16切到FP8不只是把数据类型换成8位就完事。硬件上MFMA指令的形态、LDS的共享带宽、L2的调度策略全都会变。很多人照着FP16的kernel模板改结果发现性能不仅没涨反而因为量化scale的额外搬运和寄存器布局不匹配而变慢就是这个原因。CDNA4这一代对FP8的重视程度从指令集的支持就能看出来。FP8矩阵乘指令不再只是“能用”的状态而是硬件真正围绕它优化过的结果。如果你的内核结构还是按FP32时代的老思路设计等于开着一台高性能车却一直挂在二挡。那么FP8 GEMM到底难在哪我认为有三个核心点第一是数值格式的溢出控制和scale设计第二是指令级别的分块与寄存器映射这直接决定你能发挥多少硬件算力第三是整个K循环里的数据搬移调度L2和LDS的带宽分配在这里往往比计算指令本身更敏感。本文后面的内容基本就是围绕这三点展开的。1.2 方案选型从库函数到MFMA指令在CDNA4上实现FP8 GEMM基本有三条路可以走。我分别试过说下各自的位置。第一层是直接用hipBLASLt这类高层库函数。它适合快速验证基线常规shape下的性能已经做得不错驱动更新后还会自动优化。但对定制化shape或者需要把量化scale、激活函数融合进GEMM内部时库函数的通用设计反而成了限制性能和灵活性都达不到要求。第二层是用rocWMMA这类WMMA封装库。它在MFMA之上做了一层C接口让你不必直接面对汇编指令开发效率高很多。但封装层同时意味着编译器替你做的决策不一定最优而且早期实现里块边界处理、网格计算这些小问题着实不少我在调试时踩过几回坑。第三层是直接写MFMA指令的内核这也是我最终选择的主路线。虽然开发成本高但你能完全控制寄存器分配、LDS使用和等待周期性能潜力也是三层里最大的。对于每天都在跑的大模型shape来说自己做一套内核的成本摊薄下来非常划算。我这边的实用建议是不要一上来就手写。先用hipBLASLt建立一个可信基线既能验证你的kernel数值对不对也能作为性能锚点。再在这个基础上用rocWMMA或手写MFMA逐步替换关键计算部分。这样既保证了开发效率又不至于让性能失控。1.3 整体设计框架我的GEMM内核整体采用经典的“大K循环 输出块划分 双缓冲LDS”结构。每个workgroup负责一个输出块比如128x128内部循环K方向每次处理K32或K64的切片。A块和B块分别异步搬运到LDS计算单元再从LDS读取数据送入MFMA累加结果始终保持在FP32寄存器中只在最终写回时考虑是否做反量化或激活融合。这段结构里有两个决定成败的关键。一是双缓冲加异步加载让LDS搬运和MFMA计算尽量重叠二是wavefront之间的调度策略CDNA4上的L2带宽和等待周期处理得好不好直接能拉开百分之二三十的差距。后面我会把这两块的细节单独展开讲。2. FP8格式细节与硬件原语要点2.1 E4M3与E5M2怎么选Scale又该怎么排FP8并不是只有一个格式它有E4M3和E5M2两种变体。差别简单说就是精度和动态范围的取舍。格式符号位指数位尾数位最大有限值典型用途E4M3143448前向推理中的权重与激活E5M215257344梯度、误差累积等动态范围大的场景E4M3的尾数多一位精度更好但指数范围窄特别容易溢出。E5M2反过来动态范围大很多但精度相应下降。对于GEMM前向推理来说A和B两个输入矩阵基本都用E4M3因为权重和激活经过量化之后数值范围相对可控精度更重要。FP8还有一个反直觉的坑它没有NaN表示。这意味着一旦数值溢出不会像FP16/FP32那样产生NaN让你立刻发现问题而是会变成Inf或者很大的有限值乘出来之后整个结果直接发散。很多人在FP8 GEMM调试时遇到“偶尔有几个数值离谱”的问题十有八九是scale没选好导致中间结果溢出。Scale的组织方式从大到小有per-tensor、per-column、per-row三种。per-tensor最简单整个矩阵共享一个scale但遇到数值范围分布不均衡的场景误差大per-column、per-row的误差小但需要在加载阶段把scale矩阵一并准备好额外搬运开销不可忽视。我的经验是先用per-tensor跑通全流程再根据量化误差决定要不要升级到per-column。一上来就做细粒度scale往往会在后续调试时分不清是scale问题还是内核问题。另外一个非常容易被忽视的点scale的乘法位置。A矩阵和B矩阵各自的scale要在数据从LDS读入寄存器、准备进入MFMA之前就乘好。不要等到累加之后再反算那样不仅增加额外指令还会让FP32累加器的数值和硬件预期不一致。2.2 MFMA指令原语与矩阵寄存器分配CDNA4上的FP8矩阵运算核心是MFMA指令。这类指令以wavefront为执行单位一条指令内部完成一个小矩阵的乘法累加。常见的形态包括v_mfma_f32_16x16x32_f8和v_mfma_f32_32x32x16_f8命名里的数字分别代表这个指令负责的M、N、K维度。拿16x16x32举例一个wavefront的64个线程协同计算一个16x16的输出子块K方向一次吃32个元素输入是FP8累加输出是FP32。为什么输出必须是FP32这是硬件设计决定的也符合数值稳定的需求。FP8的动态范围和精度都不足以做累加保持FP32累加是整个GEMM精度的底线。64个线程如何分配这些数据每个线程持有A矩阵和B矩阵的不同寄存器片段。这种映射关系如果你不理解后面调寄存器分配时会非常被动。简单说一个wavefront内部的线程被组织成特定的矩阵布局每个线程拿的是A的一部分行和B的一部分列通过MFMA指令内部的硬件互联完成乘积和累加。你手动写kernel时必须保证写入寄存器的数据位置和指令隐含的布局一致否则结果会完全错乱。两种常用指令形态的实际差异主要还是分块策略的取舍。16x16x32的K方向更长单指令的数据复用更好适合K维度受限的重计算场景32x32x16的M/N方向大单指令覆盖的输出面积更大指令数更少适合输出块较大、LDS带宽较宽裕的场景。没有绝对好坏要看你的整体tile大小和目标shape来做组合。CDNA4在指令并发上还有一个值得注意的点我实测下来MFMA和普通VALU指令可以在一定条件下同时发射。这意味着你可以在等待MFMA结果的同时用空闲的向量单元去处理下一块LDS数据的scale乘法或者地址计算。这条特性用好之后整个K循环的等待周期能明显压缩下来。但这个行为依赖驱动和编译器版本不同环境的表现会有差异别把它当成永远可靠的条件。2.3 分块策略与L2复用GEMM内部分块设计本质上是拿寄存器、LDS和L2的容量去换数据复用率。FP8的优势在这里体现得很彻底同样大小的tileFP8只占FP32四分之一的存储空间所以你能在LDS里放更深的K切片。我常用的输出块大小是128x128配合K切片32或者64。这个尺寸下每个workgroup的LDS占用大约在16KB到32KB还有余量做双缓冲甚至三缓冲。双缓冲保证本轮MFMA计算期间下一轮的A/B数据已经在异步搬运途中。三缓冲能进一步缓解LDS读写竞争但实际收益要看是否增加了LDS bank冲突。说到LDS bank冲突这是FP8 GEMM最容易踩的坑。A矩阵和B矩阵的tile从global memory搬到LDS之后每个wavefront以特定步长读取。如果行宽正好是某些对齐数值多个线程可能同时访问同一个bank产生冲突。解决方式是在LDS布局里对每一行做padding比如在行尾补几个字节把访问点错开。这个操作看似不起眼但实测能带来接近15%的带宽提升。再往上一层是L2复用。GEMM的K循环里同一块A slab会和右侧的每一块B slab相乘所以A slab在L2里待的时间越长越好。设计时我会尽量把workgroup的N方向调度做得宽一些让同一块A在L2里被多个B slab命中而不是频繁从global memory重新加载。CDNA4的L2带宽很大但只有你的访问模式匹配它的调度方式才能真正吃满。3. 实操过程与核心实现3.1 工具链检查与基线建立动手写内核之前先把环境确认清楚。CDNA4 GPU需要配套版本的驱动和ROCm软件栈我习惯先用rocminfo查看设备信息再用clang --version确认编译器版本避免后面出现指令不支持或者汇编映射错乱这种莫名其妙的问题。编译内核的命令很简单常规优化参数就行hipcc -O3 -stdc17 -marchnative -o test_gemm test_gemm.cpp如果后面引用了rocWMMA再额外链接对应的头文件路径和库文件即可。基线这一步我强烈建议别省。直接用hipBLASLt调用一个与你业务shape一致的GEMM拿到正确输出和耗时。这一步为什么重要因为手写内核调试时你需要一个“官方答案”来判断自己的计算结果是否可信。你要是连基线都懒得跑真出了数值问题都不知道该跟谁对齐。计时我统一用hipEvent来做。注意要连续运行多次取中位数第一次调用通常会包含初始化开销不能拿它当参考。3.2 基于rocWMMA的FP8 GEMM改版如果你是第一次写FP8 GEMM不建议直接跳到汇编层面。先用rocWMMA把整个流程跑通理解tile加载、MFMA执行和数据存储的先后顺序再考虑手写。一个典型的rocWMMA内核骨架长这样#include rocwmma/rocwmma.hpp using namespace rocwmma; // 定义输出分块尺寸 constexpr uint32_t BLOCK_M 128; constexpr uint32_t BLOCK_N 128; constexpr uint32_t BLOCK_K 32; // 每个线程持有的矩阵片段 fragmentmatrix_a, BLOCK_M, BLOCK_N, BLOCK_K, precision, layout_a a_frag; fragmentmatrix_b, BLOCK_M, BLOCK_N, BLOCK_K, precision, layout_b b_frag; fragmentaccumulator, BLOCK_M, BLOCK_N, BLOCK_K, float c_frag; // 加载、乘累加、存储三步 load_matrix_sync(a_frag, a_ptr, lda); load_matrix_sync(b_frag, b_ptr, ldb); mma_sync(c_frag, a_frag, b_frag, c_frag); store_matrix_sync(c_ptr, c_frag, ldc, memory_order::seq_cst);这段代码逻辑上是通的但它只是一个毛坯。实际跑的时候要处理几个现实问题。第一BLOCK_M、BLOCK_N并不一定是底层MFMA指令直接支持的尺寸rocWMMA内部会帮你拆解但拆解效率取决于编译器版本和驱动调度。第二scale乘法必须插入在load_matrix_sync之后、mma_sync之前否则进到MFMA里的数据还是原始FP8结果就是纯数学错误。第三K循环需要自己管理双缓冲swap的时机要看上一轮MFMA是否已经消费完数据。我在用rocWMMA时踩过的最深坑是早期版本在较大tile下的边界处理bug。具体表现是小shape测试全部正确一上大shape就会出现个别block结果错乱而且复现概率不稳定。后面定位到是block index计算里某个参数溢出但这个问题很隐蔽因为编译器不会报错。教训就是库封装的便利背后藏着调度细节的黑盒关键路径上你还是要清楚它内部做了什么。3.3 手写MFMA内核与性能调优参数手写内核的核心是把K循环里的每一步都控制在自己手里。我的结构大致是// 伪代码workgroup负责128x128输出K方向每轮处理64 // 双缓冲LDSa_buf[2][...] b_buf[2][...] // 每轮 // 1. 异步加载下一段K到LDSglobal - LDS // 2. 从上一段LDS读取A/B到VGPR并带上scale因子 // 3. 执行16条mfma指令累加结果到FP32寄存器 // 4. 等待异步加载完成swap buffer进入下一轮这个循环设计的关键是第1步的加载和第2、3步的计算尽量重叠。MFMA读LDS的时候加载单元还在往另一块LDS缓冲区写互不干扰。等到本轮MFMA执行完下一轮的数据也刚准备好等待周期被压到很小。scale合并我会放在步骤2里也就是从LDS读到VGPR时顺手乘上。这样MFMA拿到的数据就是正确数值不需要额外做反量化或预处理等于把scale的100%开销都藏进了加载流水线。性能调参的部分我记录过一组相对基线的数据。测试配置是MN4096K8192使用E4M3、per-tensor scale注意这组数据只代表趋势不同驱动和编译器版本会有出入。输出块wavefront数LDS双缓冲内循环K段相对hipBLASLt性能64x641开启K160.94128x642开启K321.08128x1284开启K321.26256x1288开启K641.19可以看到输出块从64x64涨到128x128时性能提升很明显这是因为更大的输出块意味着更高的数据复用率LDS和L2的带宽利用率上去了。但继续涨到256x128反而下降原因是每个workgroup的寄存器压力过大部分累加开始溢出到local memory性能直接崩掉。这个拐点位置因卡而异但大方向是一致的。3.4 正确性验证与结果记录手写GEMM的正确性验证不能拿FP16的标准去套。FP8量化本身就会引入1e-2级别的相对误差你拿手写内核的输出和FP32 CPU参考做逐元素对比肯定一片飘红。我的做法是分两级验证。第一级是纯数值一致性。用相同的输入分别跑hipBLASLt和手写内核比较输出。两者都是FP8输入、FP32累加如果scale设计一致输出应该非常接近。这个对比能验证你的指令映射和寄存器布局没有摆错。第二级是推理端到端验证。把GEMM内核接回真实模型里跑一遍典型输入观察loss曲线或者下游指标是否正常。量化误差只要没有异常放大通常都能在可接受范围。这一步是最终裁决因为纯矩阵数值误差和模型容错能力是两个维度。我在大K和大shape下做性能记录时会跑一组不同shape的结果。大K的GEMM上手写内核的优势最明显可以比基线快25%到30%。中等K shape下优势缩小到10%左右。小K shape反而会略输给hipBLASLt因为kernel启动和tail循环的开销占比变大手写的精细优化没机会发挥。这也是我后来坚持做kernel变体的原因不能一个版本打天下。4. 常见问题与排查技巧实录4.1 高频问题速查表问题表现可能原因解决方案结果出现Inf或极大值FP8溢出scale偏小重新校准scale检查absmax小shape正确、大shape错乱网格计算溢出或边界判断错误检查block index对应关系性能低于hipBLASLt输出块太小L2复用不足增大输出块尝试128x128性能原地踏步甚至下降寄存器溢出贷款到LDS减少每线程累加寄存器调整块尺寸结果偶尔错乱、复现率低异步加载和MFMA之间的wait缺失检查s_waitcnt位置补上依赖等待数值相比基线差异过大scale乘法的位置或顺序不对统一为LDS读入VGPR时乘scale指令数很多但性能不高MFMA等待LDS加载流水线未重叠确认双缓冲真正生效调整K段长度4.2 用性能工具定位瓶颈的排查思路当性能不符合预期时我习惯先用rocprof抓一遍指令分布和性能计数器然后再看代码。首先看Wavefronts数量和MFMA指令总量判断计算指令本身是否达到预期。如果MFMA指令数充足但执行时间很长问题多半在等待周期如果指令数本身就偏少说明你的分块策略没把计算吃满。等待周期的问题是排查中最常见的。你可以临时把s_waitcnt替换成最保守的写法也就是每个MFMA之前都强制等待所有加载完成看看性能变化。如果性能反而上升说明你原本的等待策略有bug如果性能下降说明双缓冲和async加载原本是对的瓶颈在别处。还有一个折中的定位技巧我先写一个只包含单个workgroup的最小版本固定输出块尺寸然后逐步扩大线程数和块数。每一步都确认性能和正确性没退化再做下一个改动。这样做的好处是哪怕最后出了问题也能很快缩小到某个具体改动上不用从头查一遍。4.3 独家避坑经验手写FP8 GEMM这么久有几点心得特别想分享。第一别过度迷信理论峰值。CDNA4的FP8峰值是基于最理想的纯MFMA计算得出来的实际运行时还要挤出一部分带宽给scale加载、结果写回以及LDS的读写。以我的经验能把FP8峰值跑到80%以上已经是非常优秀的kernel了。如果你看到网上有“逼近峰值100%”的说法基本要么是理想情况要么是拿特化shape刷出来的参考价值有限。第二权重矩阵数据布局值得提前处理。我最初把转置逻辑放在kernel内部结果小K shape下的转置开销直接吃掉了手写优化的全部收益。后来改成在模型加载阶段把权重矩阵按kernel需要的布局预先摆好kernel内部只做连续读取效果立竿见影。这个思路同样适用于量化scale矩阵尽量提前算好、铺好不要在热点循环里做格式转换。第三lib库能跑就绝不手写这始终是我的第一原则。手写只针对那些库函数覆盖不好、但你又高频使用的定制shape。我这次做FP8 GEMM也是因为实际项目中遇到了库函数处理不了的特殊结构才下手的。如果只是普通的大K GEMMhipBLASLt已经够好没必要重复造轮子。5. 一些个人体会与后续方向真要说这一轮优化最大的体会就是FP8 GEMM的性能瓶颈往往不在计算本身而在你用什么样的心态对待数据流水。寄存器、LDS、L2这三层存储的协同调度比单纯抠一条MFMA指令的周期重要得多。很多人在写kernel时习惯性把注意力放在“指令效率”上但真正拉开差距的是数据搬到哪、什么时候搬、scale在哪一步乘进去。还有一个很实际的经验第一版kernel不要写复杂。先写一个简单、正确、性能平庸的版本确保所有数据路径没问题再逐段替换成优化实现。我见过太多人一上来就上双缓冲加异步加载结果出了问题根本不知道从哪一行开始查。性能优化是连续逼近的过程不是一步到位的魔术。我后续打算把这个内核继续往两个方向延伸一个是batch GEMM也就是多头注意力里那种多个小矩阵并列计算的场景它跟单大GEMM的调度策略差异很大另一个是fusion把FP8 GEMM和后面的激活函数、LayerNorm融合在同一个kernel里省掉中间结果的往返带宽。如果后面有空我会再整理一篇更完整的batch GEMM实践出来。