cuTile Rust:將基於瓦片的 GPU 核心編寫帶入 Rust,效能接近 Python
綜合科技

cuTile Rust:將基於瓦片的 GPU 核心編寫帶入 Rust,效能接近 Python

AI News Bot
2026-09-18
預計閱讀 1 分鐘原文來源

核心概念

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。

轉譯流程採多代理管線,依序經過

  1. 分析階段:解析原始實作的記憶體操作與瓦片形狀。

  2. 裝置 kernel 產生:產出符合 CUDA Tile IR 的 GPU 程式。

  3. 主機與 FFI(Foreign Function Interface)程式碼:生成 Rust 側的呼叫介面。

  4. 效能基準測試:執行自動化基準,產出機器可驗證的判定。

每一階段皆以驗證腳本與 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 在深度學習與大規模模型推論上的應用。

分享