PyTorch-Triton 3.7 版本引入了 Triton 插件扩展系统，这是一个用于在运行时动态加载自定义编译器 passes、方言 ops 和 DSL 扩展的框架，无需分支或重新编译。

在此处阅读技术深度解析和入门指南：https://bit.ly/4fATzAu

即开即用的 TLX 与自定义编译器 Passes – PyTorch

Source: https://pytorch.org/blog/triton-plugin-extensions-enabling-tlx-and-custom-compiler-passes-out-of-the-box/ Blog (https://pytorch.org/blog/category/blog/)

Triton 插件扩展：即开即用的 TLX 与自定义编译器 Passes

精选项目

PyTorch logo (https://pytorch.org/projects/pytorch/)

Triton 插件扩展

TLDR

PyTorch-Triton 3.7 版本引入了 Triton 插件扩展系统，这是一个用于在运行时动态加载自定义编译器 passes、方言（包括其 ops）和 DSL 扩展的框架，无需分支或重新编译。作为该系统的首个主要使用者，Meta 的 Triton 语言扩展（TLX）现已即开即用，为原生 Triton 带来了持久 GEMM 内核和细粒度硬件控制，在 NVIDIA H100 和 AMD MI350 上的性能均达到或超越厂商库。

问题：为何需要扩展？

编写高性能 GPU 内核通常需要超越默认 Triton 编译器流水线所提供的功能。自定义优化 passes、硬件特定的 intrinsic 以及专业的内存管理模式对于在生产负载上榨取最后一丝性能至关重要。在此之前，启用这些功能意味着需要维护一个 Triton 分支，而这样的分支会带来实际的成本。

分支会迅速落后于上游。每次上游更新都可能带来合并冲突、API 损坏以及需要仔细协调的细微行为变化。绑定到分支版本团队会发现他们被困在过时的版本上，无法利用上游的错误修复、新硬件支持和社区改进。维护负担会随着时间的推移而加重，分支反而成为瓶颈而非加速器。

我们需要一种方法来扩展 Triton 的编译器流水线，添加 passes、ops 甚至整个方言，而完全不需要修改核心 Triton。一个在运行时动态加载扩展的插件系统将允许研究人员和工程师全速迭代自定义功能，始终运行在最新的上游版本上，并且无需等待更改合并到主仓库即可发布结果。

Triton 插件扩展系统

PyTorch Triton 3.7 版本正是提供了这样的功能：一个通用的插件扩展系统，贯穿整个编译流水线并内置于上游 Triton 中。插件是共享库（.so 文件），通过 TRITON_PLUGIN_PATHS 环境变量在运行时发现并加载。安装插件包后，只需将该环境变量指向它，无需重新编译 Triton，扩展即可立即使用。

可覆盖的编译器流水线

该系统的核心是一组嵌入在 Triton 后端 compiler.py 阶段中的钩子。这些钩子提供了对 MLIR pass 流水线的细粒度控制，涵盖从高级 Triton IR (TTIR) 经 TritonGPU IR (TTGIR) 到 LLVM IR 和目标特定汇编 (PTX, AMDGCN) 的每个降级层次。利用这些钩子，插件可以：

• 在任意阶段的任意点插入一个或多个自定义 passes。
• 禁用阶段内的特定 passes。
• 用专门的自定义实现替换现有 passes（例如，自定义 warp specialization 策略）。
• 覆盖整个阶段或整个流水线。

此功能在 NVIDIA 和 AMD 后端均可用。

自定义 Ops、方言和降级

插件 API 设计用于补充 PyBind11，支持三个级别的可扩展性：

1. 自定义变换 passes：可在流水线任意点插入的单个 passes，无需关联方言。
2. 自定义 MLIR 方言和转换 passes：单独编译的方言加载到 Triton 中，结合插件 passes 将标准 Triton IR 模式重写为自定义方言 ops，以实现专业降级。
3. 自定义顶层 DSL ops：新的 Python 级别语法和语义，无需修改 Triton 本身即可实现全新的编程抽象。

每内核控制

插件可以在内核级别动态开启和关闭。在内核代码中设置的编译器钩子会激活一个自定义流水线，该流水线适用于钩子设置后调用的所有内核，直到取消设置。可以定义的自定义流水线数量没有限制，插件负责实现自己的哈希策略用于内核缓存管理——确保仅在需要时触发重新编译。这一切完全由 utlx 库为用户处理。

``

Enabling the TLX plugin is as simple as setting an environment variable

import os import sysconfig

dist_packages = sysconfig.get_paths()[“purelib”] libutlx_path = os.path.join(dist_packages, “utlx_plugin”, “libutlx.so”) os.environ[“TRITON_PLUGIN_PATHS”] = libutlx_path ``

TLX：Triton 语言扩展，现已内置

Triton 语言扩展 (TLX) (https://github.com/facebookexperimental/triton) 是 Meta 开发的一组硬件感知操作，用于显式内存管理和异步计算/加载流水线。TLX 为内核作者提供了对共享内存分配、数据移动和指令调度的直接控制——这些能力对于编写能饱和现代 GPU 硬件的持久内核至关重要。

核心 TLX 操作包括：

操作描述tlx\.local\_alloc\(shape, dtype, num\_buffers\)为软件流水线分配共享内存缓冲区。tlx\.local\_view\(buffers, index\)查看分配内的特定缓冲区。tlx\.async\_load\(src, dst, mask\)发起从全局内存到共享内存的异步加载。tlx\.async\_load\_commit\_group\(tokens\)提交一组异步加载。tlx\.async\_load\_wait\_group\(n\)等待异步加载组完成。tlx\.async\_dot\(a, b, acc\)异步矩阵乘累加。tlx\.async\_dot\_wait\(n, acc\)等待异步 dot 操作完成。tlx\.local\_store\(dst, src\)将数据存储到共享内存。tlx\.local\_load\(src\)将数据从共享内存加载到寄存器。

之前，使用 TLX 需要从 Meta 的实验性 Triton 分支构建。借助插件扩展系统，TLX 现在作为独立的 Python 包（utlx）分发，可与未经修改的上游 Triton 配合使用。从 PyTorch-Triton 3.7 开始，TLX 将在未来的所有 Triton 版本中默认启用。

跨硬件：NVIDIA H100 和 AMD MI350

TLX 的关键优势之一是，相同的编程模型可在不同硬件供应商间工作，同时在底层映射到供应商特定功能。

NVIDIA H100 (Hopper) — 持久 GEMM

在 Hopper GPU 上，TLX 映射到硬件原生的 TMA (Tensor Memory Accelerator) 异步加载和 WGMMA (Warp Group