triton-lang / triton

triton-lang/triton

open_in_new前往仓库

Triton语言与编译器的开发仓库

项目概览

项目概述

Triton 是一个用于编写高效自定义深度学习原语的语言和编译器。其目标是提供一个开源环境，以比 CUDA 更高的生产力编写快速代码，同时比现有 DSL 具有更高的灵活性。该项目基于 MAPL2019 论文《Triton: An Intermediate Language and Compiler for Tiled Neural Network Computations》。

核心功能

  • 提供类似 Python 的 DSL，用于编写 GPU 内核。
  • 基于 MLIR 的编译器，支持 NVIDIA、AMD GPU，CPU 后端正在开发中。
  • 支持自动调优（autotuning）、Warp Specialization、软件流水线等高级优化。
  • 提供解释器模式，无需 GPU 即可调试内核。
  • 支持插件系统，允许扩展自定义 passes、dialects 和后端。

适用与不适用场景

适用场景：

  • 编写高性能的自定义深度学习算子（如 FlashAttention、GEMM）。
  • 需要比 CUDA 更高生产力的 GPU 编程。
  • 跨平台（NVIDIA、AMD）的 GPU 内核开发。

不适用场景：

  • 需要完全控制硬件底层细节（如寄存器分配、指令调度）的场景，可考虑 Gluon 或 CUDA。
  • 非 GPU 平台（CPU 后端尚不成熟）。
  • 需要稳定 API 的生产环境（Triton 仍处于快速迭代中）。

技术架构与依赖

  • 语言：MLIR（仓库语言标记），但主要使用 Python 和 C++。
  • 依赖：LLVM、MLIR、PyTorch（部分测试）、CMake、Ninja。
  • 架构：前端（Python DSL）→ Triton IR → MLIR 方言 → LLVM IR → PTX/AMDGPU。
  • 支持插件机制，允许第三方后端（如 Intel、CPU）通过插件集成。

安装与快速开始

从 pip 安装稳定版：

pip install triton

从源码安装：

git clone https://github.com/triton-lang/triton.git
cd triton
pip install -r python/requirements.txt
pip install -e .

使用自定义 LLVM 构建：

make dev-install-llvm

典型使用方法

以下是一个简单的 Triton 内核示例（来自官方文档）：

import triton
import triton.language as tl

@triton.jit
def add_kernel(x_ptr, y_ptr, output_ptr, n_elements, BLOCK_SIZE: tl.constexpr):
    pid = tl.program_id(axis=0)
    block_start = pid * BLOCK_SIZE
    offsets = block_start + tl.arange(0, BLOCK_SIZE)
    mask = offsets < n_elements
    x = tl.load(x_ptr + offsets, mask=mask)
    y = tl.load(y_ptr + offsets, mask=mask)
    output = x + y
    tl.store(output_ptr + offsets, output, mask=mask)

配置与部署要点

  • 环境变量：TRITON_INTERPRET=1 启用解释器，MLIR_ENABLE_DUMP=1 转储 IR，TRITON_ALWAYS_COMPILE=1 强制编译。
  • 构建选项：TRITON_BUILD_WITH_CLANG_LLD=true 使用 clang/lld，TRITON_BUILD_WITH_CCACHE=true 使用 ccache。
  • 缓存目录：TRITON_HOME 可更改 .triton 缓存位置。
  • 支持平台：Linux，NVIDIA GPU（Compute Capability 8.0+），AMD GPU（ROCm 6.2+）。

限制、风险与许可证

  • 许可证：MIT（仓库资料提供）。
  • 限制：CPU 后端仍在开发中；Windows 支持不完善；API 可能变动。
  • 风险：依赖 LLVM 的特定版本，构建复杂；社区维护，但核心维护者来自 OpenAI 等。

官方链接

信息来源和分析时间

  • 信息来源：GitHub 仓库 README、CONTRIBUTING.md、docs/meetups 下的会议记录、最新发布信息。
  • 分析时间：2026年6月18日（基于最新发布）。