cuTile Rust: 安全でデータレースのないGPUカーネル

cuTile Rust (cutile-rs) は、Rust の所有権と借用規則を GPU の起動境界を越えて拡張することで、メモリ安全でデータレースのない GPU カーネルの開発を可能にします。ホスト側で可変テンソルを互いに素な断片に分割することにより、システムは各 GPU タイルプログラムが特定のデータセグメントに対して排他的な可変参照を受け取ることを保証し、コンパイル時にデータレースを排除します。

メモリ安全性とタイルベースのモデル

cuTile Rust は、手動での共有メモリ管理や unsafe コードを必要とすることが多い従来の SIMT (Single Instruction, Multiple Threads) プログラミングを、安全なタイルベースのプログラミングモデルに置き換えます。

起動境界を越えた所有権

システムは、厳格なアクセス規則を強制することで Rust の安全性保証を維持します：

  • 可変テンソル (Mutable Tensors): これらはカーネルが起動される前に、互いに素な断片に分割されます。各タイルには、割り当てられたパーティションに対する排他的な &mut ビューが許可されます。
  • 不変テンソル (Immutable Tensors): これらはすべてのタイル間で読み取り専用の & 参照として共有されます。
  • ランチャ (Launchers): 生成されたランチャは、GPU の作業が実行されている間も所有権を維持し、同期起動、非同期パイプライン、および CUDA graph replay をサポートします。

CUDA Tile IR による JIT コンパイル

カーネルは #[cutile::module] マクロを使用して定義され、これは Rust AST をキャプチャしてホストバイナリに埋め込みます。実行時に、cuTile Rust はこの AST を CUDA Tile IR を通じて GPU cubin へ JIT コンパイルします。安全なサーフェス API がほとんどの操作を処理しますが、Tile IR 操作サーフェスへの直接アクセスが必要な場合は、unsafe イントリンジックを使用して低レベルの制御が可能です。

パフォーマンス・ベンチマーク

NVIDIA B200 で実施された評価では、cuTile Rust は測定可能な実行時オーバーヘッドなしに安全性を提供し、手動で最適化された低レベル実装と同等のパフォーマンスレベルに達することが示されています。

  • 要素ごとの操作 (Element-wise Operations): 7 TB/s に達し、ピークメモリ帯域幅の約 91% を記録しました。
  • GEMM (General Matrix Multiply): 2 PFlop/s に達し（高密度 f16 ピークの 92%）、低レベル Tile IR バリアントの 0.3% 以内の性能で動作し、cuBLAS と同等の競争力を維持しています。
  • 推論パフォーマンス: Hugging Face との提携により、Grout 推論エンジン（cuTile Rust で構築）は、RTX 5090 上の Qwen3-4B で 171 tokens/s、B200 上の Qwen3-32B で batch-1 decode 中に 82 tokens/s を達成しました。

技術要件と互換性

cuTile Rust は、計算能力 sm_80 (Ampere) 以上の NVIDIA GPU 用に設計されています。

ハードウェアおよびソフトウェアの依存関係

  • GPU アーキテクチャ: 最小 sm_80。sm_100+ のサポートには CUDA 13.1+ が必要であり、sm_8x には CUDA 13.2、sm_90 には CUDA 13.3 が必要です。
  • CUDA Toolkit: FP4 パッキングや block-scaled MMA を利用するために、バージョン 13.3 が推奨されます。
  • Rust バージョン: 1.89+。
  • OS: Linux (Ubuntu 24.04 でテスト済み)。

他の Rust GPU プロジェクトとの比較

cuda-oxide のような他のプロジェクトが SIMT スタイルのカーネルに対して Rust から CUDA へのコンパイルを提供しているのに対し、cuTile Rust はタイルベースのアプローチに焦点を当てています。このトレードオフにより、明示的な warp プリミティブや手動の共有メモリ制御を放棄する代わりに、コンパイル時に安全性を検証可能なセマンティクスモデルを提供します。

プロジェクトのステータスとエコシステム

cuTile Rust は、現在、活発な開発が行われている初期段階の研究プロジェクトです。ユーザーは API の破壊的変更や機能の不完全さを想定しておく必要があります。プロジェクトは crates.io で cutile として利用可能であり、環境構築を簡微化するために Nix flake でサポートされています。

ワークスペース構成

  • cutile: カーネルの作成と実行のための主要なユーザー向け crate。

  • cutile-compiler: cutile-ir を介して Rust カーネルを実行ファイルへコンパイルする処理を担当します。

  • cuda-async / cuda-core: それぞれ非同期 CUDA 実行および、慣用的な安全な CUDA API を提供します。

Sources