triton-lang/triton

什么是 Triton？

Triton 是一个开源编程语言和编译器，允许您为深度学习操作编写自定义 GPU 内核。它位于 PyTorch 等高级框架与低级 CUDA/ROCm 代码之间，提供 比原始 CUDA 更高的生产力，同时仍保留 手动优化性能关键原语（例如注意力机制、矩阵乘法、卷积）的灵活性。

核心理念

理念	对您意味着什么
GPU 内核的 DSL	使用类似 Python 的语法（`tl.*` 操作）编写内核。编译器将其转换为高效的 GPU 机器码。
基于 MLIR 的后端	编译器使用 LLVM 的 MLIR 基础设施，使管道更易于扩展，并支持 NVIDIA（NVPTX）和 AMD（AMDGPU）GPU。
自动调优	Triton 可自动探索分块大小和其他参数，以在当前硬件上找到最快版本的内核。
解释器模式	设置 `TRITON_INTERPRET=1` 可在无 GPU 的 CPU 上运行内核，便于调试。
插件式流水线	可通过 `add_stages_inspection_hook` API 注入自定义编译阶段。

快速上手

# 从 PyPI 安装最新稳定版（Python 3.10–3.14）
pip install triton

提供了适用于 Linux 且支持 NVIDIA Compute Capability 8.0+ 和 AMD ROCm 6.2+ 的二进制轮子。*

如需从源码构建（例如使用自定义 LLVM），README 中列出了步骤：

克隆仓库并安装 Python 构建依赖。
运行 make dev-install-llvm 以获取并编译 Triton 所需的 LLVM 版本，或参照手动 LLVM 构建说明。
以可编辑模式安装包（pip install -e .）。

典型工作流（Python 示例）

import triton
import triton.language as tl

@triton.jit
def matmul_kernel(A, B, C, M, N, K,
                 BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr):
    pid = tl.program_id(0)
    # 计算块索引 …
    # 加载分块，计算，存储结果 …

# 启动
matmul_kernel[(grid_x,)](A, B, C, M, N, K,
                         BLOCK_M=128, BLOCK_N=128, BLOCK_K=32)

@triton.jit 装饰器告诉 Triton 将函数编译为 GPU 内核。分块大小（BLOCK_*）可交由自动调优器处理，也可手动设置。

使用 Triton 可以构建的内容

自定义注意力内核 – 例如 flash-attention 实现。
优化的矩阵乘法 – 替代 cuBLAS 以适配特定形状。
实验性算子 – 在主流库中出现前尝试新原语。
研究原型 – 无需编写 CUDA 即可快速迭代内核想法。

社区与资源

文档： https://triton-lang.org – 安装指南、编程指南和调试教程。
会议： 每年一次的 Triton 开发者大会（2026 年详情见仓库），包含演讲、闪电演讲和海报展示。
谜题： 一组小型挑战（srush/Triton-Puzzles），可通过解释器运行，非常适合学习。
贡献： 欢迎提交拉取请求；详见 CONTRIBUTING.md 中的指南。
开发容器： triton-dev-containers 仓库提供了现成的 Docker/VSCode 开发容器，用于可复现的开发环境。

兼容性与硬件支持

操作系统： 当前版本仅支持 Linux。
GPU： 支持 Compute Capability 8.0 及以上版本的 NVIDIA GPU，以及 ROCm 6.2 及以上版本的 AMD GPU。CPU 支持正在开发中。

TL;DR

Triton 是一种 Python 友好的 DSL 和编译器，让您无需使用原始 CUDA/ROCm 即可为深度学习工作负载编写高性能 GPU 内核。通过 pip install triton 安装，使用 @triton.jit 装饰器编写内核，由基于 MLIR 的后端负责代码生成、自动调优和跨厂商支持。该项目积极维护，拥有完整的文档网站、年度开发者大会和友好的开源社区。