昨天介紹 MLIR 為什麼使用多層 IR。今天把這個概念放進 Triton JIT 路徑，同一個 Python 矩陣乘法，如何在第一次 launch 時依參數與 GPU 產生 specialization，再一路編譯成 GB10 能執行的 machine code。

接下來幾天會逐層打開 TTIR、TTGIR、LLVM IR、PTX、cubin 與 SASS，最後用 Nsight Systems 確認 kernel 怎麼被送進 GPU，這次實驗機器是搭載 NVIDIA GB10 Grace Blackwell 的 ASUS Ascent GX10。

實驗統一使用

• M=N=K=256
• FP16 輸入與輸出、FP32 accumulator
• tile 為 64 × 64 × 32
• num_warps=4、num_stages=2
• 一維 grid，共 16 個 Triton program instance

本篇大綱

• 拆開 @triton.jit 從呼叫、specialization、編譯、cache 到 launch 的完整流程
• 說明 TTIR、TTGIR、LLVM IR、PTX、cubin／SASS 各自解決什麼問題
• 確認 GPU、CUDA、PyTorch 與 Triton 版本
• 實際執行一支使用 tl.dot 的矩陣乘法
• 用乾淨 cache 保留後續要讀的 compiler artifact

@triton.jit 到底做了什麼

先看 Triton kernel 最常見的兩段語法

@triton.jit
def matmul_kernel(...):
    ...

matmul_kernel[grid](..., num_warps=4, num_stages=2)

@triton.jit 會把 Python function 包成 JITFunction，讓 Triton runtime 可以取得原始程式碼、參數資訊與後端編譯入口。執行到 decorator 時，GPU machine code 通常還沒產生。後面的 matmul_kernel[grid](...) 才進入 JIT runtime，綁定引數、建立 specialization、查 cache，需要時啟動 compiler，最後呼叫 launcher。

這裡的 JIT（Just-In-Time compilation）可以拆成四個環節

Python kernel + launch arguments + compile options
  ↓ 1. bind arguments
判斷 runtime type、tl.constexpr、alignment 等 specialization 資訊
  ↓ 2. build cache key
先查目前 process 的 compiled-kernel cache
  ↓ miss
再進入 compiler，查 TRITON_CACHE_DIR 的硬碟 cache
  ↓ miss
Python AST → TTIR → TTGIR → LLVM IR → PTX → cubin
  ↓ 3. load
cubin bytes → CUmodule → CUfunction
  ↓ 4. launch
grid + num_warps + shared memory + arguments + stream
  ↓
GPU 執行 kernel

前兩個 cache 屬於不同位置。Process 內的 cache 讓相同 specialization 在同一個 Python 程式中直接重用 CompiledKernel，硬碟 cache 讓新的 Python 程式也能重用先前產生的 artifact。即使兩者都命中，runtime 仍要準備 grid、引數與 stream，接著送出 kernel launch。Cache 省掉的是編譯工作，不會省掉每次執行。

Specialization：JIT 為何要看這次怎麼呼叫

Triton kernel 同一個 Python function，可以依呼叫條件編出多個版本。這次矩陣乘法的 specialization 會納入

• A、B、C 等 runtime argument 的型別與 pointer attribute
• m、n、k、stride、block_size_m/n/k 等 tl.constexpr 的值
• num_warps=4、num_stages=2 等 backend options
• 目前 target GPU 與 NVIDIA backend 資訊

tl.constexpr 會在編譯時期成為常數。例如 K=256、BLOCK_SIZE_K=32，compiler 已知 K loop 走八次，也能針對固定 tile shape 做 lowering，一般 tensor pointer 傳入的位址與 tensor 內容則留到 runtime，不會因為 A 的每個元素改變就重新編譯。

所以下面的修改會產生另一個 specialization，或至少形成不同的 compiler cache key

# tile 與 pipeline options 改了
BLOCK_SIZE_M = 128
num_warps = 8
num_stages = 3

Specialization 的代價是第一次使用新組合時會遇到 cold compile latency，好處是 compiler 能利用已知的 shape、dtype、alignment 與 launch options 產生較具體的程式。實務上需要控制 specialization 數量，如果把變動頻繁的值都變成編譯時期常數，cache 中會出現很多 kernel variant。

為什麼中間需要這麼多層 IR

一個 Triton kernel 同時包含幾種距離很遠的資訊，Python 的迴圈與 pointer 運算、tile 的 shape、warp/lane 的資料分工、NVIDIA memory space、virtual instruction，以及特定 GPU 的 machine instruction，一次從 Python 直接改寫成 SASS，會讓所有決策擠在同一個步驟，也很難判斷某個 shape、layout 或 instruction 是在哪裡選錯。

Triton 3.6.0 的 NVIDIA backend 依序註冊 ttir、ttgir、llir、ptx、cubin stages。轉換過程會逐步加入 target 資訊：TTIR 保留 tile 運算，TTGIR 加入 GPU layout，LLVM IR／PTX 再把它展開成 NVIDIA 指令。各階段的分工如下

確認實驗的軟硬體

這次實驗使用 ASUS Ascent GX10，軟硬體版本使用

architecture:       aarch64
GPU:                NVIDIA GB10
driver:             580.95.05
CUDA toolkit:       13.0
ptxas:              13.0, V13.0.88
Python:             3.12.3
PyTorch:            2.11.0+cu130
Triton:             3.6.0
compute capability: (12, 1)

執行前可以用 nvidia-smi、ptxas --version 與 Python 套件版本確認環境。後續結果都以這組版本為準，Triton 或 CUDA 版本不同，IR 格式與選用的指令可能改變。

這幾天用的矩陣乘法

這次使用的 kernel 核心如下

@triton.jit
def matmul_kernel(
    a_ptr, b_ptr, c_ptr,
    m: tl.constexpr, n: tl.constexpr, k: tl.constexpr,
    stride_am: tl.constexpr, stride_ak: tl.constexpr,
    stride_bk: tl.constexpr, stride_bn: tl.constexpr,
    stride_cm: tl.constexpr, stride_cn: tl.constexpr,
    block_size_m: tl.constexpr,
    block_size_n: tl.constexpr,
    block_size_k: tl.constexpr,
):
    pid = tl.program_id(axis=0)
    num_pid_n = tl.cdiv(n, block_size_n)
    pid_m = pid // num_pid_n
    pid_n = pid % num_pid_n

    offsets_m = pid_m * block_size_m + tl.arange(0, block_size_m)
    offsets_n = pid_n * block_size_n + tl.arange(0, block_size_n)
    offsets_k = tl.arange(0, block_size_k)

    a_ptrs = a_ptr + offsets_m[:, None] * stride_am \
        + offsets_k[None, :] * stride_ak
    b_ptrs = b_ptr + offsets_k[:, None] * stride_bk \
        + offsets_n[None, :] * stride_bn

    accumulator = tl.zeros(
        (block_size_m, block_size_n), dtype=tl.float32
    )
    for k_start in range(0, k, block_size_k):
        a = tl.load(
            a_ptrs,
            mask=(offsets_m[:, None] < m)
            & (k_start + offsets_k[None, :] < k),
            other=0.0,
        )
        b = tl.load(
            b_ptrs,
            mask=(k_start + offsets_k[:, None] < k)
            & (offsets_n[None, :] < n),
            other=0.0,
        )
        accumulator = tl.dot(a, b, accumulator)
        a_ptrs += block_size_k * stride_ak
        b_ptrs += block_size_k * stride_bk

    c = accumulator.to(tl.float16)
    c_ptrs = c_ptr + offsets_m[:, None] * stride_cm \
        + offsets_n[None, :] * stride_cn
    c_mask = (offsets_m[:, None] < m) & (offsets_n[None, :] < n)
    tl.store(c_ptrs, c, mask=c_mask)

一個 program instance 計算 C 的 64 × 64 tile。M 與 N 各切成四塊，所以 grid 是 4 × 4 = 16。K 維則以 32 為一步，迴圈走八次。

先手算一個 program instance

第一次看這個程式，最容易卡在 pid、tile 和矩陣座標的關係，我們用 pid=6 手算一次

num_pid_n = 256 / 64 = 4
pid_m = 6 // 4 = 1
pid_n = 6 % 4  = 2

所以第 6 個 program instance 負責

C[64:128, 128:192]

計算這塊 C 時，需要沿 K 維讀取

A[64:128, 0:256]
B[0:256, 128:192]

K 維一次只處理 32 個元素，因此其中一次迴圈會做

A tile：64 × 32
B tile：32 × 64
局部結果：64 × 64

這個例子是後面讀 IR 的座標，看到 tensor<64x32xf16>、tensor<32x64xf16> 和 tensor<64x64xf32> 時，可以直接把它們對回 A tile、B tile 和 accumulator。

逐段讀 kernel

第一段把一維 program id 轉成 C 的二維 tile 座標

pid = tl.program_id(axis=0)
num_pid_n = tl.cdiv(n, block_size_n)
pid_m = pid // num_pid_n
pid_n = pid % num_pid_n

第二段建立這個 tile 要使用的 row、column 與 K index

offsets_m = pid_m * block_size_m + tl.arange(0, block_size_m)
offsets_n = pid_n * block_size_n + tl.arange(0, block_size_n)
offsets_k = tl.arange(0, block_size_k)

offsets_m 與 offsets_n 各有 64 個 index；offsets_k 有 32 個。[:, None] 和 [None, :] 再把一維 index 擴成二維 pointer grid

offsets_m[:, None]：64 × 1
offsets_k[None, :]： 1 × 32
相加並 broadcast：  64 × 32  ← A tile

B 的 shape 則是 32 × 64。例如 tensor<64x32xf16> 代表這個 program instance 的 A tile 有 64 × 32 個 FP16 元素，這個數字和 CUDA thread 數量沒有直接對應。

K loop 每次載入一對 A/B tile

a = tl.load(a_ptrs, mask=..., other=0.0)
b = tl.load(b_ptrs, mask=..., other=0.0)
accumulator = tl.dot(a, b, accumulator)

Mask 為 false 時填入 0，因此邊界 tile 不會讀取矩陣外的位址，也不會污染 dot 的結果。每次 tl.dot 都把新的部分乘積累加到 FP32 accumulator，迴圈結束後才轉成 FP16 寫回 C。

Launch 與 correctness check 可以寫成

M = N = K = 256
a = torch.randn((M, K), device="cuda", dtype=torch.float16)
b = torch.randn((K, N), device="cuda", dtype=torch.float16)
c = torch.empty((M, N), device="cuda", dtype=torch.float16)

grid = (triton.cdiv(M, 64) * triton.cdiv(N, 64),)
matmul_kernel[grid](
    a, b, c,
    M, N, K,
    a.stride(0), a.stride(1),
    b.stride(0), b.stride(1),
    c.stride(0), c.stride(1),
    64, 64, 32,
    num_warps=4,
    num_stages=2,
)
torch.cuda.synchronize()
torch.testing.assert_close(
    c, torch.matmul(a, b), rtol=1e-2, atol=1e-2
)

用全新 cache 跑一次

為了觀察完整編譯流程，我將 TRITON_CACHE_DIR 指到一個剛建立的空目錄，再執行程式。只指定固定路徑不會自動清除舊 cache，這裡用 mktemp -d 確保每次得到新的目錄

TRITON_CACHE_DIR="$(mktemp -d)" \
python matmul.py

程式會做三件事

1. 第一次 launch 觸發 Triton JIT。
2. 用 torch.matmul 檢查結果。
3. 將各層 artifact 寫入這次的乾淨 cache。

這次保留的編譯結果

Triton cache 中可以找到 .source、TTIR、TTGIR、LLVM IR、PTX、cubin 與 metadata。在 Triton 3.6.0 裡，.source 是 AST 轉成 TTIR 後、正式 TTIR passes 執行前的初始 snapshot，.ttir 則是 TTIR stage 完成後的結果。兩者都屬於 Triton dialect，差別在 pass pipeline 的位置。Metadata 記錄了這次 specialization

{
  "name": "matmul_kernel",
  "target": {"backend": "cuda", "arch": 121, "warp_size": 32},
  "num_warps": 4,
  "num_ctas": 1,
  "num_stages": 2,
  "shared": 8192,
  "tmem_size": 0
}

JIT 與 cache 要分開理解

@triton.jit 不一定每次呼叫都重新編譯，要看 cache-key 有沒有相同，常見的 cache-key 變化來源包括 kernel source、argument type 與 alignment specialization、tl.constexpr、num_warps、num_stages、目標 GPU，以及 backend options。Triton compiler/backend 版本、ptxas、extern library、debug/instrumentation 設定與部分 cache-invalidating environment variables 也可能改變 cache key。

這次刻意指定空的 TRITON_CACHE_DIR，是為了觀察完整編譯。如果要量穩定執行時間，應先 warm up 並重用 cache，如果要研究編譯器，反而要把 cache state 固定並記錄下來。

今天的實驗結論

• 在 GB10（compute capability 12.1）上完成簡單的 tl.dot matmul。
• 同一個 specialization 的 TTIR、TTGIR、LLIR、PTX 與 cubin 都已保存。
• Compiler artifact、module loading 與 kernel execution 是不同階段。
• shared=8192、tmem_size=0 等 metadata 只是線索，實際用了什麼指令，還要繼續往 TTGIR、PTX 和 SASS 驗證。

明天先從 TTIR 開始，從原始程式碼逐項找到 program id、pointer、mask、K loop 與 tt.dot。

參考資料

• Triton triton.jit API
  https://triton-lang.org/main/python-api/generated/triton.jit.html
• Triton matrix multiplication tutorial
  https://triton-lang.org/main/getting-started/tutorials/03-matrix-multiplication.html
• Triton 3.6.0 JIT runtime：specialization、process cache 和 launch
  https://github.com/triton-lang/triton/blob/v3.6.0/python/triton/runtime/jit.py
• Triton 3.6.0 compiler driver：cache 和 stage pipeline
  https://github.com/triton-lang/triton/blob/v3.6.0/python/triton/compiler/compiler.py
• Triton NVIDIA backend compiler stages
  https://github.com/triton-lang/triton/blob/v3.6.0/third_party/nvidia/backend/compiler.py
• CUDA Driver API
  https://docs.nvidia.com/cuda/cuda-driver-api/