在 RTX 5090 (SM120) 上补全 NVFP4 量化相关 kernel

本文部分内容由 GPT 5.4 根据大纲写成，我进行了校对和风格修改。如仍有生硬之处，敬请谅解。本文中的 vibe coding 过程使用 GPT 5.4 / Claude Opus 4.6 完成，使用了 Cursor 和 Claude Code 作为不同阶段的辅助工具。

最近组里同学在折腾 5090 上的原生低精度预训练，也就是说，不是 BF16 训练完再量化。英伟达的 Transformer Engine 提供了许多不同数据格式的 recipe 的支持，也就是对于每一层的前向、反向，究竟是用什么数据格式，并为此提供了优化过的 kernel。然而虽然它声称为 Blackwell 优化，但其实指的是 SM100 家族，也就是 B200 / GB200 这类数据中心卡；而 RTX 5090 这样的消费级卡（SM120 家族，还包括 RTX Pro 专业卡、DGX Spark 等），并不存在开箱即用的体验。

MXFP8 的部分似乎还好，只需修改编译选项，为 SM120 编译 kernel 就能工作。但一到 Blackwell 专属的 NVFP4 部分，我们立刻就遭遇了老黄精准的刀法。受到老板的感召，我花了不到两天的时间，几乎完全靠 vibe coding，补全了 TE 2.12 在 SM120 家族上的 NVFP4 kernel 支持，主要包括两个技术：stochastic rounding（随机舍入）和 Random Hadamard Transform GEMM（随机 Hadamard 变换矩阵乘）。

所有的代码都已经在 fp4_sm120 开源。本文简单记录过程和一些心得。

FP4 量化技术

显然，FP4 的动态范围和精度都很小。以 NVFP4 用的 E2M1 为例，能表示的数（绝对值）只有 0, 0.5, 1, 1.5, 2, 3, 4, 6。除了常见的 block-wise scale 以外，Transformer Engine 在 NVFP4 的训练 recipe 里增加了以下的两个技术，来提供更高的训练稳定性和更好的精度。

随机舍入

目前硬件上进行量化时，舍入方法（也是 IEEE 754 规定的）通常是 round-to-nearest，也就是把一个值 x 量化成离它最近的可表示值。然而，当可表示的值只有这么几个，并且间距并不相同时，这样很可能引入系统性的误差。比如 1.6 和 2.4 都会被量化成 2，

而随机舍入（stochastic rounding）的做法是，如果一个值 x 落在两个可表示值 a 和 b 之间，那么应该按比例随机地向上或向下取：

P(round to b) = (x - a) / (b - a)

这样做的好处是，量化的行为在期望上是无偏的。在 SM100 上，NVIDIA 提供了 cvt.rs.satfinite.e2m1x4.f32 这条指令来实现硬件级别的随机舍入，它的格式是 cvt.rs.satfinite.e2m1x4.f32 %0, {%1, %2, %3, %4}, %5，其中 %1 到 %4 是四个输入的 FP32 值，%5 是一个随机 32 位整数，输出 %0 是四个值量化成 NVFP4 后 pack 成一个 16-bit 寄存器的结果。这条指令一次性能对四个数进行量化；当然，所有输入值需要事先经过 block-wise scale 处理，本身不能超出 NVFP4 的表示范围。在 Transformer Engine 的 ptx.cuh 中，如下的函数使用内联 PTX 汇编的方式，调用此指令实现了量化：

• mul_cvt_bf16_to_fp4_4x_with_stochastic_rounding
• mul_cvt_fp32_to_fp4_4x_with_stochastic_rounding
• mul_cvt_bf16_to_fp4_8x_stochastic_rounding

随机 Hadamard 变换

随机 Hadamard 变换（Random Hadamard Transform，RHT）是一种在量化前对数据进行随机旋转的技术。它的主要作用是将原本集中在少数通道上的异常值（outlier）打散，使得后续的 block-wise scale 更容易工作。虽然形式上它能写成一个很大的矩阵乘法 $G_{\text{rht}} = GH$（被称为 RHT GEMM），但实际实现上 $H$ 是一个很小（如 $16 \times 16$）Hadamard 矩阵（元素均为 +1 或 -1 的正交矩阵），作用在沿着最后一个维度分块后的 $G$ 上。换句话说，相当于沿着最后一维把 $G$ 切成许多个 16 元素的小块，每个小块都被一个（每次迭代随机生成的）Hadamard 矩阵所打散。这使得 RHT GEMM 并非一个真正 $O(mnk)$ 的 GEMM 操作，而是降低到了 $O(\vert G \vert)$ 的复杂度，也就是说，它几乎一定是个 memory bound 的算子。

Transformer Engine 在 SM100 上为 RHT 提供了专门的 kernel，叫做 rht_gemm_ntt_w_sfc。它的输入是高精度（如 BF16）的矩阵，输出是 NVFP4 的矩阵；在这个 kernel 里，随机旋转和量化被融合了。也就是说，输入值在被乘以 Hadamard 矩阵打散以后，会直接调用前面提到的 cvt.rs 指令进行随机舍入量化。因此，RHT kernel 的实现除了 GEMM 本身，也依赖于 cvt.rs 的支持。而这里的乘法使用的是 WGEMM，也就是基于 warp-group 的矩阵乘法，依赖于 SM100 上的 tcgen05 指令。

RTX 5090 (SM120) 现状

虽然 SM120 支持 NVFP4 矩阵乘法（由 CUTLASS / cublas 提供），但要用上完整的 NVFP4 recipe，还有以下的问题：

• SM120 架构不存在 cvt.rs.satfinite.e2m1x4.f32 指令。事实上，它完全不支持 .rs 指令后缀。
• SM120 架构虽然配备了第五代 Tensor Core，但不支持 SM100 上的 tcgen05 指令 / UMMA / TMEM 路径，shared memory 也更小（99KB vs 232KB）。因此 Transformer Engine 现有的 RHT kernel 完全无法使用，只能改走 Hopper WMMA 的实现。

但这些并不是不可逾越的障碍。前者完全可以通过软件来模拟，而后者也可以 fallback 到更古老的实现。考虑到这两个 kernel 都是 memory bound 的，大概也不会有什么性能损失。

随机舍入指令：（半）软件模拟

通过多次尝试，我（通过古法）发现 SM120 上最接近 cvt.rs.satfinite.e2m1x4.f32 的指令是 cvt.rn.satfinite.e2m1x2.f32，它一次只能处理两个输入值，并且不支持随机舍入。这并不困难，因为随机舍入的核心其实就是在输入值上加一个小噪声，使得它在两个可表示值之间随机变化，而后依然可以使用 round-to-nearest 的指令来实现随机舍入的效果。具体地，先根据 E2M1 的 ULP 给每个输入值加上一个对称的小噪声，再调用两次 cvt.rn.satfinite.e2m1x2.f32，最后手工把结果 pack 成原来 e2m1x4 的布局。这样虽然底层不是原生 cvt.rs，但最终满足的概率分布是一致的。

除了依赖 SM120 指令的版本，我还要求 Claude 生成了一个纯软件版本（不使用任何 cvt 指令），方便在别的架构（如果还有的话）上验证语义。我同时也增加了 SM100 的原生实现来对拍，结果是：

• SM120 polyfill 和原生 cvt.rs 可以做到 bit-exact（控制提供的随机 bit 一致）；
• 纯软件版本在统计意义上和 cvt.rs 等价；

在这过程中，我被各种细节坑过，特别是指令的输入、输出顺序反复修改了好几遍。最终是通过在 SM100 上对拍了几组数据，才确保了实现的正确性。令我震惊的是，一个有错误输出顺序的实现，居然能让真实训练的 loss 曲线在前几个 iteration 上下降不少，到后面才暴露出问题。

在接入 RS 支持后，NVFP4 的训练稳定性有了显著的提升，并且性能没有可观的下降。毕竟它是 element-wise 的操作，这也是预期之内的。

RHT GEMM：从头重写

接下来，需要把 Transformer Engine 的 rht_gemm_ntt_w_sfc 在 SM120 上重写。具体做法包括：

• 用 WMMA 分块做 16x16 Hadamard 矩阵乘法；
• NVFP4 stochastic rounding 直接复用上一节提供的 polyfill；

在 RTX 5090 上，较大尺寸的测试能跑到大约 1270 GB/s，差不多是显存峰值带宽的 71%。我也把它拿到 SM100 上和 TE 做了二进制对拍；在 fast-math 模式下，数值对比只需要放宽到 $<0.1\%$ 的 FP4 差异容忍度就能通过。

当然这个过程也绝非一帆风顺：虽然 GEMM vibe 起来很轻松，但还是有不少细节经过了多次修复，尤其是对 NaN 的传播处理花费了不少时间。我也要求 Claude 自己通过 nsys 和 ncu 来分析性能表现，进行了多轮优化，才获得了目前的性能数据。

在 vibe 完成后，我们也把 RHT GEMM kernel 接入了 TE 的训练 recipe，然而它的表现却十分诡异：性能有符合预期的少许下降，然而 loss 并没有显著变好，反而是最后一层输出的范数剧烈飙升。我一开始以为是实现问题，就 PUA AI 反复进行验证都没看出问题；直到后面我在 GB200 上用原版 TE 测试了一下，发现几乎完美复刻了 5090 上的曲线，才意识到大概不是 kernel 的问题。

Vibe Coding 过程

注：这一节基本也是 vibe 出来的，所以味比较冲。

其实本来我让 Claude 和 ChatGPT 各自实现了一个版本，仓库早期的提交里，甚至直接就有 cvt.chatgpt.cu 和 cvt.claude.cu 两个文件。我一开始做的事情非常朴素：把背景、目标指令和期望语义喂给模型，让它先给出一个能编译、能跑 benchmark、最好还能带上测试的初版。这个阶段 AI 的效率确实很高，尤其适合：

• 快速铺出大块 CUDA/C++ 样板代码；
• 把分散的 PTX 细节和 wrapper 函数拼起来；
• 顺手生成一堆测试和 benchmark 骨架。

但接下来就进入了传统工程阶段。很多 bug 其实都不是“算法错了”，而是非常 GPU、非常细碎的那种：

• 随机数取法有 bias；
• cvt 的操作数顺序写反了；
• pack 出来的 nibble 顺序不对；
• Philox counter layout 和 TE 的实现并不一致；
• fminf / fmaxf 会吞掉 NaN，而 reference 选择保留 NaN；
• fast-math 模式下，本来就不该追求逐 bit 一致。

也就是说，AI 很适合先把 70 分的代码搭出来；但从 70 分到能真正替代原始 kernel，靠的还是对拍、读 PTX、看硬件行为，以及不断修正那些一眼看不出来的小坑。

时间线

从提交记录来看，这个项目的推进速度其实相当符合“先 vibe 出原型，再用测试把它打磨成工程”的模式：

时间	提交	发生了什么
03-24 14:43	e4419c8	初始提交，生成 cvt.chatgpt.cu 和 cvt.claude.cu 两个版本。
03-24 15:48	1a53719	修掉 rand_byte 的 bias，说明问题已经从“能跑”进入“语义是否正确”。
03-24 21:30	9c5ba8a / d162a57	一边修 PTX 操作数顺序，一边开始加入 SM100 对拍测试。
03-25 00:47	0e25ab8	第一版能在 RTX 5090 上工作的 RHT GEMM 跑起来了。
03-25 01:03	9cd64ab	做了 multi-group-per-block 优化，把带宽推到约 1270 GB/s。
03-25 12:29	1500443	把自己的 kernel 和 TE reference 放进同一个 binary，在 SM100 上正面对拍。
03-25 23:21	705eecc	修 NaN 传播和 Philox RNG 布局，开始解决真正影响训练语义的细节。
03-25 23:49	36de459	给 fast-math 模式加上 <0.1% 容忍度，项目基本收尾。

如果把上面这些 commit 连起来看，会发现这次工作并不是“一句话让 AI 自动写完了 NVFP4 kernel”，而更像是：先让模型把原型和脚手架堆出来，然后人拿着真实硬件、reference kernel 和测试，一路把语义修到能用。

总结

这次 vibe coding 整体还算顺利，尤其是考虑到我对 NVFP4 相关的细节并不熟悉。现代 AI 模型在代码实现、性能分析方面的能力也让我比较满意。然而我也深刻感受到，（至少目前）在此类工程问题上，人类经验依旧是不可或缺的；否则，只会变成白白浪费 token 让模型在一些细节上反复纠结，生产出一坨又一坨的意大利面。