跳到主内容
快讯直播
AI智模界
AI 词典

Triton 内核:用 Python 替代手写 CUDA

一句话定义: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 如何命名,以官方文档为准。

AI 生成本文由 AI 基于公开信息自动生成,仅供参考。