AI 基礎設施
ThunderKittens 支援 Vera Rubin 新指令,NVFP4 GEMM 在樣品 GPU 突破 22 PFLOPS
Together AI 將開源 CUDA 核心框架移植至 NVIDIA Vera Rubin,重新安排資料重用與記憶體管線。16K 方形矩陣測試達 22.24 PFLOPS,但數字來自資格樣品與尚未成熟的 CUDA 13.4。

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/推論工作負載的重現結果。