初识 Triton

❗️此坑还没填完，等到后面用到 triton 再补充

既生 CUDA， 何生 triton

CUDA 编程昂贵上手门槛促使 triton 的诞生^[1]。

• 语法福利 相比 CUDA C++ like 的设计风格，triton 使用 python。语法回避 C++ 模板编程和指针；环境集成比起 Pytorch-C++-CUDA 少了几层。
• 简化编程 将许多 GPU 并行编程优化比如 Memory Coalescing、Shared Memory Management 等自动化管理^[2]。

与 CUDA 通过 NVCC 编译成 PTX^[3]类似，Triton 通过 LLVM 最终编译成 PTX。最近也有其他厂商添加自己的 triton backend^[4]。

NVIDIA 也有类似 python 管理 GPU 的 python wrap^[5]，待日后调研学习。

triton 编程模型

GPU 编程最小粒度是线程，而 GPU 硬件上每 32 个线程分为一个 wrap，每次接受相同的指令，硬件调度最小粒度是 32 的 SIMD。编程模型无法直接和硬件模型相对应，NVIDIA 向开发者隐藏了转换的细节，区分 SIMD NVIDIA 将这种范式称作 SIMT。个人理解 SIMT 是软件和硬件共同表现而非纯硬件架构分类^[6]NVIDIA底层到底怎么实现的，谁知道呢。

CUDA 和 triton 编程都主要包括俩个函数，但编程粒度有所不同

• 每个子单元执行的计算（kernel function）
• 调度多个子单元并行计算（helper function / warper function）

参考 triton tutorial 的 vector addition mini example^[7]，编程 内容概括如下：

• kernel function
  • 获取标志当前 "program" 的信息（pid）pid = tl.program_id(axis=0)
  • 根据 "program" 信息计算输入输出数据地址范围 block_start = pid * BLOCK_SIZE; offsets = block_start + tl.arange(0, BLOCK_SIZE)
  • 根据指针载入输入数据 x = tl.load(x_ptr + offsets, mask=mask)
  • 调用模块计算 output = x + y
  • 将输出数据写回 tl.store(output_ptr + offsets, output, mask=mask)
• helper function
  • 根据执行软件计算维度（和硬件规格），划分硬件资源 grid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE']), )
  • 调用 kernel function 并行计算 add_kernel[grid](x, y, output, n_elements, BLOCK_SIZE=1024)

CUDA 编程以 grid 和 block 俩个 dim3 参数划分 thread，而每个 kernel function 内部通过设置的 Dim 维度数和 Id 计算当前线程唯一标识符，如果涉及多个 thread 间的交互则需要调用同步方法；而 triton 编程 kernel function 内部本身也带有一定的并行度，体现在 BLOCK_SIZE 通过 tl.arrange() 展开并行计算。

CUDA / triton 编程模型对比

	CUDA	triton
核函数配置参数	2 个 dim3 维度变量 gridDim, blockDim（最大划分 6 层）	任意数量 program 维度 + 任意数量每个 program 并行粒度 BLOCK_SIZE
传参方式	kernel_function<<<grid, block>>>(*args)	kernel_function[dim_func](*args, BLOCK_SIZE_1=,BLOCK_SIZE_2=, ...)
最小粒度	thread	program
核函数唯一标识符	通过 gridId, blockId, threadId 和配置参数计算	通过 programId 和 program 维度配置参数计算
获取 id	内建变量 gridId, blockId, threadId	tl.program_id(axis=)
核函数内部并行度	无并行度	并行度由 BLOCK_SIZE 定义
跨 thread 协同方法	调用跨 thread 同步方法	自动完成

使用函数包装配置参数 dim_function(dict)->tuple， 接受 triton config 入参 autotune 输出配置参数。

---

1. https://github.com/openai/triton ↩︎

2. https://openai.com/index/triton/ ↩︎

3. Parallel Thread Execution, 相当于 GPU 的汇编。CUDA 代码先经过 NVCC 编译器翻译成 PTX，再经过驱动翻译成可执行机器码 ↩︎

4. 寒武纪和微软的 triton-linalg 转换流程 https://github.com/Cambricon/triton-linalg ↩︎

5. https://github.com/NVIDIA/warp ↩︎

6. 该结论主要参考《计算机体系架构：量化研究方法》中 GPU 章节以及网上论坛得出 ↩︎

7. https://triton-lang.org/main/getting-started/tutorials/01-vector-add.html
   https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html ↩︎