Repository files navigation

The Triton Conference is happening again on September 17th, 2024 in Fremont (CA)!

If you are interested in attending, please fill upthis form.

Documentation	Nightly Wheels

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!

Theofficial documentation contains installation instructions and tutorials. See also these third-partyTriton puzzles, which can all be run using the Triton interpreter -- no GPU required.

You can install the latest stable release of Triton from pip:

pip install triton

Binary wheels are available for CPython 3.8-3.12 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/triton-lang/triton.git;cd triton;pip install ninja cmake wheel pybind11; # build-time dependenciespip install -e python

Or with a 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 dependenciespip install -e python

Triton uses LLVM to generate code for GPUs and CPUs. Normally, the Triton builddownloads 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 anarbitrary LLVM version.

1. Find the version of LLVM that Triton builds against. Checkcmake/llvm-hash.txt to see the current version. For example, if it says:49af6502c6dcb4a7f7520178bd14df396f78240c

   This means that the version of Triton you have builds againstLLVM 49af6502.

2. git checkout LLVM at this revision. Optionally, make additionalmodifications to LLVM.

3. 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;llvm" -DLLVM_TARGETS_TO_BUILD="host;NVPTX;AMDGPU"$ ninja

4. Grab a snack, this will take a while.

5. 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>$ LLVM_INCLUDE_DIRS=$LLVM_BUILD_DIR/include \  LLVM_LIBRARY_DIR=$LLVM_BUILD_DIR/lib \  LLVM_SYSPATH=$LLVM_BUILD_DIR \  pip install -e python

• SetTRITON_BUILD_WITH_CLANG_LLD=true as an environment variable to use clangand lld. lld in particular results in faster builds.

• SetTRITON_BUILD_WITH_CCACHE=true to build with ccache.

• SetTRITON_HOME=/some/path to change the location of the.tritondirectory where Triton's cache is located and downloads are storedduring the build. By default, this is the user's home directory. Itcan be changed anytime.

• Pass--no-build-isolation topip install to make nop builds faster.Without this, every invocation ofpip install uses a different symlink tocmake, and this forces ninja to rebuild most of the.a files.

• vscode intellisense has some difficulty figuring out how to build Triton's C++(probably because, in our build, users don't invoke cmake directly, butinstead use setup.py). Teach vscode how to compile Triton as follows.

  • Do a local build. Run commandpip install -e python
  • Get the full path to thecompile_commands.json file produced by the build:find python/build -name 'compile_commands.json' | xargs readlink -f.You might get a full path similar to/Users/{username}/triton/python/build/cmake.macosx-11.1-arm64-cpython-3.12/compile_commands.json
  • In vscode, install theC/C++extension,then open the command palette (Shift + Command + P on Mac, orShift + Ctrl + P on Windows/Linux) and openC/C++: Edit Configurations (UI).
  • Open "Advanced Settings" and paste the full path tocompile_commands.json into the "Compile Commands" textbox.

There currently isn't a turnkey way to run all the Triton tests, but you canfollow the following recipe.

# 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.$ littest

You may find it helpful to make a symlink to the builddir and tell your localgit to ignore it.

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

Then you can e.g. rebuild and run lit with the following command.

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

For detailed instructions on how to debug Triton's frontend, please refer to thistutorial. The following includes additional tips for hacking on Triton's backend.

Helpful environment variables

• MLIR_ENABLE_DUMP=1 dumps the IR before every MLIR pass Triton runs, for allkernels. UseMLIR_ENABLE_DUMP=kernelName to dump for a specific kernel only.

  • Triton cache can interfere with the dump. In cases whereMLIR_ENABLE_DUMP=1 does not work, try cleaning your triton cache:rm -r ~/.triton/cache/*

• LLVM_IR_ENABLE_DUMP=1 dumps the IR before every pass run over the LLVM IR.

• TRITON_INTERPRET=1 uses the Triton interpreter instead of running on theGPU. You can insert Python breakpoints in your kernel code!

• TRITON_ENABLE_LLVM_DEBUG=1 passes-debug to LLVM, printing a lot ofdebugging information to stdout. If this is too noisy, run with justTRITON_LLVM_DEBUG_ONLY instead to limit the output.

  An alternative way to reduce output noisiness is running withLLVM_IR_ENABLE_DUMP=1, extract the IR before the LLVM pass of interest, andthen run LLVM'sopt standalone, perhaps passing-debug-only=foo on thecommand line.

• TRITON_LLVM_DEBUG_ONLY=<comma-separated> is the equivalent of LLVM's-debug-only command-line option. This limits the LLVM debug output tospecific pass or component names (which are specified using#define DEBUG_TYPE throughout LLVM and Triton) in order to allow the debug output tobe less noisy.TRITON_LLVM_DEBUG_ONLY allows for one or more commaseparated values to be specified (egTRITON_LLVM_DEBUG_ONLY="tritongpu-remove-layout-conversions orTRITON_LLVM_DEBUG_ONLY="tritongpu-remove-layout-conversions,regalloc").

• USE_IR_LOC={ttir,ttgir} reparses the IR such that the location informationwill be the line number of the IR file with that particular extension,instead of line number of the python file. This can provide a direct mappingfrom the IR to llir/ptx. When used with performance tools, it can provide abreakdown on IR instructions.

• TRITON_PRINT_AUTOTUNING=1 prints out the best autotuning config and total timespent for each kernel after autotuning is complete.

• DISABLE_LLVM_OPT will disable llvm optimizations for make_llir and make_ptxif its value is true when parsing as Bool. Otherwise, it will be parsed as a listof flags to disable llvm optimizations. One usage case isDISABLE_LLVM_OPT="disable-lsr"Loop strength reduction is known to cause up to 10% performance changes forcertain kernels with register pressure.

• TRITON_ALWAYS_COMPILE=1 forces to compile kernels regardless of cache hit.

• MLIR_ENABLE_TIMING dumps the timing information for each MLIR pass.

• LLVM_ENABLE_TIMING dumps the timing information for each LLVM pass.

• TRITON_DEFAULT_FP_FUSION overrides the default behavior of allowing fp fusion (mul+add->fma).

• MLIR_ENABLE_REMARK enables the performance warnings that are emitted as remarks.

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 atgithub. For more detailed instructions, please visit ourcontributor's guide.

Supported Platforms:

• Linux

Supported Hardware:

• NVIDIA GPUs (Compute Capability 7.0+)
• AMD GPUs (ROCm 5.2+)
• Under development: CPUs