NVIDIA、GPU向けCUDAタイルカーネルをPythonからRustへ移植するAIエージェント手法を公開

要点

• NVIDIAは、GPU向けカーネルを「cuTile Python」や「Triton-TileIR」から安全性の高い「cuTile Rust」へ自動変換するAIエージェントのワークフローを発表した。

• 共通の中間表現「CUDA Tile IR」を活用し、コード構造の等価性をIRの差分比較によって機械的に検証しながら移植を進める仕組みを構築した。

• カーネルライブラリ「TileGym」の公開オペレータ全24種（約40カーネル）を移植し、NVIDIA DGX B200環境でcuTile Python比平均99.5%の実行性能を達成した。

• 開発された変換用エージェントスキルや移植済みコードは、オープンソースのTileGymリポジトリを通じて公開されている。

• 米NVIDIAは2026年9月16日、GPUカーネル記述言語である「cuTile Python」および「Triton-TileIR」のコードを、Rust言語向けの「cuTile Rust」へと変換するAIエージェント技術の詳細を公式技術ブログで発表した。同社が提供するカーネルライブラリ「TileGym」の全24オペレータを対象に移植を実施し、参照元のPython実装とほぼ同等となる平均99.5%の性能を達成したという。本手法を実現するエージェントスキルは、GitHubのTileGymリポジトリ上で一般に公開されている。

共通中間表現「CUDA Tile IR」を軸にした確実な移植基盤

GPUプログラミングにおいて、メモリ安全性や所有権モデルを重視したRust言語の導入が進められている。NVIDIAの研究チームが開発した「cuTile Rust（cutile-rs）」は、Rustの所有権モデルをタイル単位のGPUカーネル（GPU上で並列実行される処理単位）記述に拡張し、可変な出力データを重複しない領域へ分割しながらホスト側の所有権規約を保つシステムである。低レベルな制御が必要な場合にはローカルでオプトアウトし、直接IR（中間表現）操作を実行することも可能だという。

一方で、同社のカーネルライブラリ「TileGym」には、これまで「cuTile Python」や「Triton-TileIR（nvtriton）」で書かれた多数のプロダクション向けカーネルが蓄積されていた。これらをRustへ移行するにあたり、チームは再最適化を一から行うのではなく、3つの言語フロントエンドが同一の中間表現「CUDA Tile IR（cuda_tileダイアレクト）」を共有している点に着目した。

すべてのフロントエンドは同じ「tileiras」コンパイラを通じてGPUバイナリを出力するため、コンパイラ最適化や性能特性の前提は共通している。そのため、今回の変換は「同じタイル処理を安全なホスト言語で再表現する」作業となり、生成されたIRの構造（メモリ操作ファミリ、タイル形状、リダクション演算など）が一致しているかを機械的に比較・検証できる基盤が整っていたと説明されている。

機械的判定を挟むマルチエージェントの変換パイプライン

変換プロセスには、各段階で検証判定を行う限定的なマルチエージェントワークフローが採用された。処理は大きく「分析」「デバイスカーネルの生成」「ホストおよびFFIコードの生成」「ベンチマーク検証」という4つのステージに分かれている。

各ステージの終了時には、検証スクリプトによる判定や、参照元実装とのTile IR差分比較（IR diffing）が実行される。機能テストを実施する前の段階でIRの構造的等価性を確認し、明確な合格判定が出た場合のみ次のステージへ進むアーキテクチャとなっており、AIモデルの出力を盲信せずに機械的チェックで担保する設計が採られた。

言語間の大きな技術的ギャップとして、カーネルの「特化（Specialization）」の扱いが挙げられている。cuTile PythonではJIT（実行時）コンパイルによって呼び出し時に暗黙的にカーネルが特化されるのに対し、cuTile Rustでは型やテンソルの形状制約をコンパイル時にRustコンパイラ（rustc）へ明示的に認識させる必要がある。そのため、カーネルのシグネチャ（関数定義）において明示的な特化宣言を行うコード変換が施された。

また、変換後のRustカーネルは、データのコピーや新たなメモリ割り当てを発生させずにテンソル記述子を受け渡すC-ABI（言語間インターフェース）レイヤーを介してTileGymと統合されている。これにより、遅延コンパイルをサポートしつつ、既存の開発サイクルを損なわない連携が可能になったとしている。

DGX B200で99.5%の性能を実証、エージェントスキルも一般公開

本システムを適用した実証実験では、TileGymに含まれる24個の公開オペレータ（合計約40個のGPUカーネル）の移植が行われた。対象には、単純な要素ごとの演算（element-wise）から、Flash-Attention decode、Multi-head Latent Attention（MLA）、さらにはMixture-of-Experts（MoE）モデルに至るまで、多様な演算が含まれている。

最新のAIアクセラレータシステム「NVIDIA DGX B200」を用いて行われたベンチマーク測定では、移植された全24オペレータがcuTile Python比で幾何平均0.95倍以上の性能しきい値を達成した。全体の平均性能比（幾何平均）は0.995に達しており、安全性の高いRustへの移行でありながら、元のPython実装に対して実質的な性能パリティ（同等性）を維持していることが確認された。

NVIDIAは、今回構築した変換用エージェントスキル「tilegym-converting-python-to-rust」および変換済みのcuTile RustオペレータをGitHubのTileGymリポジトリで公開しており、開発者が自身で保有するカーネルのRust化にも応用可能であると案内している。