- Together AIがThunderKittensにNVFP4/FP8 GEMMのVera Rubin対応を追加
- Vera RubinではKステップが64バイトに倍増し、テンソルメモリは576列に拡大
- 2x1タイル化とパイプライン深化でNVFP4は22PFLOPS超に到達
何が発表されたのか
Together AIのカーネルチームは、NVIDIAの新しいデータセンター向けプラットフォーム「Vera Rubin NVL72」へのアクセスを得て、数日間かけて新しい命令セット（ISA）を調査したと報告しています。その成果として、GPUカーネルライブラリであるThunderKittensに、Vera Rubin上でNVFP4とFP8のGEMMを書くための機能を追加しました。GEMMは行列積のことで、AIの学習や推論で最も計算量の大きい処理の一つです。
RubinはBlackwellのプログラミングモデルを引き継いでいるため、既存のBlackwell向けカーネルはそのまま動作します。ただし、そのまま実行した場合、NVFP4カーネルは理論性能（ロフライン）の約42.1%、FP8カーネルは約44.4%しか出ず、最適化の余地が大きいとしています。元のBlackwell向けカーネルは、スレッドブロッククラスタによるTMAマルチキャストでのオペランド共有、ワープの役割分担、常駐実行といった工夫を前提にしていました。
Vera Rubinの新機能
テンソルコアのKステップが2倍になりました。tcgen05.mma命令はMxNxKのタイルに対してC = A@B + Cを計算し、K方向に一定バイト数を消費します。Blackwellでは32バイトだったこのステップが、Vera Rubinでは64バイトまで増やせます。命令自体の所要サイクルは変わらないため、同じ命令ウィンドウにより多くの計算を詰め込めるようになります。ThunderKittensではmma_ABt<64>(...)のようにテンプレートパラメータで指定します。
テンソルメモリはBlackwellの512列から576列に増え、32KiB分の余裕が生まれました。ただし追加された列はPTX 9.4で導入された.exclusive修飾子を通じてのみ到達可能で、SMあたり1つの有効な割り当てに限られます。共有メモリも228KiBから328KiBへと拡大し、ホスト側の設定で有効化します。
さらに、Bタイル側のコレクタバッファ（MMAの一時ステージング領域）と、Aタイルを早く解放する命令が加わりました。コレクタではFILL／USE／LASTUSE／DISCARDの4つのラベルで再利用の権限を指定します。Early A releaseは、MMAがAオペランドを共有メモリから読み終えた時点でバリアを発火させ、次のステージのロードを早く開始できるようにするものです。
性能を引き出す工夫
32バイトのKステップのままでは、ISA上の上限が約16.8PFLOPSで、既存のNVFP4カーネルはすでに14.7PFLOPS（上限の88%）に達していました。そこで64バイトのKステップに切り替えますが、単に切り替えるだけでは性能はほとんど向上しません。テンソルコアがオペランドを消費する速度は倍になっても、供給する速度は変わらないためです。
対策として、同じCTAペアにM方向へ2つ目の出力タイルを重ねる2x1タイリングを採用し、Bチャンクを共有して転送量を減らしました。Blackwellではテンソルメモリが512列に制限され、2つのM256xN256アキュムレータで使い切ってしまうため、この形は取りにくかったとしています。Vera Rubinの追加64列によってスケーリング係数を格納できるようになり、この構成が可能になりました。
並行して、拡大した共有メモリを使ってパイプラインを深くし、タイルを先読みして転送時間を隠蔽しました。加えてB側コレクタでBの読み出しを1回に減らし（1〜3%の改善）、大規模な64k・128kのNVFP4 GEMMではEarly A releaseによりそれぞれ13.5%、22.1%の高速化を確認しています。AオペランドへのEVICT_LASTヒントもわずかな改善に寄与したとしています。
結果と残る不明点
こうした最適化の結果、NVFP4カーネルは22PFLOPS超に達し、cuBLASやCuTE DSLと競合する性能になったと報告されています。測定はNVIDIA CUDA 13.4と、Qualification Sample（QS）と呼ばれる評価用GPUで行われたとのことです。
一方で、参照した原文ではベースライン性能の詳細な比較が途切れており、各最適化の寄与を完全には確認できません。また、ThunderKittens側の該当機能がどのバージョンから利用できるのか、一般の開発者がVera Rubin NVL72にアクセスできるのかは明記されていません。ロフライン比などの具体的な数値も、示されている範囲に限られます。
| 項目 | Blackwell | Vera Rubin |
|---|---|---|
| テンソルコアのKステップ | 32バイト | 64バイト |
| テンソルメモリ | 512列 | 576列 |
| 共有メモリ | 228 KiB | 328 KiB |
| コレクタバッファ | Aタイルのみ | AタイルとBタイル |
The core issue we find is that while Rubin enables tensor cores to consume operands twice as fast, our old Blackwell kernel does not feed them fast enough.
私たちが見つけた核心的な問題は、Rubinがテンソルコアに2倍の速さでオペランドを消費させる一方で、既存のBlackwellカーネルが十分な速さでそれを供給できていないという点です。
今後の見通し
Together AIの報告は、Vera RubinがBlackwellの延長線上にあることを示す一方、新しい命令を有効化するだけでは性能が伸びず、タイル形状やパイプライン設計の見直しが不可欠であることも示しています。今後は、cuBLASやCuTE DSLとの詳細なベンチマーク比較、FP8を含む各カーネルの最終的なロフライン比、そしてThunderKittensの該当機能が一般に公開される時期が明らかになれば、実務で採用できるかどうかを判断しやすくなるでしょう。Vera Rubin NVL72自体の提供時期や価格が公表されれば、導入検討の材料になるとみられます。
日本の開発者・IT企業にとっての意味
この内容は、日本のIT企業や開発者にとって二つの意味を持ちます。一つは、GPUの世代が変わっても、既存のカーネルがそのまま動く保証はあっても性能が自動で上がるわけではないという点です。Vera RubinはBlackwellのプログラミングモデルを維持していますが、Kステップを64バイトに広げる、タイル形状を2x1に変える、パイプラインを深くするといった作業を経て初めて22PFLOPS超に到達しています。自前で推論・学習基盤を運用する企業は、新しいGPUを導入しても既存カーネルのままだと理論性能の半分以下しか出ない可能性を考慮すべきです。もう一つは、こうした最適化ノウハウがThunderKittensのようなライブラリに蓄積されることで、cuBLASやCuTE DSLに依存しない選択肢が増えるという点です。ただし対応状況や公開時期は不明なため、当面はベンダー製ライブラリとの併用を前提に検討するのが現実的とみられます。
気になる点
- 既存のBlackwell向けカーネルはVera Rubinでもそのまま動きますか
- はい。RubinはBlackwellのプログラミングモデルを維持しているため、既存カーネルはそのまま動作すると原文にあります。ただし、そのままではNVFP4でロフラインの約42.1%、FP8で約44.4%にとどまり、性能を引き出すにはKステップの拡大やタイル形状の変更といった最適化が必要です。
- ThunderKittensのVera Rubin対応は誰でも使えますか
- 現時点では公表されていません。原文はTogether AIのカーネルチームがVera Rubin NVL72へのアクセスを得て機能を追加したと述べるにとどまり、対応バージョンや公開時期、ライセンス条件には触れていません。利用可否は今後の公式発表を待つ必要があります。
- 性能の数値はどの環境で測ったものですか
- NVIDIA CUDA 13.4と、Qualification Sample（QS）と呼ばれる評価用GPUで測定されたと原文にあります。16k平方のGEMMでのパイプライン深さの検証や、64k・128k平方でのEarly A releaseの効果などワークロード別の数値が示されていますが、原文ではベースラインの詳細な比較は確認できません。
用語解説
- GEMM
- 行列積（General Matrix Multiply）の略。AIの学習・推論で中心となる計算で、GPU性能を測る指標としても使われる。
- テンソルコア
- 行列積を高速に実行するGPU内の専用演算回路。世代ごとに命令の出し方やデータの置き場所が変わる。
- ロフライン
- メモリ帯域や演算器の上限から算出される理論性能の上限値。実測性能をこの値で割った比率をロフライン比と呼ぶ。
- テンソルメモリ
- Blackwellで導入された、テンソルコアが直接読み書きできる専用メモリ領域。アキュムレータをレジスタの代わりに置く。
- TMA
- Tensor Memory Accelerator。GPUのメモリ転送を専用ハードウェアで非同期に行う仕組み。
- NVFP4
- 4ビットの浮動小数点形式。FP8よりさらに小さいデータ型で、量子化によりメモリ消費と転送量を削減できる。
出典
To Infinity and Beyond: ThunderKittens Now on NVIDIA Vera Rubin NVL72!
https://together.ai/blog/to-infinity-and-beyond-thunderkittens-now-on-nvidia-vera-rubin-nvl72
AI導入も顧問も、お任せください
何から手をつけるか、どこまでAIに任せるか。自社の業務に合わせて整理し、導入から運用まで伴走します。相談だけでも構いません。
AI導入について相談する