triton-lang/triton-ext
GitHub: triton-lang/triton-ext
为 Triton 编译器提供树外扩展开发框架,支持在不修改核心代码库的前提下动态加载自定义 MLIR pass、dialect、后端和语言扩展。
Stars: 30 | Forks: 9
# Triton Extensions
一个为 Triton 编译器提供的树外扩展(out-of-tree extensions)集合,包括 pass、dialect、backend 以及语言扩展。
## 概述
本代码库提供了一个框架,用于开发和构建 Triton 编译器扩展,这些扩展无需修改核心 Triton 代码库即可扩展功能。扩展被构建为共享库,可以在运行时由 Triton 动态加载。
扩展构建于 [上游][triton-plugins] 文档中记载的 Triton 插件基础设施之上。
2026 年 1 月 Triton 社区Meetup 的幻灯片:
[Triton Community Meetup: Triton Extensions Jan 2026](https://docs.google.com/presentation/d/1dnm8uhvabdwqsQAsaPM7IRpEh2tktQ91E1d40r91n1M/edit?usp=sharing)
### 结构
每个子目录的 `CMakeLists.txt` 负责构建其各自的扩展。
- **[`backend/`](./backend/)**: 用于实现 backend 扩展
(目前仅为脚手架)。
- **[`dialect/`](./dialect/)**: 用于自定义 MLIR dialect 扩展
(目前仅为脚手架)。
- **[`language/`](./language/)**: 用于语言扩展的实现
(目前仅为脚手架)。
- **[`pass/`](./pass/)**: 包含 MLIR pass 扩展。每个 pass 扩展都
实现为可以动态加载的共享库。Pass
扩展包含一个 `triton-ext.toml` 清单,用于指定扩展
名称和状态。
- **[`extensions/`](./extensions/)**: 包含独立的插件扩展,将 dialect、pass 和语言绑定打包为自包含的共享库,可由 Triton 在运行时加载。
- **[`utlx` (µTLX)](./extensions/utlx/)**: Triton 语言扩展插件,
提供了 Meta 的大部分 TLX 功能,而无需修改 Triton
fork。包含本地内存操作(`local_alloc`、`local_view`、
`local_store`、`local_load`、`alloc_barriers`)、自定义 pass(例如
PingPong、PruneUnusedBarriers)、TLX dialect、转换模式以及
Python DSL。构建生成 `libutlx.so`。
- **[`support/`](./support/)**: 包含扩展基础设施代码,用于
向 Triton 自动注册扩展。
## 前置条件
- 支持 C++17 的 C++ 编译器
- CMake
- GitHub CLI ([`gh`]),用于下载预构建的依赖项(可选)
- Ninja
- Python 3,用于测试和构建脚本;使用以下命令安装依赖
`pip install -r requirements.txt`
- Triton,使用 `TRITON_EXT_ENABLED=ON` 构建
注意:在 3.7 及以上版本的 Triton 发行版中,默认已启用扩展。
## 构建
此扩展代码库设计为支持树外构建。它需要指定
LLVM (`LLVM_INSTALL_DIR`) 和 Triton (`TRITON_INSTALL_DIR`) 的路径。
这两个目录都应为安装目录,即先构建,然后
通过 `make install` 打包。
要构建扩展:
1. **构建 LLVM**:将 LLVM 构建为共享库并将其安装到一个已知
位置;参见 CI [action][build-llvm] 作为参考。或者,
从 GitHub 下载预构建的 LLVM 二进制文件:运行
`ci/download-artifact.py llvm` [^list-artifacts]。
2. **构建 Triton**:构建 Triton 并将其安装到一个已知
位置;参见 CI [action][build-triton] 作为参考。或者,
从 GitHub 下载预构建的 Triton 二进制文件:运行 `ci/download-artifact.py triton` [^list-artifacts]。
\[^list-artifacts\]: GitHub artifacts 仅对有限范围内的
commit、操作系统和硬件架构可用。要列出可用的 artifacts,请运行
`ci/fetch-artifacts.py`。
1. **构建扩展**:
export LLVM_INSTALL_DIR=/path/to/llvm/install
export TRITON_INSTALL_DIR=/path/to/triton/install
make build
请注意,如果未设置 `LLVM_INSTALL_DIR` 和 `TRITON_INSTALL_DIR`,
`Makefile` 将会自动在项目目录中[搜索][search]它们。参见
CI [workflow](./.github/workflows/ci.yml) 作为参考。扩展被构建为
位于 `build/lib` 下的共享库。
## 测试
运行测试套件以验证扩展是否正常工作:
```
make test
```
## 使用
Triton 在运行时使用 `TRITON_PLUGIN_PATHS`
环境变量加载扩展(参见 [Triton 插件][triton-plugins]):
```
export TRITON_PLUGIN_PATHS=/path/to/libmy_pass.so
python your_script.py
```
某些扩展可以从 Python 中访问:例如,µTLX 插件在被导入时会自动
注册为 `triton.language.extra.tlx`,因此不需要文件系统
软链接:
```
export TRITON_PLUGIN_PATHS=/path/to/triton-ext/build/lib/libutlx.so
python your_script.py
```
要加载多个插件,请使用 `:` 分隔路径:
```
export TRITON_PLUGIN_PATHS=build/lib/libutlx.so:build/lib/libother_plugin.so
```
标签:Bash脚本, MLIR, SOC Prime, Triton, 开发工具, 编译器, 编译器扩展, 逆向工具