Triton 大会将于 2024 年 9 月 17 日在加利福尼亚州弗里蒙特再次举办！

如果您有兴趣参加，请填写此表单。

Documentation	Nightly Wheels

Triton

这是 Triton 的开发代码库，Triton 是一种用于编写高效自定义深度学习原语的语言和编译器。Triton 的目标是提供一个开源环境，以便以比 CUDA 更高的生产力编写快速代码，同时比其他现有 DSL 具有更高的灵活性。

该项目的基础在 MAPL2019 会议的以下论文中有所描述：Triton: An Intermediate Language and Compiler for Tiled Neural Network Computations。如果您使用 Triton，请考虑引用这项工作！

官方文档包含安装说明和教程。另请参见这些第三方Triton puzzles，所有这些都可以使用 Triton 解释器运行——无需 GPU。

快速安装

您可以通过 pip 安装 Triton 的最新稳定版本：

pip install triton

适用于 CPython 3.8-3.12 和 PyPy 3.8-3.9 的二进制 wheel 已提供。

以及最新的 nightly 版本：

pip install -U --index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/Triton-Nightly/pypi/simple/ triton-nightly

从源代码安装

git clone https://github.com/triton-lang/triton.git;
cd triton;

pip install ninja cmake wheel pybind11; # build-time dependencies
pip install -e python

或者使用 virtualenv：

git clone https://github.com/triton-lang/triton.git;
cd triton;

python -m venv .venv --prompt triton;
source .venv/bin/activate;

pip install ninja cmake wheel pybind11; # build-time dependencies
pip install -e python

使用自定义 LLVM 进行构建

Triton 使用 LLVM 为 GPU 和 CPU 生成代码。通常情况下，Triton 构建过程会下载预构建的 LLVM，但您也可以从源代码构建 LLVM 并加以使用。

LLVM 没有稳定的 API，因此 Triton 构建无法在任意 LLVM 版本上正常工作。

1. 确定 Triton 构建所依赖的 LLVM 版本。查看 cmake/llvm-hash.txt 以了解当前版本。例如，如果该文件内容为： 49af6502c6dcb4a7f7520178bd14df396f78240c

   这意味着您当前使用的 Triton 版本是基于 LLVM 49af6502 版本构建的。

2. 通过 git checkout 命令将 LLVM 切换到上述修订版本。您也可以选择对 LLVM 进行额外的修改。

3. 构建 LLVM。例如，您可以运行以下命令：

   $ cd $HOME/llvm-project  # 您的 LLVM 克隆目录。
   $ mkdir build
   $ cd build
   $ cmake -G Ninja -DCMAKE_BUILD_TYPE=Release -DLLVM_ENABLE_ASSERTIONS=ON ../llvm -DLLVM_ENABLE_PROJECTS="mlir;llvm" -DLLVM_TARGETS_TO_BUILD="host;NVPTX;AMDGPU"
   $ ninja
   

4. 稍作等待，构建过程需要一些时间。

5. 按照上述步骤构建 Triton，但需设置以下环境变量。

   # 根据您的 LLVM 构建目录进行适当修改。
   $ export LLVM_BUILD_DIR=$HOME/llvm-project/build
   
   $ cd <triton install>
   $ LLVM_INCLUDE_DIRS=$LLVM_BUILD_DIR/include \
     LLVM_LIBRARY_DIR=$LLVM_BUILD_DIR/lib \
     LLVM_SYSPATH=$LLVM_BUILD_DIR \
     pip install -e python
   

构建提示

• 将环境变量 TRITON_BUILD_WITH_CLANG_LLD 设置为 true，以使用 clang 和 lld。特别是 lld 可以显著加快构建速度。

• 设置 TRITON_BUILD_WITH_CCACHE=true 以使用 ccache 进行构建。

• 设置 TRITON_HOME=/some/path 可以更改 .triton 目录的位置，该目录用于存放 Triton 的缓存以及构建过程中下载的文件。默认情况下，此目录位于用户的主目录下。您可以随时更改此设置。

• 向 pip install 命令传递 --no-build-isolation 参数，可以加快无操作构建（nop builds）的速度。如果不使用此参数，每次调用 pip install 都会使用不同的 cmake 符号链接，这将迫使 ninja 重新构建大部分 .a 文件。

• VS Code 的 IntelliSense 在解析 Triton 的 C++ 代码构建方式时可能会遇到困难（这可能是因为在我们的构建流程中，用户不是直接调用 cmake，而是使用 setup.py）。您可以按照以下步骤配置 VS Code 以正确编译 Triton：

  • 进行本地构建。运行命令 pip install -e python
  • 获取构建过程中生成的 compile_commands.json 文件的完整路径： find python/build -name 'compile_commands.json' | xargs readlink -f。 您可能会得到类似 /Users/{username}/triton/python/build/cmake.macosx-11.1-arm64-cpython-3.12/compile_commands.json 的完整路径。
  • 在 VS Code 中，安装 C/C++ 扩展，然后打开命令面板（在 Mac 上为 Shift + Command + P，在 Windows/Linux 上为 Shift + Ctrl + P），并打开 C/C++: Edit Configurations (UI)。
  • 打开“高级设置”，并将 compile_commands.json 的完整路径粘贴到“Compile Commands”文本框中。

运行测试

目前尚无一键运行所有 Triton 测试的方法，但您可以按照以下步骤操作。

# One-time setup.  Note we have to reinstall local Triton because torch
# overwrites it with the public version.
$ pip install scipy numpy torch pytest lit pandas matplotlib && pip install -e python

# Run Python tests using your local GPU.
$ python3 -m pytest python/test/unit

# Move to builddir.  Fill in <...> with the full path, e.g.
# `cmake.linux-x86_64-cpython-3.11`.
$ cd python/build/cmake<...>

# Run C++ unit tests.
$ ctest -j32

# Run lit tests.
$ lit test

你可能会发现，创建一个指向 builddir 的符号链接并让本地 git 忽略它是很有帮助的。

$ ln -s python/build/cmake<...> build
$ echo build >> .git/info/exclude

然后你可以使用以下命令重新构建并运行 lit。

$ ninja -C build && ( cd build ; lit test )

开发提示

有关如何调试 Triton 前端的详细说明，请参考本教程。以下是针对 Triton 后端开发的额外提示。

有用的环境变量

• MLIR_ENABLE_DUMP=1 会在 Triton 运行的每个 MLIR passes 之前转储所有内核的 IR。使用 MLIR_ENABLE_DUMP=kernelName 仅为特定内核转储。

  • Triton 缓存可能会干扰转储。如果 MLIR_ENABLE_DUMP=1 不起作用，请尝试清理 Triton 缓存：rm -r ~/.triton/cache/*

• LLVM_IR_ENABLE_DUMP=1 会在对 LLVM IR 运行每个 passes 之前转储 IR。

• TRITON_INTERPRET=1 使用 Triton 解释器而非在 GPU 上运行。您可以在 kernel 代码中插入 Python 断点！

• TRITON_ENABLE_LLVM_DEBUG=1 向 LLVM 传递 -debug，将大量调试信息打印到标准输出。如果输出过于冗长，可仅使用 TRITON_LLVM_DEBUG_ONLY 来限制输出。

  另一种减少输出噪音的方法是运行 LLVM_IR_ENABLE_DUMP=1，提取感兴趣的 LLVM pass 之前的 IR，然后单独运行 LLVM 的 opt，并可在命令行上传递 -debug-only=foo。

• TRITON_LLVM_DEBUG_ONLY=<comma-separated> 等效于 LLVM 的 -debug-only 命令行选项。这将 LLVM 调试输出限制为特定的 pass 或组件名称（这些名称通过 LLVM 和 Triton 中的 #define DEBUG_TYPE 指定），以使调试输出不那么冗长。TRITON_LLVM_DEBUG_ONLY 允许指定一个或多个逗号分隔的值（例如 TRITON_LLVM_DEBUG_ONLY="tritongpu-remove-layout-conversions 或 TRITON_LLVM_DEBUG_ONLY="tritongpu-remove-layout-conversions,regalloc"）。

• USE_IR_LOC={ttir,ttgir} 重新解析 IR，使得位置信息将是具有特定扩展名的 IR 文件的行号，而非 Python 文件的行号。这可以提供从 IR 到 llir/ptx 的直接映射。与性能工具一起使用时，它可以提供 IR 指令的细分情况。

• TRITON_PRINT_AUTOTUNING=1 在自动调优完成后，打印出每个内核的最佳自动调优配置和花费的总时间。

• DISABLE_LLVM_OPT 如果解析为 Bool 时其值为 true，将为 make_llir 和 make_ptx 禁用 llvm 优化。否则，它将被解析为禁用 llvm 优化的标志列表。一个使用案例是 DISABLE_LLVM_OPT="disable-lsr"。 已知循环强度折减（Loop strength reduction）会导致某些寄存器压力较大的内核性能变化高达 10%。

• TRITON_ALWAYS_COMPILE=1 强制编译内核，无论是否命中缓存。

• MLIR_ENABLE_TIMING 转储每个 MLIR pass 的时间信息。

• LLVM_ENABLE_TIMING 转储每个 LLVM pass 的时间信息。

• TRITON_DEFAULT_FP_FUSION 覆盖允许 fp 融合（mul+add->fma）的默认行为。

• MLIR_ENABLE_REMARK 启用作为备注发出的性能警告。

更新日志

2.0 版本已发布！新功能包括：

• 大量错误修复
• 性能改进
• 后端使用 MLIR 重写
• 支持包含连续矩阵乘法的内核（例如 flash attention）

贡献指南

非常欢迎社区贡献，无论是修复错误还是添加新功能，都可以通过 github 进行。有关详细说明，请参阅我们的 贡献者指南。

兼容性

支持的平台：

• Linux

支持的硬件：

• NVIDIA GPU（计算能力 7.0+）
  • AMD GPU（ROCm 5.2+）
  • 开发中：CPU

Warp 专业化支持

Warp 专业化通过利用异步执行模型来提升内核性能，其中内核的不同部分由独立的硬件单元处理。在 H100 上，这些单元之间通过共享内存进行的数据通信效率很高。考虑到这一点，我们开发了一种自动 warp 专业化优化，将用户内核划分为异步任务（映射到 NVIDIA GPU 上的 warp 组），这些任务自然地并发执行，充分利用硬件的多任务 warp 调度器。以下部分详细介绍了为实现 warp 专业化而开发的编译器功能。

异步任务

Warp 专业化建立在将用户程序划分为异步任务（以下部分称为“async tasks”或“tasks”）的概念之上。每个异步任务将由支持的硬件上的独立 warp 组执行，以实现指令级并行。尽管编译器自动优化异步任务的划分仍然是一个挑战，但我们的自动任务划分方法已被证明对类似于 GEMM 和 Flash Attention 等典型示例的内核有效。

要启用 warp 专业化，用户只需指定某些自动调优标志，即 num_consumer_groups 和 num_buffers_warp_spec。例如，warp 专业化的 GEMM 实现可能如下所示。完整示例可参见 09-persistent-matmul.py。

@triton.autotune(
    configs=[
        triton.Config(
            {
                "BLOCK_SIZE_M": 128,
                "BLOCK_SIZE_N": 256,
                "BLOCK_SIZE_K": 64,
                "GROUP_SIZE_M": 8,
            },
            num_stages=2,
            num_warps=4,
            num_consumer_groups=2,
            num_buffers_warp_spec=3,
        ),
    ],
    key=["M", "N", "K"],
)
@triton.jit
def matmul_persistent_ws_kernel(
   a_ptr, b_ptr, c_ptr, M, N, K,
   stride_am, stride_ak, stride_bk, stride_bn, stride_cm, stride_cn,
   BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr,
):
   pid = tl.program_id(axis=0)
   num_pid_m = tl.cdiv(M, BLOCK_M)
   num_pid_n = tl.cdiv(N, BLOCK_N)
   pid_m = pid // num_pid_m
   pid_n = pid % num_pid_n
   offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
   offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
   offs_k = tl.arange(0, BLOCK_K)
   a_ptrs = a_ptr + (offs_m[:, None] * stride_am + offs_k[None, :] * stride_ak)
   b_ptrs = b_ptr + (offs_k[:, None] * stride_bk + offs_n[None, :] * stride_bn)
   acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
   for k in range(0, tl.cdiv(K, BLOCK_K)):
       a = tl.load(a_ptrs)
       b = tl.load(b_ptrs)
       acc += tl.dot(a, b)
       a_ptrs += BLOCK_K * stride_ak
       b_ptrs += BLOCK_K * stride_bk
   c = acc.to(tl.float16)
   c_ptrs = c_ptr + stride_cm * offs_m[:, None] + stride_cn * offs_n[None, :]
   tl.store(c_ptrs, c)

编译器会自动确定如何利用一个生产者 warp 组和两个消费者 warp 组来执行内核。它首先为某些锚点操作分配任务 ID，这些锚点操作会影响其余操作的任务分配。一旦锚点任务被标注，编译器会按以下方式将非锚点操作分配给任务：

• 某个锚点操作独有的控制依赖项会被包含在与该锚点操作相同的任务中。
• 某个锚点操作独有的数据依赖项会被包含在与该锚点操作相同的任务中，除非这些依赖项是另一个锚点操作。
• 任务之间共享的控制或数据依赖项会被包含在所有相关任务中。

对于上述 GEMM 示例，编译器会计算出一个任务方案，并使用 MLIR 属性在 IR 中对其进行标注。为了更清晰地说明这一点，我们使用源代码标注。任务传播后：

@triton.jit
def matmul_persistent_ws_kernel(
   a_ptr, b_ptr, c_ptr, M, N, K,
   stride_am, stride_ak, stride_bk, stride_bn, stride_cm, stride_cn,
   BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr,
):
   pid = tl.program_id(axis=0) # async_task 0, 1
   num_pid_m = tl.cdiv(M, BLOCK_M) # async_task 0, 1
   num_pid_n = tl.cdiv(N, BLOCK_N) # async_task 0, 1
   pid_m = pid // num_pid_m # async_task 0, 1
   pid_n = pid % num_pid_n # async_task 0, 1
   offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) # async_task 0, 1
   offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) # async_task 0, 1
   offs_k = tl.arange(0, BLOCK_K) # async_task 0
   a_ptrs = a_ptr + (offs_m[:, None] * stride_am + offs_k[None, :] * stride_ak) # async_task 0
   b_ptrs = b_ptr + (offs_k[:, None] * stride_bk + offs_n[None, :] * stride_bn) # async_task 0
   acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) # async_task 1
   for k in range(0, tl.cdiv(K, BLOCK_K)): # async_task 0, 1
       a = tl.load(a_ptrs)   # async_task 0
       b = tl.load(b_ptrs)   # async_task 0
       acc += tl.dot(a, b)   # async_task 1
       a_ptrs += BLOCK_K * stride_ak # async_task 0
       b_ptrs += BLOCK_K * stride_bk # async_task 0
   c = acc.to(tl.float16) # async_task 1
   c_ptrs = c_ptr + stride_cm * offs_m[:, None] + stride_cn * offs_n[None, :] # async_task 1
   tl.store(c_ptrs, c) # async_task 1

数据分区

为进一步提升性能，编译器会将相同的工作负载拆分到两个异步任务中。这样，当一个任务因繁重计算（例如点积运算）而阻塞时，另一个任务组可以并行执行其他操作。编译器会确定如何在两个任务之间分配工作以实现性能最大化。在 H100 GPU 上，编译器默认会尝试沿 M 维度拆分输入张量 A，使每个消费者独立计算输出张量的一半。这种方法称为协作分区。如果这种拆分没有优势（例如，导致 wgmma 指令小于原生指令），编译器将转而尝试沿 N 维度拆分。

上述 GEMM 内核（配置的分块大小为 [128, 256, 64]）的转换代码如下所示（为便于说明，使用源代码注释而非 IR）。

@triton.jit
def matmul_persistent_ws_kernel(
   a_ptr, b_ptr, c_ptr, M, N, K,
   stride_am, stride_ak, stride_bk, stride_bn, stride_cm, stride_cn,
   BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr,
):
   pid = tl.program_id(axis=0) # async_task 0, 1, 2
   num_pid_m = tl.cdiv(M, BLOCK_M) # async_task 0, 1, 2
   num_pid_n = tl.cdiv(N, BLOCK_N) # async_task 0, 1, 2
   pid_m = pid // num_pid_m # async_task 0, 1, 2
   pid_n = pid % num_pid_n # async_task 0, 1, 2
   offs_m_1 = pid_m * BLOCK_M + tl.arange(0, BLOCK_M // 2) # async_task 0, 1, 2
   offs_m_2 = pid_m * BLOCK_M + tl.arange(BLOCK_M // 2, BLOCK_M) # async_task 0, 1, 2
   offs_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_N) # async_task 0, 1, 2
   offs_k = tl.arange(0, BLOCK_K) # async_task 0
   a_ptrs_1 = a_ptr + (offs_m_1[:, None] * stride_am + offs_k[None, :] * stride_ak) # async_task 0
   a_ptrs_2 = a_ptr + (offs_m_2[:, None] * stride_am + offs_k[None, :] * stride_ak) # async_task 0
   b_ptrs = b_ptr + (offs_k[:, None] * stride_bk + offs_n[None, :] * stride_bn) # async_task 0
   acc_1 = tl.zeros((BLOCK_M // 2, BLOCK_N), dtype=tl.float32) # async_task 1
   acc_1 = tl.zeros((BLOCK_M // 2, BLOCK_N), dtype=tl.float32) # async_task 2
   for k in range(0, tl.cdiv(K, BLOCK_K)): # async_task 0, 1, 2
       a_1 = tl.load(a_ptrs_1)   # async_task 0
       a_2 = tl.load(a_ptrs_2)   # async_task 0
       b = tl.load(b_ptrs)   # async_task 0
       acc_1 += tl.dot(a_1, b)   # async_task 1
       acc_2 += tl.dot(a_2, b)   # async_task 2
       a_ptrs_1 += BLOCK_K * stride_ak # async_task 0
       a_ptrs_2 += BLOCK_K * stride_ak # async_task 0
       b_ptrs += BLOCK_K * stride_bk # async_task 0
   c_1 = acc_1.to(tl.float16) # async_task 1
   c_2 = acc_2.to(tl.float16) # async_task 2
   c_ptrs_1 = c_ptr_1 + stride_cm * offs_m_1[:, None] + stride_cn * offs_n[None, :] # async_task 1
   c_ptrs_2 = c_ptr_2 + stride_cm * offs_m_2[:, None] + stride_cn * offs_n[None, :] # async_task 2
   tl.store(c_ptrs_1, c_1) # async_task 1
   tl.store(c_ptrs_2, c_2) # async_task 2

代码划分

我们假设所有操作都已标记有任务 ID 列表。首先，我们找出 warp 组之间所需的所有通信。每个通信都始于一个带有单个任务 ID 的加载操作，并结束于属于不同任务 ID 的加载操作的直接使用者。对于包含通信通道的 ForOp，我们添加额外的参数：phase 和 bufferIndex。

我们引入了一个调优配置：num_buffers_warp_spec。对于每个通信通道，如果它位于 forOp 内，我们会使用 SMEM 中的一个缓冲区数组来保存结果，数组的大小由 num_buffers_warp_spec 决定。对于每个位于 ForOp 内的通信通道，我们还使用一个屏障数组。在此过程中，引入了四个新操作以正确同步生产者和消费者：ProducerAcquireOp、ProducerCommitOp、ConsumerWaitOp 和 ConsumerReleaseOp。这四个新操作均接受一个令牌和一个缓冲区索引作为参数。ProducerAcquire 和 ConsumerWait 还接受一个额外的阶段操作数。

对于具有多个任务 ID 的 ForOp，我们为每个任务 ID 克隆一个副本，每个副本包含具有特定任务 ID 的操作。最后，我们创建多个 IfOp，每个可能的任务 ID 对应一个。我们遍历函数体，为每个附加的任务 ID 克隆操作，并将克隆的操作放入正确的 IfOp 中。

为了调整寄存器使用，我们引入了两个新操作：RegAllocOp 和 RegDeallocOp，两者都接受一个整数操作数。对于每个 warp 组，我们决定插入 RegAllocOp 或 RegDeallocOp。当前的启发式规则很简单：如果任务 ID 为 0，则添加 RegDeallocOp，否则使用 RegAllocOp。寄存器调整的数量可以通过 reg_dec_producer 和 reg_inc_consumer 进行调优。

此过程还会将 loadOp 降级为 AsyncTMACopyGlobalToLocalOp 或 AsyncCopyGlobalToLocalOp，以便通过 SMEM 来表示通信。对于 TMA，生产者将变为 ProducerAcquire -> barrier_expect -> AsyncTMACopyGlobalToLocalOp，消费者将包含 wait_barrier -> 操作 -> ConsumerRelease。对于非 TMA 加载，生产者将变为 ProducerAcquire -> AsyncCopyGlobalToLocalOp -> ProducerCommitOp，消费者将包含 ConsumerWaitOp -> 操作 -> ConsumerRelease。