A Coding Guide to NVIDIA’s Tile-Based GPU Programming: From cuTile and Triton Kernels to Flash Attention

NVIDIA 近期推出了一項名為 TileGym 的 GPU 程式教學，透過循序漸進的實作方式，帶領開發者認識基於 tile（區塊）的 GPU 程式設計模型。這份教學涵蓋了從環境設定、硬體探測到實際撰寫多種核心運算的全部過程，並且特別針對不同 GPU 環境提供了 cuTile 與 Triton 兩種後端支援，讓學習者無論是在高階 Ampere 架構 GPU 或標準 Colab T4 上都能順利執行。 教學的第一步是檢查執行環境。程式先偵測是否有可用的 CUDA GPU，並讀取 Compute Capability 與 CUDA 版本。由於 NVIDIA 的 cuTile 後端僅能在 Compute Capability 8.0 以上且 CUDA 13 以上的環境中運作，標準 Colab 提供的 T4 GPU（Compute Capability 7.5）無法直接使用，因此教學設計了自動降級機制：若 cuTile 無法導入，就自動切換至 Triton 後端。Triton 是 OpenAI 開發的開放原始碼 tile 編譯器，支援相同的高階程式設計模型，並且在 T4 等主流 GPU 上表現穩定。 整個教學的核心概念是「tile 程式模型」。傳統 CUDA 程式設計（SIMT）要求開發者為每一個執行緒撰寫程式碼，自行管理全域索引與邊界檢查，一次只操作一個元素。而 tile 模型則讓開發者專注於整個區塊（例如 1024 個元素或 128×128 的子矩陣），透過載入區塊、對整個區塊進行運算、再將結果寫回，剩下的執行緒與張量核心的對應則交給編譯器處理。教學中反覆強調：cuTile 與 Triton 的語法雖有差異，但背後是同一套思維。cuTile 使用 `ct.bid(0)`、`ct.load`、`ct.store` 等原語，Triton 則用 `tl.program_id`、`tl.load`、`tl.store`、`tl.dot`，兩者可以互相對照。 在實際撰寫核心的環節，教學依序實作了向量加法、融合 GELU 激活函數、行式 softmax、平舖矩陣乘法以及 Flash Attention。每個核心都先以 Triton 撰寫，並與 PyTorch 的標準實作進行正確性驗證與效能基準測試。向量加法是最基礎的範例，展示如何用 tile 載入兩個向量並相加後儲存。融合 GELU 則進一步結合了線性變換與高斯誤差線性單元的近似計算，在一個核心內完成整個運算，減少記憶體往返。 行式 softmax 是 transformer 模型中常用的操作，教學在 Triton 中以 `tl.max` 與 `tl.exp`、`tl.sum` 等區塊運算實現，並使用 `other=-float("inf")` 處理邊界。平舖矩陣乘法則展示了如何將大矩陣切分為多個 BM×BN 的 tile，並在內部迴圈中透過累加方式計算，這是所有 transformer 注意力機制的基礎。 最後的重頭戲是 Flash Attention 核心。Flash Attention 是一種高效的注意力計算方法，能在不儲存完整注意力矩陣的情況下，逐塊計算輸出，大幅降低記憶體使用。教學中的 Triton 版本實作了經典的 online softmax 技巧：為每個 tile 維護當前的最大值（m_i）與分母（l_i），並在每次載入新的 key/value tile 後更新，最後寫入正確的結果。這個核心的複雜度明顯高於前幾個，但也最貼近實務應用。 整個教學在 Colab 環境中以互動式筆記本的形式公開，使用者可以直接複製程式碼並運行。教學最後也強調，透過這套工作流程，開發者可以學會如何將傳統的逐元素心態轉換為 tile 心態，進而善用現代 GPU 的張量核心與快取架構，寫出更高效的深度學習運算核心。 對於有興趣深入 GPU 程式設計的開發者來說，這份教學提供了一個從基礎到應用的完整路徑，並且特別適合那些在有限硬體資源（如標準 Colab GPU）下仍想學習業界最新 tile 程式技巧的使用者。