Together AI 在 NVIDIA Vera Rubin NVL72 上优化 ThunderKittens GEMM 至 22 PFLOPS

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 的 tile 上计算 C = A@B + C，每一步沿 K 消耗固定数量的字节。在 Blackwell 上，这一步是 32 字节，但在 Vera Rubin 上可以增加到 64 字节。MMA 本身仍然需要相同的周期数，因此 K 翻倍让我们能在相同的指令窗口内打包两倍的工作量。

在 ThunderKittens 中，我们通过为现有 mma 操作新增一个模板参数来表达这一点。

mma_ABt    (...);   // Blackwell Default: 32-byte K step
mma_ABt<64>(...);   // Vera Rubin: 64-byte K step

2. Tensor memory 增长到 576 列

Blackwell 引入了 tensor memory 的概念，这是一个 128 lane x 512 column x 32-bit 的空间，tensor core 可以直接读写。在 Vera Rubin 上，该空间增加到 576 列，额外提供了 32 KiB 的 tensor memory 可供使用。

注意，这些额外的列只能通过 .exclusive 限定符访问，这是 PTX 9.4 新增的特性，确保一个 SM 上只有一个活跃的 tensor memory 分配。非独占分配仍然上限为 512，且必须是 2 的幂。

在 ThunderKittens 中，用户可以通过向我们的 tensor memory 分配器传入模板参数来请求这一点，指定该分配为独占。

template<int _nblocks_per_sm, int _ncta, bool _managed = true, bool _exclusive = false>
struct tensor_allocator { .... }

tensor_allocator<1, C::CLUSTER_SIZE, false> tm; // Blackwell default: 512 columns
tensor_allocator<1, C::CLUSTER_SIZE, false, true>  tm; // Rubin: Up to 576 columns

3. Shared memory 增长到 328 KiB

Hopper 和 Blackwell 提供了 228 KiB 的 shared memory，而 Vera Rubin 引入了超大 shared memory 模式，可动态增加到 328 KiB。这是一个主机端规范，可按如下方式调用。

CUfunction function = nullptr;
cudaGetFuncBySymbol(&function,reinterpret_cast<const void*>(kernel));
cuFuncSetAttribute(function, CU_FUNC_ATTRIBUTE_SHARED_MEMORY_MODE,
                   CU_SHARED_MEMORY_MODE_ALLOW_OVERSIZED_SHARED_MEMORY);

4. B 侧 Collector

Blackwell 引入了 collector buffer 的概念，这是一个小型 MMA 暂存缓冲区，可以锁存到一个 A tile 上，使下一条指令从那里获取数据，而不是从 shared memory 获取。Vera Rubin 通过 .collector::b::* 将这一功能扩展到 B tile。

为了利用这一点，我们为每个 MMA 的操作数标注四个标签之一，描述其对 collector buffer 的操作。

• “FILL”从 shared memory 读取操作数并锁存
• “USE”从缓冲区读取
• “LASTUSE”从缓冲区读取并释放
• “DISCARD”是默认值，跳过锁存

这些标签是复用的权限限定符，而非保证。这意味着即使 Tensor Core 有复用权限，它仍可能重新加载矩阵。

现在任一操作数都可以驻留在 collector buffer 中，我们可以尝试新的模式。例如在一个 2x2 块上，两侧都进行收集可将四个 MMA 从八次操作数读取减少到仅五次。

在 ThunderKittens 中，我们可以通过以下方式暴露这一点。

mma2_ABt_chunk<64, false, false, collector::FILL,collector::DISCARD>(C[0][0], a0, b0, ...);
mma2_ABt_chunk<64, false, false, collector::LASTUSE, collector::FILL   >(C[0][1], a0, b1, ...);
mma2_ABt_chunk<64, false, false, collector::FILL,    collector::LASTUSE>(C[1][1], a1, b1, ...);
mma2_ABt_chunk<64, false, false, collector::LASTUSE, collector::DISCARD>(C[1][0], a1, b0, ...);

5. 提前释放 A

当 MMA 处理完成后，一个 tcgen05.commit 会到达 mbarrier，通知生产者某个 stage slot 可以复用了。PTX 9.4 引入了 tcgen05.commit.sync_restrict::shared::read::mma::a，这是一条新指令，允许我们为 A tile 提前发出信号。我们不必等待 MMA 完成，而是可以在 MMA 完成从 shared memory 读取其 A 操作数后立即触发 barrier，从而通知 tma loader 开始存储下一 stage 的内存。

ThunderKittens 引入了一种新的 commit 类型，供用户表达这一点。

tensor_commit<2> (inputs_finished[stage], mask); // arrives when MMA retires
tensor_aread_commit<2>(A_finished[slot], mask); // arrives when MMA finishes reading A

构建 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 内核开启更宽的编码，我们注意到性能只有微幅提升，而非预期的 2 倍。虽然更宽的 MMA 使张量核心消耗操作数的速度翻倍，但它无助于我们供应操作数的速度。为了让 double-K 发挥作用，我们拉动两个杠杆来让核心保持满足：移动更少的字节，并加深流水线以确保这些取数操作被重叠。

读取更少的字节：

为了取更少的字节，我们在同一对 CTA 上沿 M 维度堆叠第二个输出瓦片。由于两个累加器仅在 M 上不同，我们能够为它们共享同一个 B 块。我们最初的 Blackwell NVFP4 内核采用了 1x1 的瓦片格式，这意味着覆盖 M512xN256 的输出需要两个配对任务，每个任务独立传输自己的 B 副本。通过转向 2x1 格式，我们可以只取一次 B 并覆盖相同的输出，从而减少操作数流量。

这种 2x1 瓦片格式在我们的 NVFP4 Blackwell 内核中并不容易实现，因为张量内存被限制在 256 KiB。两个 M256xN256 累加器已经占用了 512 列，这意味着块缩放 MMA 没有空间存储它们的 A 和 B 缩放因子。虽然程序员可以通过让尾声 warp 仅加载累加器列的一个子集，然后发出 MMA 空闲状态信号，从而让下一个 K 瓦片的 MMA 开始，来绕过这一限制，但这引入了一小部分无法隐藏的延迟。幸运的是，借助 Vera Rubin 额外的 64 列，我们可以存储缩放因子，而无需进行这种周旋。

加深流水线：

改变后的瓦片格式减少了操作数流量，但无助于减少每次取数的耗时。下一个挑战是让张量核心持续获得供给。Vera Rubin 更大的共享内存允许我们创建更深的流水线，提前暂存更多瓦片，并给传输更多时间来完成。对我们的 NVFP4 和 FP8、16k 方阵 GEMM 进行环形深度扫描，我们观察到以下结果。

NVFP4 16k 方阵 GEMM：

FP8 (E4M3) 16k 方阵 GEMM：

虽然最大的增益似乎来自这最后一步，但我们注意到这些增益依赖于我们之前的优化。以下是独立扫描 K 步长、共享内存流水线和瓦片策略的结果。

最后的润色：

为了进一步推动我们的内核，我们尝试了几个最后的调节旋钮

1. 内核配置调优：我们进一步针对不同工作负载调优内核以获得最大性能。关于 CTA 配对大小，我们发现 Blackwell 内核原有的 1x1 瓦片格式在较小的方阵工作负载上性能最佳。对于更大的形状，我们采用 2x1 CTA 配对瓦片，以及调优后的 2、4 或 8 个 CTA 的簇大小。 此外，在所有形状上，我们排列瓦片光栅化顺序以改善内存局部性和性能。
2. B 侧收集器：由于我们遵循 2x1 瓦片格式，我们可以利用 B 侧收集器。 通过在一个 MMA 上指定“FILL”，在下一个上指定“LASTUSE”，我们可以将 B 读取减少到只发生一次而不是两次。我们测得这大约有 1-3% 的提升。
3. 利用 sync_restrict::shared::read::mma::a：对于更大的 64k 和 128k 方形 NVFP4 GEMM，我们观察到提前释放 A 分别带来了 13.5% 和 22.1% 的加速。我们发现该指令在较大尺寸下很有用，此时 A 瓦片会与其他资源争夺驻留空间，导致其行在重用之间被逐出，并迫使加载器等待它们。这使我们能够实现提前重用 A 所带来的收益。在较小尺寸下，A 瓦片从不离开 L2，这意味着读取已经足够快，无需提前释放 A。为了利用该指令，我们修改了传统的环形顺序逻辑。 在普通 GEMM 中，A 和 B 瓦片属于同一个环，并在单次提交下同步运行。然而，为了让提前释放 A 生效，我们需要解耦这两个瓦片，使 A 的加载能够独立进行。 提前释放 A 需要在更早的信号上释放 A 的槽位，因此我们给 A 分配了自己的环以及自己的到达/完成屏障对。
4. L2 逐出提示：我们用 EVICT_LAST 标记 A 操作数，以鼓励 L2 驻留，供后续重用它们的作业使用。重用收益来自跨作业，而非集群内部，并能带来零点几个百分点的提升。

结果：

我们注意到，上述所有测量均使用 NVIDIA CUDA 13.4 在 Qualification Sample（QS）GPU 上完成。我们预计所有基线的性能将随着 Vera Rubin 软件版本的发布而继续提升。

结论：

我们希望其中一些内容对您有用，并且我们很期待大家很快开始试用它们。从 LUT GEMM、硬件原生 megakernel 到新的引擎优化，还有大量有趣的片段可以分享。更多内容即将推出！

Together AI 的 kernels 和性能团队正在积极招聘！如果您想了解更多关于这些 kernel 的信息，或与我们一起开发下一组更新，请联系 Simran 或 Dan！

• Simran：[email protected]
• Dan：[email protected]