返回首頁

AI 基礎設施

ThunderKittens 支援 Vera Rubin 新指令,NVFP4 GEMM 在樣品 GPU 突破 22 PFLOPS

Together AI 將開源 CUDA 核心框架移植至 NVIDIA Vera Rubin,重新安排資料重用與記憶體管線。16K 方形矩陣測試達 22.24 PFLOPS,但數字來自資格樣品與尚未成熟的 CUDA 13.4。

Authors of the study: Nicholas Kluge Corrêa Camila Galvão James William Santos Carolina Del Pino Edson Pontes Pinto Camila Barbosa Diogo Massmann Rodrigo Mambrini Luiza Galvão Edmund Terem Nythamar de Oliveira · CC BY 4.0 · Image source
zh-Hant

Together AI 於 9 月 10 日公開 ThunderKittens 的 Vera Rubin 支援,加入 Rubin 專用 PTX 指令及 NVFP4、FP8 GEMM。舊 Blackwell 核心雖能直接執行,卻只能達到 NVFP4、FP8 理論 roofline 的 42.1% 與 44.4%;問題不在 Tensor Core 算力不足,而是既有資料路徑無法跟上更快的運算單元。

團隊首先把 `tcgen05.mma` 的 K 步幅由 32 bytes 擴至 64 bytes,使每個相同週期的指令處理兩倍資料。Rubin 的 Tensor Memory 也從 512 增至 576 columns,配合 `.exclusive` 配置後,可容納兩組累加器及量化 scale;核心因此能採用 2×1 tile,讓兩個輸出區塊重用同一份 B 資料。新增的 B-side collector 可把 B tile 暫存在 MMA staging buffer,減少重複讀取,官方量測單獨帶來約 1% 至 3% 改善。

另一項關鍵是 oversized shared-memory 模式:每個 SM 的可用量由 228 KiB 增至 328 KiB,讓核心能預載更多 tile、加深環形管線。在 16K 方形 NVFP4 GEMM 中,管線由三階段加至五階段後,效能從 17,054 升至 22,239 TFLOPS;FP8 五階段則達 11,995 TFLOPS。ThunderKittens 已把 Rubin 支援併入公開儲存庫,為不依賴封閉函式庫的核心開發提供可讀範例。

不過這不是端到端模型吞吐量,也不能直接換算為推論成本。結果只涵蓋特定 GEMM 形狀,且在 Qualification Sample GPU、CUDA 13.4 上取得;功耗、跨 GPU 通訊、較小或不規則矩陣,以及正式量產硬體上的 cuBLAS 改進,都可能改變相對表現。工程團隊下一步應關注正式 CUDA 工具鏈、更多核心類型與實際 MoE/推論工作負載的重現結果。

來源

  1. To Infinity and Beyond: ThunderKittens Now on NVIDIA Vera Rubin NVL72!
  2. ThunderKittens: Tile primitives for speedy kernels