Triton Compiler 教學

這裡提到的 Triton 是 OpenAI 在 2021 年
Triton: an intermediate language and compiler for tiled neural network computations 發表的 Framework

而不是 triton inference server,這同名又都是 ai 相關的熱門詞其實搜尋時需要花一點點力氣排除

其中包含了

  • Triton-C:由 Cuda-C 擴展出來為了支援 triton 相關特性
  • Triton-IR:由 LLVM-IR 擴展出來為了支援 triton 相關特性
  • Triton-JIT compiler

主要功能是

  1. 提供近似 cuBLAS 加速程度的運算,但簡化記憶體存取的語法
    2.利用 compiler 機制進行自動優化

Triton 與前後模組的架構關聯

圖片引用自 https://superjomn.github.io/posts/triton-mlir-publish/

概念上是依照下列的方式將 triton compiler 分成三個階段
file

實際上操作是實作多種不同的 Pass 然後轉換成 LLVM IR
file
LLVM 在轉換成 Cuda 定義的 IR PTX 然後轉成 cubin 機碼

而 triton compiler 教學中使用的主要操作就是使用者使用 Frontend 的介面來改寫原生的 python kernel 實作變為 triton 的 kernel 實作

以下是官方教學連結
https://triton-lang.org/main/getting-started/tutorials/index.html

這邊需要注意,網頁上的教學文字 在下載下來的 jupyter note 裡面也有 (只是圖片可能會遺漏)
建議直接下載 jupyter notebook 跟著說明跟 code 來看會比較簡單

我的學習紀錄

2024/07/13 社群:
Philippe Tillet:
Triton tutorials (especially the flash attention one) are not (no longer) meant to be canonical implementations of common kernels

https://github.com/microsoft/triton-shared
https://cloud.tencent.com/developer/article/2372957
file

https://www.lei.chat/zh/posts/mlir-linalg-dialect-and-patterns/

延伸問題

Reference

關於

AI Computing / 武術 / 登山 / IT / - 貪多而正努力咀嚼的人生小吃貨