作者：Jie Xin和Jonathan Bentz

NVIDIA CUDA Tile是基于 GPU 的編程模型，其設(shè)計目標(biāo)是為 NVIDIA Tensor Cores 提供可移植性，從而釋放 GPU 的極限性能。CUDA Tile 的一大優(yōu)勢是允許開發(fā)者基于其構(gòu)建自定義的 DSL。

本文介紹了 NVIDIA 正在將 CUDA Tile 集成為OpenAI Triton的后端的相關(guān)工作。OpenAI Triton 是一個開源的 Python DSL，專門用于編寫 GPU 上的 DL kernels。它支持分塊計算，可將數(shù)據(jù)與計算任務(wù)劃分為較小的 Block。Triton 內(nèi)置基于 MLIR 的編譯器，能夠生成PTX代碼，這使得即使不具備 CUDA 經(jīng)驗的研究人員也能編寫出高效的 GPU 代碼。

什么是CUDA Tile和CUDA Tile IR？

CUDA Tile 是對 CUDA 編程模型的擴(kuò)展，對 Tile 編程提供原生支持。CUDA Tile 在CUDA 13.1版本首次推出，標(biāo)志著 GPU 編程的一次重要演進(jìn)。與要求開發(fā)者基于 SIMT 模型、以單個線程為單位進(jìn)行思考不同，基于 Tile 的模型可使開發(fā)者在更高抽象層級上表達(dá)計算。

開發(fā)者只需操作數(shù)據(jù)塊（即 Tile），編譯器和運行時系統(tǒng)便會自動處理線程調(diào)度、硬件映射與資源分配。這不僅顯著降低了編程復(fù)雜度，也為編譯器的大幅優(yōu)化創(chuàng)造了條件。

CUDA Tile IR是基于 MLIR 的中間表示形式及編譯器基礎(chǔ)設(shè)施。CUDA Tile 的整個開發(fā)流程由 CUDA Tile IR 規(guī)范驅(qū)動，該規(guī)范明確定義了在 NVIDIA GPU 上進(jìn)行 Tile 計算所需的形式化語義、操作和類型系統(tǒng)。

什么是Triton-to-TileIR？

Triton-to-TileIR后端是 Triton 的一個橋接層，使其能夠以 CUDA Tile IR（而非 PTX）作為目標(biāo)代碼。它擴(kuò)展了 Triton 編譯器生態(tài)，使開發(fā)者能將使用 OpenAI Triton 編寫的 GPU kernel 編譯并運行在新推出的 CUDA Tile IR 后端之上，使開發(fā)者無需重寫代碼，便可無縫利用現(xiàn)代硬件能力。

隨著 GPU 編程從傳統(tǒng)的 SIMT 模型演進(jìn)到基于 Tile 的抽象，這一集成讓開發(fā)者既能保留 Triton 易用的 Python 語法，又能獲得對 Tensor Cores 的原生 TileIR 支持以及更強的架構(gòu)可移植性。

Triton-to-TileIR 降低了這些新能力的使用門檻。值得注意的是，Triton 本質(zhì)上是基于 Tile 的編程語言，其理念是讓開發(fā)者以數(shù)據(jù)塊（即 Tile）為單位計算，而非單個線程。這一設(shè)計思想與 CUDA Tile IR 十分契合。

這種一致性使得 Triton 能夠采用一條更直接的后端編譯路徑：Triton-to-TileIR 不再將 Triton 的 Tile 級抽象轉(zhuǎn)換為線程級的 SIMT 代碼，而是保持其 Tile 級語義，并直接編譯至原生支持 Tile 粒度計算的 CUDA Tile IR。

Triton 現(xiàn)有用戶無需學(xué)習(xí)新語言或重寫代碼，即可體驗 CUDA Tile IR 帶來的性能優(yōu)勢。僅需設(shè)置環(huán)境變量，便可將編譯流程從原有的 PTX 后端切換至 CUDA Tile IR 后端，從而獲得更優(yōu)性能與面向未來的架構(gòu)兼容性。

此外，Triton 用戶將能夠根據(jù)每個 kernel 的具體需求，在應(yīng)用中靈活選擇使用 PTX 后端或 CUDA Tile IR 后端。

Triton-to-TileIR開發(fā)路線圖

Triton-to-TileIR 目前是 triton-lang 團(tuán)隊下的一個孵化項目，目前正處于積極開發(fā)階段。該代碼庫旨在作為協(xié)作平臺，共同實現(xiàn)并優(yōu)化 CUDA Tile IR 后端，并為其未來并入 Triton 主分支奠定基礎(chǔ)。

其開發(fā)路線圖主要包括以下幾項關(guān)鍵技術(shù)工作流：

核心轉(zhuǎn)換基礎(chǔ)設(shè)施：實現(xiàn) MLIR 語言轉(zhuǎn)換機(jī)制，將 Triton 操作映射到對應(yīng)的 CUDA Tile IR 操作。

測試與驗證：構(gòu)建完整的測試套件，以驗證包括控制流、內(nèi)存訪問模式、數(shù)值精度等各類邊緣情況轉(zhuǎn)換過程中的語義正確性。

性能基準(zhǔn)測試：設(shè)立性能基線，在矩陣乘法、卷積、逐元素操作及規(guī)約等多種操作中，系統(tǒng)性地對比 kernel 分別經(jīng)由 TileIR 與 PTX 編譯后的性能差異。

開源項目集成：與開源社區(qū)緊密協(xié)作，推進(jìn) CUDA Tile IR 后端在 Helion 等相關(guān)開源項目中的支持。

如何使用Triton-to-TileIR

目前，Triton-to-TileIR 僅支持從源代碼編譯，暫無預(yù)編譯的二進(jìn)制包，因此您需要在本地自行從源代碼構(gòu)建該項目。

前提條件：

CUDA版本：CUDA 13.1 或更高

GPU架構(gòu)：NVIDIA Blackwell架構(gòu) GPU（如GeForce RTX 5080）；后續(xù) CUDA 版本預(yù)計將增加對早期 GPU 架構(gòu)的支持

從源碼構(gòu)建

在滿足以上條件后，可用以下代碼構(gòu)建項目：

# Clone the repository
git clone https://github.com/triton-lang/Triton-to-tile-IR.git
cd Triton-to-tile-IR
 
# Build and install
# Specific build instructions should be followed according to the project's README
pip install -e .

需要注意的是，具體的構(gòu)建步驟可能隨項目更新而變化。可查閱Triton-to-TileIR README文件與構(gòu)建指南，獲取針對您系統(tǒng)架構(gòu)的詳細(xì)配置說明、依賴管理指引及排障方法。

驗證Tile IR編譯

隨后，您可以通過運行向量加法教程，并確認(rèn)其是否調(diào)用 Tile IR 后端，來驗證安裝是否成功：

# Navigate to the tutorial directory
cd python/tutorials
 
# Run the vector addition example with Tile IR enabled
export ENABLE_TILE=1
python 01-vector-add.py

當(dāng) Tile IR 后端處于激活狀態(tài)時，Triton 會使用 .tileIR 文件擴(kuò)展來緩存編譯后的 kernel，而不是 SIMT 后端所使用的標(biāo)準(zhǔn) cubin 文件。請檢查以下緩存文件：

# Find the Triton cache directory (typically in ~/.triton/cache)

Triton-to-TileIR的局限性

盡管 Triton-to-TileIR 展現(xiàn)了一定的潛力，但目前仍處于早期開發(fā)階段，對部分操作的支持尚不完善且存在一些臨時性的性能問題。

暫不支持的操作

目前，Tile IR 后端尚未完全覆蓋 Triton 支持的所有操作。建議查閱“當(dāng)前暫未支持或功能未完全覆蓋的操作清單”。

隨著 CUDA 版本的不斷迭代，Triton CUDA Tile IR 后端也將在兼容性與功能上實現(xiàn)同步增強。

Tensor-of-pointer性能下降

在 CUDA 13.1 的 Tile IR 后端上，Triton 的“tensor-of-pointer”模式（即用 pointer 構(gòu)成的 tensors 來描述內(nèi)存訪問模式）目前性能暫未達(dá)到最優(yōu)水平。這是一個暫時性的性能局限。對于受影響的工作負(fù)載，可以采取以下措施：

將關(guān)鍵操作暫時切換回 SIMT 后端執(zhí)行

關(guān)注后續(xù)版本更新，相關(guān)問題將在新版本中優(yōu)化

優(yōu)化代碼，改用 TMA load/store API

關(guān)于最后提到的代碼優(yōu)化，許多 kernel 中加載的 tensors 本身具有連續(xù)的 Tile 布局以及明確的形狀和步幅。因此，無需在 kernel 內(nèi)部顯式構(gòu)建 tensor-of-pointers。相反，可以直接將這些布局信息傳遞給 TMA load/store API，從而提升 Tile IR 后端性能。

以下為典型 tensor-of-pointers 模式代碼：

# Before: tensor-of-pointer style
offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
offs_k = tl.arange(0, BLOCK_K)
 
a_ptrs = a_ptr + (offs_m[:, None] * stride_am
                  + offs_k[None, :] * stride_ak)
b_ptrs = b_ptr + (offs_k[:, None] * stride_bk
                  + offs_n[None, :] * stride_bn)
 
a = tl.load(a_ptrs)
b = tl.load(b_ptrs)

a_ptrs 中的每個元素都是在 kernel 中計算出的顯式 pointer，即使該 Tile 本身是連續(xù)的，并且其布局可以通過（shape,strides,block_shape）完全描述。

使用 TMA，相同的操作可以重寫為：

desc_a = tl.make_tensor_descriptor(
    a,                             # base pointer
    shape=(M, K),
    strides=(stride_am, stride_ak),
    block_shape=(BLOCK_M, BLOCK_K) # tile size
)
desc_b = tl.make_tensor_descriptor(
    b, shape=(K, N),
    strides=(stride_bk, stride_bn),
    block_shape=(BLOCK_K, BLOCK_N)
)
 
 
offs_m = pid_m * BLOCK_M
offs_n = pid_n * BLOCK_N
 
a_tile = desc_a.load([offs_m, 0])       # [BLOCK_M, BLOCK_K]
b_tile = desc_b.load([0, offs_n])       # [BLOCK_K, BLOCK_N]
desc_c.store([offs_m, offs_n], acc)     # TMA-backed store

進(jìn)一步了解Triton-to-TileIR

Triton-to-TileIR 項目代表了 GPU 編程演進(jìn)的關(guān)鍵一步，它彌合了開發(fā)生產(chǎn)力與硬件效率之間的差距。通過將 Triton 易用、面向 Tile 的編程模型與 CUDA Tile IR 虛擬指令集連接起來，該集成有望為機(jī)器學(xué)習(xí)從業(yè)者和 GPU 開發(fā)者帶來性能提升、可移植性并為未來做好準(zhǔn)備。

TileIR 后端為現(xiàn)有 Triton 開發(fā)者提供了一條平滑的升級路徑，僅需修改極少的代碼即可體驗下一代 GPU 架構(gòu)。同時，它也向更廣闊的 GPU 編程生態(tài)展現(xiàn)了語言設(shè)計者與硬件供應(yīng)商深度協(xié)同的戰(zhàn)略價值：在不犧牲高級抽象快速迭代能力的前提下，更簡單地獲取先進(jìn)硬件功能。

隨著該項目從孵化階段逐步走向生產(chǎn)就緒，其如何推動 Triton 的采用率、如何重塑基于 Tile 的 GPU 編程范式，將成為值得持續(xù)關(guān)注的焦點。最終的衡量標(biāo)準(zhǔn)也將非常直觀：不具備深厚 GPU 專業(yè)知識的研究人員能否 NVIDIA GPU 上編寫出接近最優(yōu)性能的 Triton 代碼。

如需了解更多詳情，可查閱triton-lang/Triton-to-tile-IR的 GitHub 庫以及《CUDA Tile IR 后端性能調(diào)優(yōu)提示》文檔。

關(guān)于作者

Jie Xin 是 NVIDIA 的計算架構(gòu)師。他畢業(yè)于華中理工大學(xué)，主修計算機(jī)科學(xué)。專注于編譯器、深度學(xué)習(xí)框架和 GPU 架構(gòu)。

Jonathan Bentz 領(lǐng)導(dǎo) NVIDIA 的 CUDA 技術(shù)營銷工程團(tuán)隊，其團(tuán)隊專注于創(chuàng)建和提供引人入勝的內(nèi)容，并與 CUDA 開發(fā)者建立聯(lián)系。Jonathan 擁有愛荷華州立大學(xué)化學(xué)博士學(xué)位和計算機(jī)科學(xué)碩士學(xué)位。