English | 中文
本教程面向没有 GPU 经验的的Triton初学者,带你从基础的向量加到RoPE、matmul_ogs、topk、Gluon Attention 等大模型算子进阶学习之路。如果没有Python基础,可以通过Python编程入门教程(以在线评测平台为载体)来学习 Python 语法,或者根据本教程内容与 ChatGPT 对话直接入门 Triton。
作者邮箱:tt@bobhuang.xyz
OpenAI/Triton 是一个让你用 Python 写高性能 GPU 算子的编程语言(DSL)。目前有NVIDIA、AMD、华为昇腾、寒武纪、摩尔线程、沐曦等多个后端,一个kernel多种硬件均可以运行,具体见FlagOpen/FlagGems。
优势:保持 Python 的易用性,同时充分发挥 GPU 并行计算和内存优化能力。
应用:加速深度学习算子和自定义算子,提升大模型训练和推理性能。
本教程使用 Triton 3.4.0(released on 2025, Jul 31),只需安装 torch==2.8.0。若使用较低版本的 PyTorch,可自行升级 Triton版本。Triton具有还不错的版本兼容,大部分算子对Triton版本没有要求。
-
一、 Triton 简介
-
二、 向量加算子实战
- 1、 torch的向量加法
- 2、 单program 16个元素加法和验证
- 3、 通过mask控制元素访问
- 4、 多Block(program)运行
- 5、 使用参数化的BLOCK_SIZE
-
三、 练习1课后作业题解
- 1、 Matrix Copy 矩阵拷贝
- 2、 Color Inversion 颜色反转
- 3、 Reverse Array 数组反转
-
四、 矩阵转置实践
- 1、1D grid 进行native转置
- 2、2D grid 进行native转置
- 3、tl.trans 原语
- 4、维度扩展的1x1 block实践
- 5、通过mask控制元素访问
- 6、性能对比
-
附录
- 附录1、Triton kernel 原语覆盖情况
- 附录2、kernel 优化机制覆盖情况(未覆盖)
- 附录3、关联项目及学习资料