Triton，本文指 OpenAI Triton，先看官方介紹，

Triton is a language and compiler for parallel programming. It aims to provide a Python-based programming environment for productively writing custom DNN compute kernels capable of running at maximal throughput on modern GPU hardware.

Triton 是一個(gè)用于并行編程的 Python DSL，也是一個(gè)將計(jì)算映射到并行硬件的編譯器框架。

兩個(gè)視角，

• 算法研究：使用 Python DSL 快速實(shí)現(xiàn)高性能算子，無(wú)縫接入 PyTorch 生態(tài)
• 編譯器開發(fā)：為自研 AI 芯片提供編譯基礎(chǔ)設(shè)施，可基于 MLIR 構(gòu)建完整編譯棧

它是一座連接高級(jí)算法思想與底層硬件并行性的橋梁，兼顧了生產(chǎn)力、性能與可移植性。

本文聚焦于 AI 編譯器開發(fā)，想的是怎么把 Triton + MLIR 上手玩起來(lái)。

不過在此之前，先聊聊我寫這篇文章的初衷。

為什么寫 Triton？

我也是初次接觸 Triton。那為什么想到寫它了呢？

因?yàn)槟壳?AI 編譯器的主流方向是 Triton + MLIR，正好打算深入學(xué)習(xí)一下。

但真要寫 Triton，現(xiàn)在又有問題了，怎么寫呢？如今問一問 AI 不就好了。

所以，文風(fēng)與內(nèi)容，打算稍作變化。啥意思，會(huì)多點(diǎn)‘廢話’，多嘮叨幾句，不再簡(jiǎn)要成手冊(cè)，這就由 AI 來(lái)做吧。

我呢，多談?wù)勛约旱南敕ǎ阂皇亲尨蠹腋杏X更像是在和一個(gè)人交流；二是分享個(gè)人實(shí)踐的經(jīng)驗(yàn)，希望能對(duì)大家有所裨益。

GoCoding 定位，讓大家了解技術(shù)上新的東西，或能引導(dǎo)大家進(jìn)入新的領(lǐng)域。技術(shù)深入跟隨會(huì)議或社區(qū)進(jìn)展就好。

怎么來(lái)學(xué) Triton？

1. 官方渠道：不用多說(shuō)，當(dāng)然從官方渠道開始。不建議看我這樣的二手信息。

   那我寫文的價(jià)值呢？一，對(duì)我個(gè)人最有用；二，讓關(guān)注我的人能了解新的技術(shù)、發(fā)現(xiàn)感興趣的領(lǐng)域，一起學(xué)習(xí)。

2. AI 輔助：與 AI 聊聊我的目標(biāo)，看看它建議的學(xué)習(xí)路徑。現(xiàn)在可真方便 ??

那我們開始吧！

首先，了解一樣?xùn)|西，我不僅僅會(huì)了解它是什么、用做什么。就如文章開頭對(duì) Triton 的描述。我一般還會(huì)去了解背景：它的由來(lái)、發(fā)展歷程。挖掘得細(xì)一點(diǎn)，有時(shí)還會(huì)發(fā)現(xiàn)一些有意思的小故事。

其次，就是對(duì)比。為什么是它，它的優(yōu)勢(shì)在哪，它在往什么方向發(fā)展。也就是了解一下它“到哪里去”的問題。更遠(yuǎn)的，未來(lái)的可能性與不確定性，學(xué)習(xí)并深入時(shí)可能會(huì)想一想。

之后，才是學(xué)習(xí)。笨辦法，是我用得比較多的。就是，把一本書或官方教程，全部寫一遍、跑一遍。當(dāng)然，現(xiàn)在 AI 開道，免費(fèi)導(dǎo)師，可以試試聰明辦法了。又想起我自己說(shuō)過的那句：AI：不懂的人用，則是引導(dǎo)者 ??

現(xiàn)在，就按上述思考，開始吧！

Triton 的發(fā)展階段

Triton 的核心優(yōu)勢(shì)

Triton 核心優(yōu)勢(shì)：用 Python 寫 GPU 算子，基于 MLIR 輕松跨硬件，已是 PyTorch 高性能算子的事實(shí)標(biāo)準(zhǔn)前端。

Triton 與 TVM, XLA 的核心差異：

Triton + MLIR 跨硬件編譯架構(gòu)

Triton + MLIR 的學(xué)習(xí)路徑

學(xué)習(xí)路徑可以遵循“從用戶到開發(fā)者，從抽象到具體”的原則，分為四個(gè)層次：

層次一：作為算子開發(fā)者（用戶體驗(yàn)層）

目標(biāo)：學(xué)會(huì)使用Triton Python DSL快速實(shí)現(xiàn)高性能內(nèi)核。

1. 基礎(chǔ)入門：
   • 閱讀Triton官方教程，理解其編程模型（tl.program_id, tl.arange, tl.load/store）。
   • 親手實(shí)現(xiàn)幾個(gè)經(jīng)典算子：向量加法、矩陣乘法、Softmax。
   • 關(guān)鍵概念：塊（Block）、指針運(yùn)算、掩碼（Mask）、原子操作。
2. 進(jìn)階掌握：
   • 學(xué)習(xí)內(nèi)存優(yōu)化技巧：利用tl.cache, tl.make_block_ptr進(jìn)行向量化加載/存儲(chǔ)，理解不同內(nèi)存空間（DRAM, L2 Cache, Shared Memory）的差異。
   • 學(xué)習(xí)高級(jí)調(diào)度：如流水線（Pipeline）和異步復(fù)制（async_copy）。
   • 學(xué)習(xí)如何將Triton內(nèi)核封裝為PyTorch的torch.autograd.Function，并集成到模型中使用。
3. 實(shí)踐：嘗試復(fù)現(xiàn)或優(yōu)化一個(gè)經(jīng)典論文中的融合算子（如FlashAttention V1/V2的核心部分）。

層次二：理解編譯流程（系統(tǒng)認(rèn)知層）

目標(biāo)：理解從Triton Python代碼到目標(biāo)代碼的完整編譯過程。

1. 追蹤編譯鏈：使用torch.compile或Triton的JIT功能運(yùn)行一個(gè)簡(jiǎn)單內(nèi)核，并學(xué)習(xí)如何輸出其關(guān)鍵中間表示。
   • 查看Triton IR：了解高級(jí)Python操作如何被降低為Triton IR操作。
   • 查看MLIR：理解Triton IR是如何被轉(zhuǎn)換為MLIR中的Triton, LLVM, NVVM等Dialect的。這是理解其跨硬件能力的關(guān)鍵。
2. 學(xué)習(xí)MLIR基礎(chǔ)概念（無(wú)需深入，建立認(rèn)知即可）：
   • 什么是Dialect、Operation、Attribute、Type。
   • MLIR的核心優(yōu)勢(shì)：多層中間表示、可重用的轉(zhuǎn)換和優(yōu)化基礎(chǔ)設(shè)施。
3. 實(shí)踐：使用TRITON_DUMP_GRAPH=1等環(huán)境變量，導(dǎo)出一個(gè)內(nèi)核的編譯流程圖，對(duì)照官方架構(gòu)圖理解每一步。

層次三：探索編譯器開發(fā)（貢獻(xiàn)者/定制者層）

目標(biāo)：能夠?yàn)門riton編譯器添加新特性或?yàn)槠涠ㄖ菩掠布蠖恕?/p>

1. 深入Triton編譯器代碼：
   • 定位關(guān)鍵目錄：python/ (DSL), lib/ (核心編譯器邏輯，C++)。
   • 學(xué)習(xí)從Python AST到Triton IR的生成過程。
   • 學(xué)習(xí)基于MLIR的Pass是如何對(duì)Triton Dialect進(jìn)行優(yōu)化的。
2. 學(xué)習(xí)為MLIR添加Pass：
   • 編寫一個(gè)簡(jiǎn)單的MLIR Pass，對(duì)Triton Dialect的IR進(jìn)行模式匹配和重寫。
3. 后端開發(fā)入門（以連接自研芯片為例）：
   • 路徑一（推薦）：將Triton Dialect lowering到已有的、更通用的MLIR Dialect（如Vector, GPU），然后利用社區(qū)或自定義的后端從這些Dialect生成代碼。
   • 路徑二（更深入）：定義自己的硬件Dialect，并實(shí)現(xiàn)從Triton Dialect到該硬件Dialect的轉(zhuǎn)換規(guī)則。
   • 這需要對(duì)MLIR的方言轉(zhuǎn)換框架（Pattern Rewrite, Dialect Conversion）有扎實(shí)理解。

層次四：架構(gòu)與集成（系統(tǒng)架構(gòu)師層）

目標(biāo)：設(shè)計(jì)基于Triton+MLIR的完整編譯棧，用于產(chǎn)品化部署。

1. 性能調(diào)優(yōu)與自動(dòng)化：
   • 研究Triton的自動(dòng)調(diào)優(yōu)器（Auto-Tuner）原理，學(xué)習(xí)如何為新的硬件特性添加調(diào)優(yōu)空間。
   • 探索基于機(jī)器學(xué)習(xí)的內(nèi)核自動(dòng)生成與優(yōu)化。
2. 端到端集成：
   • 研究如何將Triton與模型級(jí)編譯器（如Torch-MLIR, IREE）結(jié)合，實(shí)現(xiàn)從PyTorch模型到混合內(nèi)核（部分由Triton生成）的無(wú)縫編譯。
   • 設(shè)計(jì)運(yùn)行時(shí)（Runtime）接口，管理自定義硬件上Triton內(nèi)核的加載、啟動(dòng)和資源分配。
3. 緊跟社區(qū)：關(guān)注Triton和MLIR的RFC、會(huì)議（如MLIR Open Design Meetings）、論文，把握技術(shù)演進(jìn)方向。

Triton + MLIR 動(dòng)手開始

學(xué)習(xí)路徑有了，那就動(dòng)手開始吧。

我根據(jù)個(gè)人情況，想法是對(duì)四個(gè)層次豎著切一刀，每一部分都實(shí)踐一個(gè)簡(jiǎn)單用例，先對(duì)整體建立概念，再繼續(xù)深入了解我關(guān)心的部分。

每個(gè)人都可以依據(jù)自己的興趣與實(shí)際情況去開始，不必非得按部就班地學(xué)。

結(jié)語(yǔ)

有興趣的小伙伴，一起來(lái)學(xué)吧。