OpenAI Triton 1.0 發佈

OpenAI 已發佈 Triton 1.0，這是一種開源且類似 Python 的程式語言與編譯器，旨在讓研究人員無需具備豐富的 CUDA 經驗，即可編寫高效能的 GPU 程式碼。Triton 讓開發者能夠使用大幅減少的程式碼，實現與專家調優的 CUDA kernel 同等的硬體效能——例如在 FP16 矩陣乘法上達到與 cuBLAS 相當的效能。

解決 GPU 編程的複雜性

為了避免在原生框架算子中建立與移動暫時張量（temporary tensors）所帶來的效能損失，編寫專用的 GPU kernel 通常是必要的。然而，手動進行 GPU 優化非常困難，因為它要求開發者管理三個主要的架構組件：

DRAM: 記憶體傳輸必須合併（coalesced）成大型交易，以利用現代介面的匯流排寬度。
SRAM: 資料必須在重複使用前手動存入 SRAM，並進行管理以避免共享記憶體銀行衝突（bank conflicts）。
ALUs: 計算必須在串流多處理器（SMs）之間及內部進行仔細的劃分與排程，以利用指令/執行緒層級的並行性與 Tensor Cores。

Triton 自動化了記憶體合併、共享記憶體管理以及 SMs 內部的排程，同時將高階演算法考量——例如分塊（tiling）與 SMs 間的同步——留給開發者處理。

Triton 程式模型

Triton 使用一種程式模型，其中 kernel 被定義為裝飾過的 Python 函數，並在執行個體網格（grid of instances）上啟動。與 CUDA 或 Numba 使用的單指令多執行緒（SIMT）模型不同，Triton 透過對區塊（blocks，維度為 2 的冪次之小陣列）的操作來揭示執行個體內的並行性。

關鍵技術差異

基於區塊的執行: 透過對區塊而非單個執行緒進行操作，Triton 抽象化了 CUDA 執行緒區塊內的並行問題，包括共享記憶體同步與 Tensor Core 排程。
指標算術 (Pointer Arithmetic): Triton JIT 將輸入視為指標而非張量，為複雜資料結構（如 block-sparse tensors）提供了必要的低階記憶體存取控制。
融合 Kernel (Fused Kernels): Triton 簡化了融合 kernel 的建立。例如，Triton 中的融合 softmax 實作可以在整個正規化過程中將行（rows）保留在 SRAM 中，從而最大化資料重複使用。這種方法可能比等效的 PyTorch 實作高出 2 倍的效率。

矩陣乘法效能

Triton 對於矩陣乘法（神經網路的核心運算）非常有效。它僅需大約 25 行 Python 程式碼，即可在 V100 tensor cores 上達到峰值效能。這種易用性讓開發者能夠自定義矩陣乘法 kernel，以包含融合轉換（例如 slicing 或 Leaky ReLU），而無需具備卓越的 GPU 編程專業知識。

系統架構與編譯器後端

Triton 的效能源自於以 Triton-IR 為核心的模組化架構，這是一種基於 LLVM 的中間表示法，其中多維區塊是其一等公民。

編譯流水線

Python AST: @triton.jit 裝飾器會遍歷 Python 函數的抽象語法樹（Abstract Syntax Tree）。
Triton-IR: 使用靜態單一賦值（SSA）建構演算法將 AST 轉換為 Triton-IR。
LLVM-IR & PTX: 編譯器後端會簡化並優化 IR，自動進行並行化，並將其轉換為 LLVM-IR，最終轉換為可在 NVIDIA GPU 上執行的 PTX。

編譯器優化

自動記憶體存放 (Automatic Memory Stashing): 編譯器會分析計算密集型區塊級操作（例如 tl.dot）的操作數，並使用存活分析（liveness analysis）自動將資料存入共享記憶體。
自動並行化: Triton 透過同時執行不同的 kernel 執行個體來實現跨 SMs 的並行化，並透過在 SIMD 單元中劃分區塊級操作的迭代空間來實現 SMs 內部的並行化。