学习: Triton

下面把 Triton 编程范式 作为一个研究方向来梳理：它不只是"用 Python 写 GPU kernel"，而是把 GPU 编程抽象从 thread-level ​ 提升到 tile / block-program level​ 的一整套模型、编译器和工程方法论。

---

1. 一句话定位

CUDA：Scalar Program, Blocked Threads（你写一条线程，硬件复制成成千上万线程）

Triton：Blocked Program, Scalar Threads（你写一个 program 如何处理一整块数据，编译器再把块拆给线程/warp/Tensor Core）

Triton = Python DSL + MLIR 编译器 + tile 抽象 + autotune 工作流，目标是：

用接近 PyTorch 的生产力，写出接近手调 CUDA 的深度学习/GPU 算子。

---

2. 核心编程范式：Tile-Centric SPMD

2.1 基本心智模型

CUDA 里你写：

int i = blockIdx.x * blockDim.x + threadIdx.x;
y[i] = a[i] + b[i];

Triton 里你写：

@triton.jit
def add_kernel(x_ptr, y_ptr, out_ptr, n, BLOCK_SIZE: tl.constexpr):
    pid = tl.program_id(0)
    offs = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    mask = offs < n
    x = tl.load(x_ptr + offs, mask=mask)
    y = tl.load(y_ptr + offs, mask=mask)
    tl.store(out_ptr + offs, x + y, mask=mask)

关键点：

• 没有 threadIdx / warpIdx
• 没有 __shared__ 声明
• 没有 __syncthreads()
• 操作对象是 一整块 tensor / tile
• 边界用 mask 表达，而不是 if-break

---

2.2 三个核心概念

概念	含义
Program​	一个并行执行实例，大致对应 CUDA 的一个 thread block / CTA
Tile​	program 一次处理的一块数据（vector / sub-matrix）
constexpr meta-param​	BLOCK_M/N/K、num_warps、num_stages 等编译期超参

程序员决定：

• 数据怎么分块
• 哪个 program 处理哪个 tile
• 算法是否融合（softmax + residual、GEMM + epilogue）

编译器决定：

• tile → warp → thread 的映射
• 访存合并
• shared memory 分配与 pipeline
• Tensor Core 指令选择（mma.sync / wgmma / mfma）
• 软件流水线、register/spill 管理

---

3. 编译器范式：MLIR 多层 IR

Triton 不是"Python 转 CUDA 字符串"，而是正经编译器：

Python + @triton.jit
   ↓ trace / AST
Triton IR (TTIR)         # tile-level, hardware agnostic
   ↓
Triton GPU IR (TTGIR)    # 加入 layout / warp / shared / MMA
   ↓
LLVM IR
   ↓
PTX / cubin  (NVIDIA)
AMDGCN / hsaco (AMD)

TTGIR 是灵魂：

• #blocked、#dot_op、#shared、#nvidia_mma、#amd_mfma
• layout propagation
• distribute-to-warps
• async copy / software pipelining
• MMA intrinsic matching

这也是为什么 同一份 Triton 源码可以跨架构生成不同机器码。

---

4. Autotune 是范式的一部分，不是附属品

在 CUDA 里：

改 launch config → 改 template → 重编译 → 手测

在 Triton 里：

@triton.autotune(
    configs=[
        triton.Config({'BLOCK_M':128,'BLOCK_N':128,'BLOCK_K':32}, num_warps=4),
        triton.Config({'BLOCK_M':64,'BLOCK_N':256,'BLOCK_K':64}, num_warps=8),
    ],
    key=['M','N','K']
)

Autotune 把"架构知识"外置成配置空间：

• tile shape
• warp 数
• pipeline stage
• 是否用 dot / atomic / block ptr

这使得 Triton 工程范式为：

写对算法 → 给配置空间 → 首次运行 benchmark → 缓存最优 kernel

---

5. Triton 最适合什么问题

✅ 强项

• 融合算子：LayerNorm + Residual、Softmax、SwiGLU、Rotary
• 内存瓶颈算子：elementwise / reduction / softmax
• 自定义 GEMM epilogue
• FlashAttention / fused attention
• 非标准形状、非标准 dtype（fp8 / bf16 / nvfp4）
• torch.compile / Inductor 后端生成代码

❌ 弱项

• 强不规则访存：hash table、spMV、graph scatter
• all-to-all particle interaction（N-body）
• 需要精细 warp 级原语、custom TMA、warp specialization
• 极端性能天花板场景（专家级 CUTLASS / CUDA 仍更强）

有论文实测：

• 规则 compute-bound GEMM：Triton 可接近甚至超过 cuBLAS
• SpMV / N-body：Triton 可能明显落后于手调 CUDA

---

6. 与 CUDA / CUTLASS / TVM 的范式对比

维度	CUDA C++	CUTLASS	Triton	TVM
抽象层级	thread	tile（C++ template）	tile（Python）	tensor program
共享内存	手写	template/layout	编译器管	schedule 管
Tensor Core	手调 PTX	CuTe layout	tl.dot 自动映射	intrinsic
生产力	低	中低	高	中
性能天花板	最高	很高	高（90--98%）	高但调参重
适合人群	底层架构师	算子库作者	算法/训练/推理工程师	编译系统研究者

一句话：

CUTLASS 是"用 C++ 把 Triton 想做的事手动做一遍"；Triton 是"把 CUTLASS 的智力塞进编译器"。

---

7. 研究层面的开放问题

如果你要做"Triton 编程范式研究"，值得盯的方向：

1. Tile abstraction 的理论边界
   • 什么算子天然不适合 tile-SPMD？
   • irregular / dynamic-shape / control-flow-heavy 如何建模？
2. Layout-aware IR 设计
   • TritonGPU IR 的 layout propagation 是否足够？
   • Gluon / ML-Triton 的 workgroup-warp-intrinsic 多级 IR 是否是未来？
3. Autotune 自动化
   • 用代价模型替代暴力 bench
   • 结合 roofline / occupancy / register pressure 预测最优 config
4. 跨架构公平性问题
   • Hopper / Blackwell / AMD / Intel XPU 上同一源码性能方差很大
   • Triton 的"可移植性"到底是源码可移植，还是性能不可移植？
5. 与 LLM 代码生成结合
   • LLM 生成 Triton kernel
   • Triton 作为"GPU 算子中间表示"比 CUDA 更适合模型生成
6. 和 PyTorch 编译栈的关系
   • Inductor → Triton
   • 未来会不会出现"Triton 作为 AI 系统通用后端"？

---

8. 典型研究题目（可直接开做）

• 《Triton tile abstraction 对不规则访存算子的表达力边界》
• 《TritonGPU IR 中 layout propagation 对 GEMM pipeline 的影响》
• 《基于代价模型的 Triton autotune 替代方案》
• 《Triton vs CUTLASS vs CuTe：tile 编程范式系统性评测》
• 《LLM + Triton：自动生成 fused attention kernel 的可靠性研究》
• 《Hopper/Blackwell 上 Triton 的 TMA / wgmma 利用率分析》

---

9. 最小研究结论

Triton 的范式创新不是"Python 写 GPU"，而是：

把线程调度、访存合并、shared memory、Tensor Core 映射这些专家知识，从"人脑"移进"编译器 + autotune 配置空间"。

它牺牲一点极限性能，换来了：

• 更低认知负担
• 更快算子迭代
• 更接近算法本身的表达
• 成为 PyTorch / LLM 基础设施的"GPU 汇编层"

---