Together 的内核团队最近获得了 NVIDIA Vera Rubin NVL72 平台的访问权限。过去几天里，我们深入研究了新的 ISA，并用微基准测试对芯片进行了探查。有许多有趣的新特性！我们已经在 ThunderKittens 中完成了一些功能的添加，以便在 Vera Rubin 上编写 NVFP4 和 FP8 GEMM，同时帮助其他小猫探索星辰。
在深入探讨 Vera Rubin 带来了什么之前，我们先快速回顾一下 Blackwell GPU GEMM，以此作为旅程的起点。
起点：NVIDIA HGX B200 GEMM
NVIDIA Blackwell 架构的第五代张量核心从根本上改变了 GEMM 编程模型。NVIDIA Hopper 架构的 wgmma
指令由 warpgroup 集体发出，而 Blackwell 的 tcgen05
指令由单个线程发出，使得一个小的生产者 warp 就能驱动张量核心。累加器也从寄存器移到了 Tensor Memory，操作数直接从共享内存读取，使得单个 MMA 可以跨越两个 SM 上的两个 CTA。
为了在 Blackwell 上达到有竞争力的性能，我们的 GEMM：
- 启动线程块集群，使每对 CTA 可以通过 TMA 多播共享操作数，将 HBM 的内存流量减半。
- 在集群内对 warp 进行特化：加载器通过 TMA 将 A 和 B 带入共享内存，单个 MMA warp 驱动张量核心，一个消费者 warpgroup 将完成的累加器从张量内存搬运到 HBM。
- 持久运行，一个 tile 的输入还在流入时，前一个 tile 的输出仍在排出。
通过这些努力，我们得到了以下结果。
在我们的早期 Together 博客文章 或 ThunderKittens 2.0 发布中，可以了解更多关于这些内核及其优化的内容！
由于 Rubin 保留了 Blackwell 编程模型，我们旧有的 GEMM 仍然可以工作。然而，在 Vera Rubin 上朴素地运行它们时，我们观察到我们的 NVFP4 和 FP8 内核仅分别达到 roofline 的约 42.1% 和 44.4%——还有很大的优化空间！
本文的其余部分分为两部分。首先，我们介绍对 GEMM 重要的 Rubin 新特性，以及如何在 ThunderKittens 中使用它们。然后，我们逐步将这些特性集成到现有的 Blackwell NVFP4 内核中，将其提升到超过 22 PFLOPS，并与 cuBLAS 和 CuTE DSL 竞争。
我们发现的核心问题是，虽然 Rubin 使张量核心消耗操作数的速度提高了一倍，但我们旧的 Blackwell 内核无法足够快地为其提供数据。要达到计算上限，我们需要让 tile 从它们已经在芯片上的数据中榨取更多的复用。
NVIDIA Vera Rubin 平台有哪些新特性？
对比厂商规格，我们看到从 Blackwell 到 Vera Rubin 的以下改进。
就编写高性能 GEMM 而言，我们特别注意到以下特性。
1. 张量核心的 K 翻倍
回顾一下，tcgen05.mma
在 MxNxK 瓦片上计算 C = A@B + C，每一步沿 K 消耗固定数量的字节。在 Blackwell 上，这一步是 32 字节，但在 Vera Rubin 上可以增加到 64 字节。MMA 本身仍然占用相同的周期数，因此 K 翻倍让我们能在相同的指令窗口内打包两倍的工作量。
在 ThunderKittens 中，我们通过为现有 mma 操作新增一个模板参数来表达这一点。
2. 张量内存增长到 576 列
Blackwell 引入了张量内存的概念，这是一个 128 通道 x 512 列 x 32 位的空间，张量核心可以直接读写。在 Vera Rubin 上，该空间增加到 576 列，额外提供 32 KiB 的张量内存可供使用。
请注意，这些额外的列只能通过 .exclusive
限定符访问，这是 PTX 9.4 的新增特性，确保一个 SM 上只有一个活跃的张量内存分配。非独占分配仍然上限为 512，且必须是 2 的幂。
在 ThunderKittens 中，用户可以通过向我们的张量内存分配器传入模板参数来请求这一点，指定该分配为独占。
3. 共享内存增加到 328 KiB
Hopper 和 Blackwell 提供 228 KiB 的共享内存，而 Vera Rubin 引入了超大共享内存模式，可以动态增加到 328 KiB。这是一个主机端规格，可以如下调用。
4. B 侧收集器
Blackwell 引入了收集器缓冲区的概念，这是一个小型 MMA 暂存缓冲区，可以锁存一个 A 瓦片，使下一条指令从那里获取它，而不是从共享内存获取。Vera Rubin 通过 .collector::b::*
将这一功能扩展到 B 瓦片。
为了利用这一点，我们用四个标签之一标注每个 MMA 的操作数，描述其对收集器缓冲区的操作。
- “FILL”从共享内存读取操作数并锁存
- “USE”从缓冲区读取
- “LASTUSE”从缓冲区读取并释放
- “DISCARD”是默认值，跳过锁存
这些标签是复用的权限限定符，而非保证。这意味着即使张量核心有复用的权限，它仍然可能重新加载矩阵。
既然任一操作数都可以驻留在收集器缓冲区中，我们就可以尝试新的模式。例如，在一个 2x2 块上，两侧都收集可以将四个 MMA 从八次操作数读取减少到仅五次。
在 ThunderKittens 中，我们可以通过以下方式暴露这一点。
5. 提前释放 A
一个 tcgen05.commit
在其处理的 MMA 完成后到达一个 mbarrier，向生产者发出信号，表明一个 stage slot 可以重新使用。PTX 9.4 引入了 tcgen05.commit.sync_restrict::shared::read::mma::a
，这是一条新指令，允许我们为 A tile 提前发出信号。我们不必等待 MMA 完成，而是可以在 MMA 完成从共享内存读取其 A 操作数后立即触发 barrier，从而使我们能够通知 tma loader 开始存储下一 stage 的内存。
ThunderKittens 引入了一种新的 commit 类型，供用户表达这一点。
构建 GEMM
我们现在有了新特性和一个 Blackwell GEMM。以下各节会逐步将这些特性集成到我们现有的 Blackwell kernel 中，并解释当我们转向 Vera Rubin 时为什么需要它们。
加宽指令：
最直观的瓶颈来自仍然依赖 Blackwell 的 32 字节 K 步长。在 Vera Rubin 上，该编码的 ISA 上限约为 16.8 PFLOPs，而我们的 NVFP4 Blackwell GEMM 开箱即已达到 14.7 PFLOPs（上限的 88%）。要进一步推进，我们必须将 MMA 处理的 K 字节数加倍。
然而，仅仅为现有 Blackwell kernel 开启更宽的编码，我们注意到性能只有轻微提升，而不是预期的 2 倍。虽然更宽的 MMA 使 tensor core 消耗操作数的速度加倍，但它无助于我们提供操作数的速度。要让 double-K 有用，我们拉动两个杠杆来让核心保持满意：移动更少字节，并加深流水线以确保这些获取被重叠。
读取更少字节：
为了获取更少字节，我们在同一 CTA 对上的 M 维度堆叠第二个输出 tile。由于两个累加器仅在 M 上不同，我们能够为它们共享同一个 B chunk。我们最初的 Blackwell NVFP4 kernel 使用 1x1 tiling 格式，这意味着覆盖 M512xN256 的输出需要两个 pair job，每个都独立传输自己的 B 副本。通过转向 2x1 格式，我们可以只获取一次 B 并覆盖相同的输出，从而减少操作数流量。
这种 2x1 tiling 格式在我们的 NVFP4 Blackwell kernel 中不容易实现，因为 tensor memory 上限为 256 KiB。两个 M256xN256 累加器已经占用 512 列，这意味着 block-scaled MMA 没有空间存储其 A 和 B scale。虽然程序员可以通过让 epilogue warp 在发出 MMA-empty 状态信号之前只加载累加器列的一个子集来绕过这一限制，从而允许下一个 K tile 的 MMA 开始，但这引入了一小部分无法隐藏的延迟。幸运的是，借助 Vera Rubin 额外的 64 列，我们可以存储缩放因子，而无需进行这种操作。
深化流水线：
改变后的分块格式减少了操作数流量，但无助于缩短每次取数所需的时间。下一个挑战是让张量核心持续获得数据供给。Vera Rubin 更大的共享内存使我们能够构建更深的流水线，提前暂存更多分块，并为传输留出更多完成时间。针对我们的 NVFP4 和 FP8、16k 方阵 GEMM 扫描环形深度，我们观察到以下结果。
NVFP4 16k 方阵 GEMM：
FP8 (E4M3) 16k 方阵 GEMM：
虽然最大的收益似乎来自这最后一步，但我们注意到这些收益依赖于我们此前的优化。下面是对 K 步长、共享内存流水线和分块策略分别进行扫描的结果。
最后的润色：
为了进一步推进我们的内核，我们尝试了几个最后的调节项
内核配置调优：我们针对不同工作负载进一步调优内核以获得最大性能。关于 CTA 对大小，我们发现 Blackwell 内核最初的 1x1 分块格式在较小的方阵工作负载上性能最佳。对于更大的形状，我们采用 2x1 CTA 对分块，以及调优后的 2、4 或 8 个 CTA 的簇大小。此外，在所有形状上，我们排列分块光栅化顺序以改善内存局部性和性能。B 侧收集器：由于我们遵循 2x1 分块格式，我们可以利用 B 侧收集器。通过在一个 MMA 上指定“FILL”，在下一个 MMA 上指定“LASTUSE”，我们可以将 B 的读取减少到只发生一次而不是两次。我们测得这大约有 1-3% 的提升。利用对于更大的 64k 和 128k 方阵 NVFP4 GEMM，我们观察到提前释放 A 分别带来了 13.5% 和 22.1% 的加速。我们发现在较大尺寸下这条指令很有用，此时 A 分块与其他资源争夺驻留空间，导致其行在重用之间被逐出，并迫使加载器等待它们。这使我们随后能够实现提前重用 A 带来的收益。在较小尺寸下，A 分块从不离开 L2，这意味着读取已经足够快，无需提前释放 A。为了利用这条指令，我们修改了传统的环形顺序逻辑。在普通 GEMM 中，A 和 B 分块属于同一个环，并在单次提交下同步运行。然而，为了让提前释放 A 生效，我们需要解耦这两个分块，使 A 的加载能够独立运行。提前释放 A 需要在更早的信号上释放 A 的槽位，因此我们给 A 自己的环以及自己的到达/完成屏障对。sync_restrict::shared::read::mma::a
：L2 逐出提示：我们用 EVICT_LAST 标记 A 操作数，以鼓励 L2 驻留，供后续重用它们的工作使用。重用收益来自跨工作而非簇内，并带来零点几个百分点的帮助。
结果：
我们注意到，上述所有测量均是在 Qualification Sample（QS）GPU 上使用 NVIDIA CUDA 13.4 完成的。我们预计，随着 Vera Rubin 软件版本的发布，所有基线的性能都将持续提升。
结论：
我们希望你觉得其中一些内容有用，并且我们很期待大家很快开始上手使用它们。从 LUT GEMM、硬件原生 megakernel 到新的引擎优化，还有大量有趣的片段可以分享。更多内容即将发布！
Together AI 的 kernels 和性能团队正在积极招聘！如果你想进一步了解这些 kernel，或与我们一起开发下一批更新，请联系 Simran 或 Dan！
- Simran：[email protected]
- Dan：[email protected]