Together AI Blog

To Infinity and Beyond: ThunderKittens Now on NVIDIA Vera Rubin NVL72!

8.5内容质量

TL;DR · AI 摘要

Together AI在NVIDIA Vera Rubin NVL72上优化ThunderKittens，实现22 PFLOPS性能，接近cuBLAS。

核心要点

  • Vera Rubin NVL72的FP8算力达17.5 PFLOPS/GPU，是Blackwell的3.89倍
  • 通过TMA multicast和warpgroup优化，内存带宽利用率提升50%
  • 优化后NVFP4内核达到22 PFLOPS，接近cuBLAS性能

结构提纲

按章节快速跳转。

  1. §引言

    介绍ThunderKittens在NVIDIA Vera Rubin平台的部署背景与目标

  2. ·Blackwell架构GEMM编程模型

    解析Blackwell第五代Tensor Core的tcgen05指令特性

  3. ›Vera Rubin新特性

    对比Blackwell与Vera Rubin的算力与内存带宽提升

  4. ·性能优化方法

    通过TMA multicast和warpgroup优化实现性能突破

  5. ›优化结果

    NVFP4内核达到22 PFLOPS，接近cuBLAS性能

思维导图

用一张图看清主题之间的关系。

查看大纲文本（无障碍 / 无 JS 友好）
  • ThunderKittens在Vera Rubin的优化
    • Blackwell架构特性
      • tcgen05指令单线程执行
    • Vera Rubin新特性
      • 35 PFLOPS/GPU算力
      • 22 TB/s内存带宽
    • 优化方法
      • TMA multicast
      • warpgroup优化

金句 / Highlights

值得收藏与分享的关键句。

#NVIDIA Vera Rubin#GEMM优化#ThunderKittens#FP8#NVFP4
打开原文

无限延伸：ThunderKittens 现已部署在 NVIDIA Vera Rubin NVL72 上

Together 的内核团队最近获得了 NVIDIA Vera Rubin NVL72 平台的访问权限。我们过去几天一直在研究新的 ISA 并用显微镜观察芯片。新平台有许多令人兴奋的功能！我们已经将部分功能集成到 ThunderKittens 中，实现了在 Vera Rubin 上编写 NVFP4 和 FP8 GEMMs 的能力，并帮助其他小猫探索星空。

天文学家小猫。太可爱了！

在深入探讨 Vera Rubin 的能力之前，我们先快速回顾一下 Blackwell GPU 的 GEMM 实现。

起点：NVIDIA HGX B200 GEMM

NVIDIA Blackwell 架构的第五代张量核心彻底改变了 GEMM 编程模型。虽然 NVIDIA Hopper 架构的 wgmma 指令由 warp group 集体发出，但 Blackwell 的 tcgen05 指令由单个线程发出，使单个生产者 warp 能够驱动张量核心。累加器也从寄存器移动到了张量内存，操作数直接从共享内存读取，使单个 MMA 能够跨两个 SM 的两个 CTA 进行扩展。

为了在 Blackwell 上达到竞争力的性能，我们的 GEMM 实现了以下优化：

  • 启动线程块集群，使每个 CTA 对可以通过 TMA 广播共享操作数，将 HBM 的内存流量减少一半。
  • 在集群内部对 warp 进行专业化分工：加载器通过 TMA 将 A 和 B 读入共享内存，单个 MMA warp 驱动张量核心，消费者 warp group 将完成的累加器从张量内存传输到 HBM。
  • 持续运行，一个 tile 的输入在流式传输时，前一个 tile 的输出仍在被处理。

通过这些优化，我们实现了以下性能提升。

更多关于这些内核及其优化的细节，请参阅我们之前的 Together 博客文章或 ThunderKittens 2.0 发布说明！

由于 Vera Rubin 保留了 Blackwell 的编程模型，我们的旧版 GEMM 仍然可以运行。然而，直接在 Vera Rubin 上运行时，我们发现 NVFP4 和 FP8 内核的性能仅达到理论上限的 42.1% 和 44.4%——这为优化留下了巨大空间！

本文的其余部分分为两部分。首先，我们将介绍 Vera Rubin 对 GEMM 重要的新特性，并说明如何在 ThunderKittens 中使用它们。然后，我们将逐步将这些特性集成到现有的 Blackwell NVFP4 内核中，使其性能超过 22 PFLOPS，并与 cuBLAS 和 CuTE DSL 竞争。

我们发现的核心问题是：虽然 Vera Rubin 使张量核心能够以两倍的速度消耗操作数，但我们的旧版 Blackwell 内核无法以足够快的速度为其提供数据。为了达到计算上限，我们需要让 tile 从已经在片上的数据中提取更多重用。

NVIDIA Vera Rubin 平台的新特性

对比厂商规格，我们发现从 Blackwell 到 Vera Rubin 的以下改进。

| NVIDIA HGX B200 | NVIDIA Vera Rubin NVL72 | |----------------|------------------------| | NVFP4 张量核心 | 9 PFLOPS/GPU | 35 PFLOPS/GPU | | FP8 张量核心 | 4.5 PFLOPS/GPU | 17.5 PFLOPS/GPU | | FP16/BF16 张量核心 | 2.25 PFLOPS/GPU | 4 PFLOPS/GPU | | 内存带宽 | 8 TB/s/GPU | 22 TB/s/GPU | | SM 数量 | 148/GPU | 224/GPU | | 峰值功耗 | 1000W/GPU | 2300W/GPU |

就编写高性能 GEMM 而言，我们特别关注以下特性。

1. 张量核心的 K 值翻倍

为了便于审查，tcgen05.mma 在 MxNxK 的瓦片上计算 C = A@B + C，每步沿 K 维消耗固定字节数。在 Blackwell 上，这一步为 32 字节，但在 Vera Rubin 上可提升至 64 字节。MMA 本身的周期数保持不变，因此双倍 K 步允许我们在相同指令窗口中装入两倍的工作量。

在 ThunderKittens 中，我们通过向现有 mma 操作添加新的模板参数来实现这一点。

cpp
mma_ABt    (. ..);   // Blackwell 默认：32 字节 K 步
mma_ABt<64>(...);   // Vera Rubin：64 字节 K 步

2. 张量内存扩展至 576 列

Blackwell 引入了张量内存的概念，这是一个 128 通道 x 512 列 x 32 位的存储空间，张量核心可以直接读写。在 Vera Rubin 上，该空间扩展至 576 列，额外提供了 32 KiB 的张量内存空间。

需要注意的是，这些额外列只能通过 .exclusive 限定符访问，这是 PTX 9.4 的新增特性，确保每个 SM 上只有一个活跃的张量内存分配。非独占分配仍限制在 512 列且必须是 2 的幂次。

在 ThunderKittens 中，用户可通过向张量内存分配器添加模板参数，指定分配为独占模式来请求此功能。

cpp
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 默认：512 列
tensor_allocator<1, C::CLUSTER_SIZE, false, true>  tm; // Rubin：最多 576 列

3. 共享内存扩展至 328 KiB

虽然 Hopper 和 Blackwell 提供了 228 KiB 的共享内存，Vera Rubin 引入了超大共享内存模式，可动态扩展至 328 KiB。这是一个主机端指定功能，调用方式如下。

cpp
CUfunction function = nullptr;
cudaGetFuncBySymbol(&function,reinterpret_cast
(kernel));
cuFuncSetAttribute(function, CU_FUNC_ATTRIBUTE_SHARED_MEMORY_MODE,
                   CU_SHARED_MEMORY_MODE_ALLOW_OVERSIZED_SHARED_MEMORY);

4. B 侧收集器

Blackwell 引入了收集器缓冲区的概念，这是一个小型 MMA 中间缓冲区，可捕获 A 瓦片以便后续指令直接从缓冲区读取而非共享内存。Vera Rubin 通过 .collector::b::* 将此功能扩展至 B 瓦片。

要使用此功能，我们需要为每个 MMA 的操作数添加四个标签之一，描述其对收集器缓冲区的操作。

  • “FILL”：从共享内存读取操作数并捕获
  • “USE”：从缓冲区读取
  • “LASTUSE”：从缓冲区读取并释放
  • “DISCARD”：默认行为，跳过捕获

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

现在两个操作数均可驻留于收集器缓冲区，我们可以尝试新的模式。例如，在 2x2 块中，双侧收集可将 8 次操作数读取所需的 4 个 MMA 减少至仅需 5 次。

在 ThunderKittens 中，我们可通过以下方式暴露此功能：

cpp
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 时，表示该 MMA 已完成处理，向生产者发出信号表明有一个阶段插槽可供重复使用。PTX 9.4 引入了 tcgen05.commit.sync_restrict::shared::read::mma::a 这一新指令，允许我们提前释放 A 瓷砖的信号。无需等待 MMA 完成处理，我们可以在 MMA 从共享内存读取完 A 操作数后立即触发屏障，从而通知 tma loader 开始存储下一阶段的内存。

ThunderKittens 为用户引入了一种新的提交类型来实现这一功能。

code
tensor_commit<2> (inputs_finished[stage], mask); // 当 MMA 退栈时到达
tensor_aread_commit<2>(A_finished[slot], mask); // 当 MMA 完成读取 A 时到达

构建 GEMM

我们现在拥有新功能和 Blackwell GEMM。以下部分将逐步将这些功能整合到现有的 Blackwell 内核中，并解释在转向 Vera Rubin 时为何需要这些功能。

指令扩展：

最直观的瓶颈仍然来自于对 Blackwell 的 32 字节 K 步的依赖。在 Vera Rubin 上，这种编码的 ISA 限制约为 16.8 PFLOPs，而我们的 NVFP4 Blackwell GEMM 原生已达到 14.7 PFLOPs（达到上限的 88%）。要实现进一步突破，我们必须将 MMAs 处理的 K 字节数翻倍。

然而，仅在现有 Blackwell 内核中启用更宽的编码时，我们发现性能仅略有提升，而非预期的 2 倍。虽然更宽的 MMA 可使张量核心消耗操作数的速度翻倍，但它无法提升我们供应操作数的速度。为了使双 K 有效，我们通过两个杠杆来保持核心的高效运行：减少传输字节数并加深流水线以确保这些获取操作能够重叠。

读取更少字节：

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

由于 Blackwell 的张量内存上限为 256 KiB，我们的 NVFP4 内核难以实现 2x1 瓷砖格式。两个 M256xN256 累加器已占用 512 列，这意味着块缩放的 MMA 没有空间存储其 A 和 B 缩放因子。虽然程序员可以通过让 epilogue 线程在发出 MMA 空闲状态前仅加载部分累加器列来绕过这一限制，从而允许下一 K 瓷砖的 MMA 开始运行，但这种方法引入了一定的延迟且无法隐藏。幸运的是，Vera Rubin 的额外 64 列使我们无需这种操作即可存储缩放因子。

深化流水线：

改变瓷砖格式减少了操作数流量，但无法缩短每次获取操作所需的时间。下一个挑战是保持张量核心的持续供应。Vera Rubin 更大的共享内存使我们能够创建更深的流水线，提前准备更多瓷砖，并为传输提供更多完成时间。对我们的 NVFP4 和 FP8 16k 方形 GEMM 进行环深度测试，结果如下：

NVFP4 16k 方形 GEMM：

共享内存瓷砖阶段

所需共享内存

实现的 TFLOPS

3

202 KiB

17,054

4

258 KiB

20,595

5

314 KiB

22,239

FP8 (E4M3) 16k Square GEMM:

Config

209 KiB

10,895

257 KiB

11,995

6

305 KiB

11,288

尽管最大的收益来自于这个最终的调整点，但需要指出这些收益依赖于我们之前的优化。以下是独立调整K步长、共享内存流水线和分块策略后的结果。

最终优化：

为了进一步提升内核性能，我们对一些关键参数进行了微调：

  • 内核配置调优：我们针对不同工作负载进一步优化内核配置。关于CTA对尺寸，我们发现Blackwell内核原始的1x1分块格式在处理较小的正方形工作负载时表现最佳。对于更大的形状，我们采用2x1 CTA对分块，并配合经过调优的集群尺寸（2、4或8个CTA）。此外，在所有形状中，我们通过调整分块光栅化顺序来提升内存局部性和性能。
  • B侧收集器：由于我们采用2x1分块格式，可以利用B侧收集器。通过在第一个MMA上指定“FILL”，在下一个MMA上指定“LASTUSE”，可以将B读取次数从两次减少到一次。我们测量到这一优化带来了约1-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分配独立的环形队列和独立的到达/完成屏障对。
  • L2逐出提示：我们使用EVICT_LAST标记A操作数，以鼓励L2缓存保留供后续重复使用的工作负载。这种重用收益跨工作负载而非集群内部，可带来约0.1-0.3%的性能提升。

结果：

需要说明的是，以上所有测量数据均使用NVIDIA CUDA 13.4在资格样本（QS）GPU上完成。我们预计随着Vera Rubin软件的发布，所有基准性能将持续提升。

结论：

希望您能从本文中获得一些有价值的见解，我们期待大家尽快开始尝试这些优化。从LUT GEMM、硬件原生兆内核到新的引擎优化，仍有大量有趣的细节值得分享。相关内容将很快发布！

Together AI的内核和性能团队正在积极招聘！如果您想了解更多关于这些内核的信息，或希望与我们合作开发下一阶段的更新，请联系Simran或Dan！