ThunderKittens移植至NVIDIA Vera Rubin NVL72：NVFP4 GEMM突破22 PFLOPS

深度学习内核库 ThunderKittens 迎来一次关键升级——团队将其移植到了 NVIDIA 最新的 Vera Rubin NVL72 平台，并围绕新硬件重构了 NVFP4 GEMM 内核。最终结果相当亮眼：性能从原先仅占 roofline（理论峰值）的 42%，一路攀升至超过 22 PFLOPS，已经能与 cuBLAS 和 CuTe DSL 这类官方高度优化方案正面竞争。

这篇文章梳理这次移植背后的技术变化，以及新硬件 ISA（指令集架构）层面带来的机会与挑战。

ThunderKittens 是什么，为什么移植值得关注

ThunderKittens 是一个面向 GPU 的轻量级深度学习内核编写框架，核心目标是让研究者和工程师能够以更简洁的抽象写出接近硬件极限的高性能内核。相比直接手写 CUDA 或依赖庞大的模板库，它强调「小而美」的编程模型，围绕 tile（分块）这一核心抽象组织计算与数据搬运。

将这样一个框架移植到全新架构，不只是换个编译目标那么简单。每一代 NVIDIA 架构在张量核心、内存层级、异步数据搬运指令上都有显著演进，内核库想要吃满新硬件的性能，往往需要针对性地重写关键路径。这次从 42% roofline 到 22 PFLOPS 的提升，正说明了「针对新架构重构」与「简单编译通过」之间的巨大差距。

NVFP4 GEMM 为何是性能核心战场

这次升级的重点落在 NVFP4 GEMM 上。GEMM（通用矩阵乘法）是几乎所有深度学习负载的性能基石，从注意力机制到前馈网络，其核心计算都可以归结为矩阵乘。谁能把 GEMM 跑到接近理论峰值，谁就掌握了整个推理与训练栈的效率上限。

NVFP4 则代表着低精度计算的新趋势。随着模型规模膨胀，4 位浮点（FP4）格式在保持可接受精度的同时，大幅压缩了内存带宽与算力需求。NVIDIA 在新一代硬件上为这类低精度格式提供了专门的张量核心支持，但要真正发挥这些单元的威力，软件侧必须精确匹配硬件对数据布局、缩放因子、累加精度的要求。

团队表示，正是围绕 Vera Rubin NVL72 的新硬件特性重建了 NVFP4 GEMM，才把性能从 roofline 的 42% 推到了竞争级别。这意味着原有内核在新架构上存在明显的「水土不服」，而重构填平了这道鸿沟。

FP4（4位浮点）格式是近年低精度计算演进的前沿。相比广泛使用的 FP16/BF16（16位），FP4 将每个数值的存储压缩至4位，理论上可将内存带宽需求降至1/4，同时允许张量核心在相同面积下堆入更多并行计算单元。然而极低的位宽也意味着动态范围极窄——标准 FP4 格式仅能表示十几个不同数值，必须配合「缩放因子（scale factor）」机制才能覆盖实际模型权重的分布范围。NVIDIA 在 Blackwell 及 Vera Rubin 架构中引入的 NVFP4 正是一种带有块级（block-wise）缩放的 FP4 变体，硬件层面原生支持缩放因子的读取与融合计算。这使得 NVFP4 GEMM 的软件实现远比普通矩阵乘复杂：内核不仅要管理主矩阵的分块加载，还要同步处理每个数据块对应的缩放因子，并确保两者在张量核心的流水线中精确对齐，否则精度与性能都会出现严重退化。

ISA 变化带来的机会

文章指出，这次性能飞跃的关键在于利用了新架构 ISA 层面的变化。虽然原文未展开所有细节，但从 GPU 架构演进的一般规律看，这类改进通常来自几个方向：

更强的张量核心与低精度路径

新一代张量核心往往为 FP4/FP8 等格式提供更高吞吐和更灵活的缩放机制。软件只要正确对接这些指令，就能在不增加精度损失的前提下获得成倍算力。

异步内存与数据搬运优化

GEMM 的性能瓶颈常常不在计算本身，而在数据能否及时喂给计算单元。新 ISA 对异步拷贝、共享内存调度、集群级（cluster）协作的支持，直接决定了能否隐藏访存延迟、逼近 roofline。

NVIDIA 从 Ampere 架构开始引入异步数据搬运指令（如 cp.async），Hopper 进一步引入 TMA（Tensor Memory Accelerator）硬件单元，允许 SM 将数据从全局内存或 NVLink 域搬运到共享内存的操作完全外包给专用引擎，主计算流水线无需等待。Vera Rubin 延续并扩展了这一路线。对 GEMM 而言，这意味着可以构造「双缓冲（double buffering）」乃至更深的流水线：当张量核心正在消费第 N 个数据分块时，TMA 已在后台预取第 N+1 块，理想情况下访存延迟被完全隐藏。ThunderKittens 的 tile 抽象天然契合这种流水线模式——框架可以将异步预取封装进调度原语，让开发者无需手工管理同步屏障，却仍能获得接近手写汇编的访存效率。

NVL72 的规模化互联

Vera Rubin NVL72 本身是一个高密度互联平台，72 这个数字暗示了大规模 GPU 域的协同。在这种规模下，内核不仅要单卡高效，还要能在多卡协作中保持扩展性。22 PFLOPS 这一量级的数字，正是建立在平台级算力聚合之上。

与 cuBLAS、CuTe DSL 同台竞技意味着什么

cuBLAS 是 NVIDIA 官方的高度优化线性代数库，CuTe DSL 则是 CUTLASS 生态中用于描述张量布局与计算的领域专用语言——两者都代表了业界内核优化的第一梯队。一个相对轻量的开源框架能在新硬件上做到与它们「competitive（有竞争力）」，本身就是对其编程模型与抽象设计的有力背书。

对开发者而言，这释放了一个积极信号：你未必需要钻进官方闭源库或复杂模板的深水区，也能通过 ThunderKittens 这样的框架在最新硬件上写出接近极限的内核。这降低了前沿硬件的性能开发门槛，对研究社区尤其重要。

CuTe DSL 是 NVIDIA CUTLASS 库中的核心子系统，使用 C++ 模板元编程为张量的逻辑布局（logical layout）与物理存储（physical stride）提供统一描述，并能在编译期生成高度特化的内核代码。其设计哲学是「零成本抽象」：开发者用高层布局语言描述计算，编译器自动推导出最优的数据搬运与指令选择策略。这套机制极为强大，但学习曲线陡峭，调试复杂。ThunderKittens 与之的差异在于：前者选择以「可读性与研究迭代速度」为优先，牺牲部分极端情况下的最优性，换取开发者能在数小时内写出可运行的自定义内核。此次在 NVL72 上接近 CuTe DSL 性能，说明这一取舍在最新硬件上并未带来无法接受的代价。

结语：硬件迭代与软件适配的持续赛跑

这次移植案例再次印证了一个事实：硬件的理论峰值从不自动转化为可用性能，中间隔着大量针对性的软件工程。从 42% 到 22 PFLOPS，差距不在芯片，而在软件是否吃透了新 ISA。

随着低精度格式和大规模互联平台成为 AI 基础设施的主流方向，像 ThunderKittens 这样敢于紧跟最新架构、快速重构核心内核的开源项目，将在性能生态中扮演越来越重要的角色。对关注 AI 算力栈的工程师来说，这类工作值得持续追踪。