We're hiring! If you are interested in working on Triton at OpenAI, we have roles open for Compiler Engineers and Kernel Engineers.
Documentation |
---|
This is the development repository of Triton, a language and compiler for writing highly efficient custom Deep-Learning primitives. The aim of Triton is to provide an open-source environment to write fast code at higher productivity than CUDA, but also with higher flexibility than other existing DSLs.
The foundations of this project are described in the following MAPL2019 publication: Triton: An Intermediate Language and Compiler for Tiled Neural Network Computations. Please consider citing this work if you use Triton!
The official documentation contains installation instructions and tutorials.
You can install the latest stable release of Triton from pip:
pip install triton
Binary wheels are available for CPython 3.7-3.11 and PyPy 3.8-3.9.
And the latest nightly release:
pip install -U --index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/Triton-Nightly/pypi/simple/ triton-nightly
git clone https://github.com/openai/triton.git;
cd triton;
pip install ninja cmake; # build-time dependencies
pip install -e python
Or with a virtualenv:
git clone https://github.com/openai/triton.git;
cd triton;
python -m venv .venv --prompt triton;
source .venv/bin/activate;
pip install ninja cmake; # build-time dependencies
pip install -e python
Triton uses LLVM to generate code for GPUs and CPUs. Normally, the Triton build downloads a prebuilt LLVM, but you can also build LLVM from source and use that.
LLVM does not have a stable API, so the Triton build will not work at an arbitrary LLVM version.
-
Find the version of LLVM that Triton builds against. Check
python/setup.py
for a line likeversion = "llvm-17.0.0-c5dede880d17"
This means that the version of Triton you have builds against LLVM c5dede880d17.
-
git checkout
LLVM at this revision. Optionally, make additional modifications to LLVM. -
Build LLVM. For example, you might run
$ cd $HOME/llvm-project # your clone of LLVM. $ mkdir build $ cd build $ cmake -G Ninja -DCMAKE_BUILD_TYPE=Release -DLLVM_ENABLE_ASSERTIONS=ON ../llvm -DLLVM_ENABLE_PROJECTS="mlir" $ ninja
-
Grab a snack, this will take a while.
-
Build Triton as above, but set the following environment variables.
# Modify as appropriate to point to your LLVM build. $ export LLVM_BUILD_DIR=$HOME/llvm-project/build $ cd <triton install>/python $ LLVM_INCLUDE_DIRS=$LLVM_BUILD_DIR/include \ LLVM_LIBRARY_DIR=$LLVM_BUILD_DIR \ LLVM_SYSPATH=$LLVM_BUILD_DIR \ pip install -e python
Version 2.0 is out! New features include:
- Many, many bug fixes
- Performance improvements
- Backend rewritten to use MLIR
- Support for kernels that contain back-to-back matmuls (e.g., flash attention)
Community contributions are more than welcome, whether it be to fix bugs or to add new features at github. For more detailed instructions, please visit our contributor's guide.
Supported Platforms:
- Linux
Supported Hardware:
- NVIDIA GPUs (Compute Capability 7.0+)
- Under development: AMD GPUs, CPUs