cuTile Rust: 安全且無資料競爭的 GPU Kernel

cuTile Rust (cutile-rs) 透過將 Rust 的所有權與借用規範擴展到 GPU 啟動邊界，實現了記憶體安全且無資料競爭的 GPU kernel 開發。藉由在主機端將可變張量（mutable tensors）分割成不相交的部分，系統確保每個 GPU tile 程式都能獲得其特定資料段的獨佔可變引用，從而在編譯時消除資料競爭。

記憶體安全與基於 Tile 的模型

cuTile Rust 以一種安全、基於 tile 的程式設計模型，取代了傳統的 SIMT (Single Instruction, Multiple Threads) 程式設計——這通常需要手動管理共享記憶體與使用 unsafe 程式碼。

跨啟動邊界的所有權

系統透過執行嚴格的存取規範來維持 Rust 的安全性保證：

  • 可變張量 (Mutable Tensors): 在 kernel 啟動前，這些張量會被分割成不相交的部分。每個 tile 都會獲得其分配分區的獨佔 &mut 視圖。
  • 不可變張量 (Immutable Tensors): 這些張量會作為唯讀的 & 引用在所有 tile 之間共享。
  • 啟動器 (Launchers): 生成的啟動器在 GPU 工作執行期間保留所有權，支援同步啟動、非同步管線以及 CUDA graph replay。

透過 CUDA Tile IR 進行 JIT 編譯

Kernel 定義使用 #[cutile::module] 巨集，該巨集會擷取 Rust AST 並將其嵌入到主機端二進位檔中。在執行時，cuTile Rust 會透過 CUDA Tile IR 將此 AST JIT 編譯成 GPU cubin。雖然安全的表面 API (surface API) 處理了大部分操作，但對於需要直接存取 Tile IR 操作層面的使用者，也可以透過 unsafe 原語 (intrinsics) 進行低階控制。

效能基準測試

在 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 合作，使用 cuTile Rust 建構的 Grout 推論引擎在 RTX 5090 上為 Qwen3-4B 達到 171 tokens/s，並在 B200 上於 batch-1 解碼時為 Qwen3-32B 達到 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: 建議使用版本 13.3，以利用 FP4 packing 與 block-scaled MMA 等功能。
  • Rust 版本: 1.89+。
  • 作業系統: Linux (已在 Ubuntu 24.04 上測試)。

與其他 Rust GPU 專案的比較

雖然像 cuda-oxide 這樣的其他專案提供了 Rust 到 CUDA 的 SIMT 風格 kernel 編譯，但 cuTile Rust 專注於基於 tile 的方法。這種權衡是放棄了顯式的 warp 原語與手動共享記憶體控制，以換取一種能在編譯時進行安全性檢查的語義模型。

專案狀態與生態系統

cuTile Rust 目前是一個處於積極開發早期階段的研究專案。使用者應預期 API 可能會發生變動，且功能尚未完整。該專案可在 crates.io 上以 cutile 名稱取得，並透過 Nix flake 提供支援，以簡化環境設定。

工作區架構

  • cutile: 用於撰寫與執行 kernel 的主要使用者端 crate。
  • cutile-compiler: 透過 cutile-ir 處理 Rust kernel 到執行檔的編譯。
  • cuda-async / cuda-core: 分別提供非同步 CUDA 執行與慣用的安全 CUDA API。

Sources