Triton 程式設計入門：實用教程

📚 內容摘要

一個全面的科學教學課程，旨在為 Triton（一種基於 Python 的語言與編譯器，用於撰寫自定義 GPU 核心）提供完整的學習路徑。本課程涵蓋程式設計模型、語言語義、數值行為與效能優化，從基本的向量加法逐步推進至現代深度學習系統中使用的融合式與分塊運算。

從基礎原理掌握高性能 GPU 核心工程的藝術。

作者： EvoClass

致謝： Triton 文件與 Triton GitHub 倉庫。

🎯 學習目標

1. 定義 Triton 及其在深度學習軟體堆疊中的角色。
2. 区分 Triton 與 CUDA、PyTorch 緩慢執行代碼及低階 GPU 組合語言的差異。
3. 釐清哪些工作負載適合使用 Triton，並理解核心融合與瓶頸的相關性。
4. 完成 Triton 環境的乾淨安裝，並驗證軟體堆疊。
5. 實作一個基本的向量複製核心，以驗證環境邏輯與核心邏輯的差異。
6. 識別並分類 GPU 瓶頸，以支持使用 PyTorch 操作融合。
7. 定義一個程式實例，並使用 cdiv 計算一維啟動網格的維度。
8. 進行指標運算，將特定的程式識別碼（pid）對映至記憶體偏移。
9. 區分 PyTorch 張量（主機端元數據）與 Triton 張量（編譯器層級的區塊）。
10. 使用 tl.arange 計算程式識別碼（pid）與特定記憶體偏移之間的對映關係。

🔹 第 1 課：Triton 入門：哲學與設計

概述： 本課介紹 Triton，一種專用領域語言與編譯器，旨在彌補高階 Python 生產力與低階 GPU 性能之間的差距。學生將探索 Triton 的核心設計理念，建立一個概念性的思維模型，了解它如何以與標準 PyTorch 或 CUDA 不同的方式處理平行運算。

學習成果：

• 定義 Triton 及其在深度學習軟體堆疊中的角色。
• 区分 Triton 與 CUDA、PyTorch 緩慢執行代碼及低階 GPU 組合語言的差異。
• 釐清哪些工作負載適合使用 Triton，並理解核心融合與瓶頸的相關性。

🔹 第 2 課：環境設定與識別 GPU 瓶頸

概述： 本課涵蓋 Triton 開發的基礎要點，重點在建立穩定、乾淨的開發環境，並透過一個基本的「合理性」核心進行驗證。學生將學習如何區分不同類型的 GPU 性能瓶頸——運算、記憶體與啟動開銷，以判斷哪些 PyTorch 操作最適合手動操作融合。

學習成果：

• 完成 Triton 環境的乾淨安裝，並驗證軟體堆疊。
• 實作一個基本的向量複製核心，以驗證環境邏輯與核心邏輯的差異。
• 識別並分類 GPU 瓶頸，以支持使用 PyTorch 操作融合。

🔹 第 3 課：Triton 程式設計模型：網格與指標

概述： 本課介紹 Triton 程式設計模型，從 PyTorch 高階抽象過渡到基於區塊的 SPMD（單一程式，多資料）方法。學生將學習 Triton 如何透過一維啟動網格與程式實例組織執行，如何操控指標來存取記憶體，以及主機端 PyTorch 張量與編譯器層級 Triton 張量之間的根本差異。

學習成果：

• 定義一個程式實例，並使用 cdiv 計算一維啟動網格的維度。
• 進行指標運算，將特定程式識別碼（pid）對映至記憶體偏移。
• 區分 PyTorch 張量（主機端元數據）與 Triton 張量（編譯器層級的區塊）。

🔹 第 4 課：核心語言語義與記憶體遮罩

概述： 本課涵蓋使用 Triton 核心語言語義，將資料在全域記憶體與 GPU 廠商寄存器之間傳輸所需的基礎運算。學生將學習如何將平行程式實例對映至特定資料索引，透過記憶體遮罩管理邊界條件，並區分編譯時期常數與執行時期變數。

學習成果：

• 使用 tl.arange 計算程式識別碼（pid）與特定記憶體偏移之間的對映關係。
• 使用 tl.load 與 tl.store 搭配邊界遮罩，實現穩健的記憶體存取。
• 解釋 tl.constexpr 對編譯器最佳化的必要性，以及在定義形狀的函數中對執行時期值的限制。

🔹 第 5 課：實作你的第一個核心：向量加法

概述： 本課引導你完成創建 Triton 核心的完整生命週期，從理論轉化為實際的向量加法實作。你將學習撰寫 GPU 端核心，設計穩健的 Python 主機端包裝函數以啟動它，並實作科學驗證協定以確保正確性。

學習成果：

• 使用 Triton 的指標運算與遮罩系統，實作完整的向量加法核心。
• 設計主機端 Python 包裝函數，管理網格啟動、記憶體安全與輸入驗證。
• 使用 torch.allclose 執行嚴謹的驗證協定，以確保在多樣輸入大小與邊界情況下結果正確。

🔹 第 6 課：效能基礎：佔用率與基準測試

概述： 本課從基本的核心語法過渡到 GPU 效能的「第一原理」，聚焦於為何邏輯正確的程式仍可能效率低下。學生將探討記憶體流量、佔用率與硬體利用率之間的關係，最終以科學方法進行基準測試與 BLOCK_SIZE 最佳化。

學習成果：

• 使用 GPU 效能的第一原理，區分計算密集型與記憶體密集型核心。
• 解釋「權衡三角」，並說明佔用率如何隱藏記憶體延遲。
• 執行科學基準測試協定，包含預熱、同步與參數掃描。

🔹 第 7 課：二維張量與布局感知核心設計

概述： 本課從一維元素運算過渡到 Triton 中的二維張量處理。重點在於多維邏輯索引與線性物理記憶體之間的基本關係，透過位移（stride）實現。學生將學習如何構建二維指標網格，並設計尊重記憶體局部性的核心。

學習成果：

• 理解二維張量如何透過基底指標與位移在記憶體中表示。
• 使用廣播偏移模式，在 Triton 中構建二維地址網格。
• 實作布局感知核心（複製、轉置、偏移加法），正確處理非連續記憶體。

🔹 第 8 課：還原運算、Softmax 與數值穩定性

概述： 本課探討從簡單的元素運算核心過渡到更複雜的還原運算。學生將學習這類核心類型的架構差異，實現列式 Softmax 的標準模式，並理解數值穩定性在硬體上的關鍵作用。

學習成果：

• 比較還原核心與點對點核心的計算模式。
• 使用 Triton 五步還原模式，實作數值穩定的列式 Softmax 核心。
• 解釋在指數運算前減去最大值的數學與硬體層面必要性，以防止數值溢出。

🔹 第 9 課：矩陣乘法與大型語言模型的操作融合

概述： 本課探討從基本元素運算核心過渡到通用矩陣乘法（GEMM）及其在大型語言模型（LLMs）中的關鍵角色。學生將學習 Triton 中分塊的思維模型，操作融合帶來的效率提升，以及生產就緒核心所需的標準。

學習成果：

• 描述 Triton GEMM 思維模型，包括程式實例與分塊尺寸。
• 在 LLM 工作流程中識別操作融合的機會，並解釋其對效能的影響。
• 設計將偏移加法融合至 GEMM 輸出的邏輯實作。

🔹 第 10 課：優化生命週期：除錯與自動調校

概述： 本課探討從撰寫功能性 Triton 程式，過渡到開發生產級、高效能核心的過程。它建立了一套系統性的「語意至效能」除錯流程，並引入進行自動調校與基準測試所需的嚴謹思維。

學習成果：

• 系統化除錯： 採用分層策略，優先確保語意正確性與數值穩定性，再處理效能瓶頸。
• 實作自動調校工作流程： 定義元參數的有效搜尋空間，並執行基準測試協定，避免常見陷阱如硬體過度擬合。
• 規劃職業成長路徑： 釐清從基本 Triton 實作到高階生產級核心開發的「實務階梯」。