ThunderKittensがVera Rubinの新命令に対応、NVFP4 GEMMがサンプルGPUで22 PFLOPSを突破

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へ拡大し、各命令が同じサイクル数で2倍のデータを処理できるようにした。RubinのTensor Memoryも512 columnsから576 columnsへ増加しており、`.exclusive`構成と組み合わせることで、2組のアキュムレータと量子化scaleを格納できる。これによりカーネルは2×1 tileを採用し、2つの出力ブロックで同じBデータを再利用できるようになった。新たに追加されたB-side collectorは、B tileをMMA staging bufferに一時保存して重複読み出しを削減する。公式の測定では、この変更単独で約1～3%の性能向上が得られた。

もう一つの重要な要素がoversized shared-memoryモードだ。SM当たりの利用可能容量が228 KiBから328 KiBへ増え、カーネルがより多くのtileをプリロードし、リングパイプラインを深くできるようになった。16K正方行列のNVFP4 GEMMでは、パイプラインを3段から5段へ増やすことで、性能が17,054 TFLOPSから22,239 TFLOPSへ向上した。FP8の5段パイプラインも11,995 TFLOPSを達成している。ThunderKittensはすでにRubin対応を公開リポジトリへマージしており、クローズドなライブラリに依存しないカーネル開発のための、読みやすい実装例を提供している。

ただし、これはエンドツーエンドのモデルスループットではなく、推論コストへ直接換算することもできない。結果は特定のGEMM形状だけを対象としており、Qualification Sample GPUとCUDA 13.4上で取得されたものだ。消費電力、GPU間通信、より小さい行列や不規則な行列、さらに量産版ハードウェアにおけるcuBLASの改善によって、相対性能は変わる可能性がある。エンジニアリングチームは今後、正式版CUDAツールチェーン、より多様なカーネル、実際のMoE／推論ワークロードにおける再現結果に注目すべきだ。