
核心概念
cuTile Rust(cutile‑rs)是以 Rust 語言撰寫、以瓦片(tile)為單位的 GPU 核心開發系統。它將 Rust 的所有權模型延伸至瓦片化的 GPU 核心,將可變輸出切割成互不重疊的片段,並在每一次 kernel 呼叫時保留主機端的所有權合約。開發者若需要更低階的控制,也可以在局部「opt‑out」直接使用 Tile IR(中間表示)指令。
轉譯工作流
TileGym CUDA 瓦片核心函式庫累積了大量以 cuTile Python(CUDA Tile Python)與 Triton‑TileIR(nvtriton)撰寫的生產級 kernel。為了讓這些資產同時可在 Rust 生態中使用,我們團隊打造了一套 AI 代理技能,自動將 Python 與 Triton‑TileIR 程式碼翻譯成 cuTile Rust。
轉譯流程採多代理管線,依序經過
-
分析階段:解析原始實作的記憶體操作與瓦片形狀。
-
裝置 kernel 產生:產出符合 CUDA Tile IR 的 GPU 程式。
-
主機與 FFI(Foreign Function Interface)程式碼:生成 Rust 側的呼叫介面。
-
效能基準測試:執行自動化基準,產出機器可驗證的判定。
每一階段皆以驗證腳本與 Tile IR 差異比對作為通關條件,確保翻譯過程不依賴人工信任。
效能與驗證
透過此技能,我們成功將 TileGym 中全部 24 個公開運算子(共約 40 個 GPU kernel)移植至 cuTile Rust,平均達到 99.5% 的 cuTile Python 效能。這些 kernel 包括從簡單的 element‑wise(逐元素)運算,到 Flash‑Attention 解碼、Multi‑head Latent Attention(MLA)以及 Mixture‑of‑Experts(MoE)模型等多樣化應用。值得注意的是,部分運算子仍需多個 kernel 變體以因應不同的硬體配置。
主要挑戰在於 cuTile Python 的 JIT(即時編譯)會在呼叫時隱式專門化(specialize)每個 kernel,而 Rust 必須在 kernel 簽名中明確列出所有專門化參數。多代理管線的機器可檢查機制正是為了解決此差異,確保最終產出的 Rust 程式碼在正確性與效能上皆與原始實作相符。
未來展望
由於 cuTile Python、Triton‑TileIR 與 cuTile Rust 皆以同一個 CUDA Tile IR 為底層表示,翻譯過程不涉及重新最佳化,而是將相同的瓦片程式以更安全的主機語言重新表述。未來,開發者可直接在 TileGym 倉庫取得此 AI 代理技能,將自家 kernel 快速搬移至 Rust 生態,享受 Rust 的記憶體安全與所有權保證,同時保留原有的 GPU 效能表現。這將進一步降低高效能運算領域對於語言切換的門檻,促進 Rust 在深度學習與大規模模型推論上的應用。