一句话定义:Triton 内核(Triton Kernel)是用 Triton 写出来的 GPU 计算函数。Triton 是一门面向 GPU 的编程语言与编译器,最初由 OpenAI 发起并开源,如今由社区持续维护。你用 Python 语法描述「对一整块数据做什么」,编译器负责把它翻译成能在显卡上跑起来的底层代码。
为什么需要它:从线程级到块级
传统写 CUDA(Compute Unified Device Architecture)算子,你得亲自决定:开多少个线程、每个线程负责哪个元素、要不要先把数据搬进共享内存、循环怎么展开、怎么用上矩阵乘加这类专用指令。像开手动挡,离合、换挡、转速全归你管,管得好很快,管不好就熄火。
Triton 换了个抽象层级:程序以「块」(block)为单位。给函数加上 @triton.jit,用 tl.arange 造一组下标,用 tl.load / tl.store 成块地读写显存,用掩码处理边界情况。至于这些块怎么映射到 GPU 的流多处理器、多少线程协作、寄存器怎么分配、要不要走共享内存和矩阵乘加指令,都交给编译器。
打个比方:写 CUDA 像自己砌墙,每块砖怎么摆、水泥怎么调都得你定;写 Triton 像把图纸交出去,说「这一面 128 行 64 列,往右挪一格再贴上去」,施工队(编译器)自己决定派几个工人、怎么排班。
和相邻概念的区别
| 概念 | 抽象层级 | 谁来优化 | 典型场景 |
|---|---|---|---|
| CUDA C++ / CUTLASS | 线程级 | 人 | 极致性能,要抠硬件细节 |
| Triton 内核 | 块级 | 编译器 + 少量调参 | 融合算子、新算法快速验证 |
| AI 词典:torch.compile">torch.compile / 算子库 | 图级或调用级 | 框架 | 不写内核,直接调用 |
两者不是替代关系。PyTorch 2 的 TorchInductor 后端在 GPU 上正是用 Triton 生成内核的,很多人其实在用它而不自知;反过来,当 Triton 编译不出理想性能时,也常退回手写 CUDA。
对从业者的实际意义
- 算法同学:想验证自定义注意力、量化反量化、归一化融合之类的写法,几十行 Python 就能跑通,不必先啃几个月 CUDA 工程。
- 工程同学:常见算子不必都手写,但要看懂生成的代码,会用
@triton.autotune扫块大小、流水级数等参数——性能调优仍是活。 - 其他职场人:你用的模型推理更快、显存更省,一部分功劳来自这类「让写算子变便宜」的工具。算子供给变多,框架才敢做更激进的融合。
具体支持哪些硬件后端、API 如何命名,以官方文档为准。
