Convert Figma logo to code with AI

triton-lang logotriton

Development repository for the Triton language and compiler

14,992
1,887
14,992
798
View on GitHub

Top Related Projects

triton-lang's avatar

triton

14,095

Development repository for the Triton language and compiler

deepspeedai's avatar

DeepSpeed

37,573

DeepSpeed is a deep learning optimization library that makes distributed training and inference easy, efficient, and effective.

pytorch's avatar

pytorch

88,135

Tensors and Dynamic neural networks in Python with strong GPU acceleration

tensorflow's avatar

tensorflow

188,828

An Open Source Machine Learning Framework for Everyone

NVIDIA's avatar

cutlass

7,173

CUDA Templates for Linear Algebra Subroutines

facebookresearch's avatar

xformers

9,220

Hackable and optimized Transformers building blocks, supporting a composable construction.

Quick Overview

Triton is an open-source programming language and compiler designed for GPU programming. It aims to simplify the development of efficient GPU kernels by providing a high-level, Python-like syntax while generating high-performance CUDA code.

Pros

  • Simplifies GPU programming with a Python-like syntax
  • Generates optimized CUDA code automatically
  • Integrates well with PyTorch for machine learning applications
  • Provides automatic memory management and optimization

Cons

  • Limited documentation and learning resources compared to more established GPU programming tools
  • Still in active development, which may lead to frequent changes and potential instability
  • Primarily focused on NVIDIA GPUs, limiting its use on other hardware
  • Smaller community compared to more mainstream GPU programming languages

Code Examples

  1. Matrix multiplication:
@triton.jit
def matmul_kernel(
    a_ptr, b_ptr, c_ptr,
    M, N, K,
    stride_am, stride_ak,
    stride_bk, stride_bn,
    stride_cm, stride_cn,
    BLOCK_SIZE: tl.constexpr
):
    pid = tl.program_id(0)
    m_start = pid // (N // BLOCK_SIZE) * BLOCK_SIZE
    n_start = (pid % (N // BLOCK_SIZE)) * BLOCK_SIZE

    accumulator = tl.zeros((BLOCK_SIZE, BLOCK_SIZE), dtype=tl.float32)
    
    for k in range(0, K, BLOCK_SIZE):
        a = tl.load(a_ptr + m_start * stride_am + k * stride_ak + tl.arange(0, BLOCK_SIZE)[:, None] * stride_am + tl.arange(0, BLOCK_SIZE)[None, :] * stride_ak)
        b = tl.load(b_ptr + k * stride_bk + n_start * stride_bn + tl.arange(0, BLOCK_SIZE)[:, None] * stride_bk + tl.arange(0, BLOCK_SIZE)[None, :] * stride_bn)
        accumulator += tl.dot(a, b)

    c = accumulator
    tl.store(c_ptr + m_start * stride_cm + n_start * stride_cn + tl.arange(0, BLOCK_SIZE)[:, None] * stride_cm + tl.arange(0, BLOCK_SIZE)[None, :] * stride_cn, c)

This example demonstrates a matrix multiplication kernel using Triton.

  1. Element-wise addition:
@triton.jit
def add_kernel(x_ptr, y_ptr, output_ptr, n_elements, BLOCK_SIZE: tl.constexpr):
    pid = tl.program_id(0)
    block_start = pid * BLOCK_SIZE
    offsets = block_start + tl.arange(0, BLOCK_SIZE)
    mask = offsets < n_elements
    x = tl.load(x_ptr + offsets, mask=mask)
    y = tl.load(y_ptr + offsets, mask=mask)
    output = x + y
    tl.store(output_ptr + offsets, output, mask=mask)

This example shows an element-wise addition kernel using Triton.

  1. Softmax function:
@triton.jit
def softmax_kernel(output_ptr, input_ptr, input_row_stride, output_row_stride, n_cols, BLOCK_SIZE: tl.constexpr):
    row_idx = tl.program_id(0)
    col_offsets = tl.arange(0, BLOCK_SIZE)
    input_ptr += row_idx * input_row_stride
    output_ptr += row_idx * output_row_stride
    
    row_minus_max = tl.zeros((BLOCK_SIZE,), dtype=tl.float32) - float('inf')
    for col in range(0, n_cols, BLOCK_SIZE):
        cols = col + col_offsets

Competitor Comparisons

triton-lang logo

triton

14,095

Development repository for the Triton language and compiler

Pros of Triton

  • More active development with frequent updates and contributions
  • Broader community support and engagement
  • Extensive documentation and examples for easier adoption

Cons of Triton

  • Potentially less stable due to frequent changes
  • May require more frequent updates to keep up with new features and changes
  • Steeper learning curve for newcomers due to expanded feature set

Code Comparison

Triton:

@triton.jit
def add_kernel(x_ptr, y_ptr, output_ptr, n):
    pid = tl.program_id(0)
    offset = pid * 128
    x = tl.load(x_ptr + offset, mask=offset < n, other=0)
    y = tl.load(y_ptr + offset, mask=offset < n, other=0)
    output = x + y
    tl.store(output_ptr + offset, output, mask=offset < n)

Triton>:

@triton.jit
def add_kernel(x_ptr, y_ptr, output_ptr, n):
    for i in range(n):
        x = tl.load(x_ptr + i)
        y = tl.load(y_ptr + i)
        output = x + y
        tl.store(output_ptr + i, output)

Note: The code comparison is hypothetical, as Triton> doesn't exist as a separate repository. The example illustrates potential differences in syntax or approach that might exist between different versions or forks of the Triton project.

deepspeedai logo

DeepSpeed

37,573

DeepSpeed is a deep learning optimization library that makes distributed training and inference easy, efficient, and effective.

Pros of DeepSpeed

  • More comprehensive ecosystem for distributed training and optimization
  • Extensive documentation and tutorials for various use cases
  • Broader adoption in industry and research communities

Cons of DeepSpeed

  • Steeper learning curve due to more complex features
  • Potentially higher overhead for simpler use cases
  • Less focus on low-level GPU programming compared to Triton

Code Comparison

DeepSpeed:

import deepspeed
model_engine, optimizer, _, _ = deepspeed.initialize(
    args=args,
    model=model,
    model_parameters=params
)

Triton:

import triton
@triton.jit
def matrix_multiply(a_ptr, b_ptr, c_ptr, M, N, K):
    # Kernel implementation
    pass

DeepSpeed focuses on high-level distributed training optimizations, while Triton provides low-level GPU kernel programming capabilities. DeepSpeed offers a more comprehensive solution for large-scale model training, including features like ZeRO optimizer and pipeline parallelism. Triton, on the other hand, allows for fine-grained control over GPU computations, enabling custom kernel development for specific performance optimizations.

DeepSpeed is better suited for researchers and practitioners working with large language models and distributed training scenarios. Triton is more appropriate for developers needing to write custom CUDA-like kernels or optimize specific computational bottlenecks in their deep learning pipelines.

pytorch logo

pytorch

88,135

Tensors and Dynamic neural networks in Python with strong GPU acceleration

Pros of PyTorch

  • Mature ecosystem with extensive documentation and community support
  • Wide range of pre-built models and datasets available
  • Supports both deep learning and traditional machine learning tasks

Cons of PyTorch

  • Larger memory footprint and slower execution compared to specialized frameworks
  • Steeper learning curve for beginners due to its extensive feature set
  • Less optimized for specific hardware architectures

Code Comparison

PyTorch:

import torch

x = torch.tensor([1, 2, 3])
y = torch.tensor([4, 5, 6])
z = x + y
print(z)

Triton:

import triton
import triton.language as tl

@triton.jit
def add_vectors(x_ptr, y_ptr, output_ptr, n):
    pid = tl.program_id(0)
    x = tl.load(x_ptr + pid)
    y = tl.load(y_ptr + pid)
    output = x + y
    tl.store(output_ptr + pid, output)

# Usage would require additional setup and kernel launch

Triton offers more low-level control and optimization potential for specific hardware, while PyTorch provides a higher-level, more user-friendly interface for general-purpose deep learning tasks.

tensorflow logo

tensorflow

188,828

An Open Source Machine Learning Framework for Everyone

Pros of TensorFlow

  • Extensive ecosystem with wide industry adoption and support
  • Comprehensive documentation and large community for troubleshooting
  • Supports multiple programming languages (Python, JavaScript, C++)

Cons of TensorFlow

  • Steeper learning curve for beginners
  • Can be slower for certain operations compared to Triton's specialized GPU optimizations
  • More complex setup and configuration process

Code Comparison

Triton:

@triton.jit
def add_kernel(x_ptr, y_ptr, output_ptr, n_elements):
    pid = tl.program_id(0)
    block_start = pid * BLOCK_SIZE
    offsets = block_start + tl.arange(0, BLOCK_SIZE)
    mask = offsets < n_elements
    x = tl.load(x_ptr + offsets, mask=mask)
    y = tl.load(y_ptr + offsets, mask=mask)
    output = x + y
    tl.store(output_ptr + offsets, output, mask=mask)

TensorFlow:

import tensorflow as tf

def add_tensors(x, y):
    return tf.add(x, y)

x = tf.constant([1, 2, 3])
y = tf.constant([4, 5, 6])
result = add_tensors(x, y)
NVIDIA logo

cutlass

7,173

CUDA Templates for Linear Algebra Subroutines

Pros of CUTLASS

  • Highly optimized for NVIDIA GPUs, offering excellent performance for matrix computations
  • Extensive library of CUDA templates for various linear algebra operations
  • Supports a wide range of data types and precision levels

Cons of CUTLASS

  • Limited to NVIDIA GPUs, lacking cross-platform compatibility
  • Steeper learning curve due to its low-level CUDA-based implementation
  • Requires more manual optimization and tuning for specific use cases

Code Comparison

Triton example (matrix multiplication):

@triton.jit
def matmul_kernel(a_ptr, b_ptr, c_ptr, M, N, K):
    pid = tl.program_id(0)
    c = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
    for k in range(0, K, BLOCK_SIZE_K):
        a = tl.load(a_ptr + pid * M * K + k)
        b = tl.load(b_ptr + k * N + pid * N)
        c += tl.dot(a, b)
    tl.store(c_ptr + pid * M * N, c)

CUTLASS example (matrix multiplication):

using Gemm = cutlass::gemm::device::Gemm<float, cutlass::layout::RowMajor,
                                         float, cutlass::layout::RowMajor,
                                         float, cutlass::layout::RowMajor>;
Gemm gemm_op;
gemm_op(args...);
facebookresearch logo

xformers

9,220

Hackable and optimized Transformers building blocks, supporting a composable construction.

Pros of xformers

  • Broader scope: Focuses on efficient transformers and attention mechanisms
  • More extensive documentation and examples
  • Larger community and active development from Facebook AI Research

Cons of xformers

  • Higher complexity and learning curve
  • Less flexibility for low-level GPU programming
  • Potentially slower execution for certain specialized tasks

Code Comparison

xformers:

from xformers.components import Attention

attention = Attention(
    dim_model=512,
    num_heads=8,
    attention_dropout=0.1,
    bias=False,
)
output = attention(input_tensor)

Triton:

import triton
import triton.language as tl

@triton.jit
def attention_kernel(q, k, v, output, M, N, K):
    # Custom low-level implementation of attention mechanism
    # ...

Summary

xformers offers a higher-level API for working with transformers and attention mechanisms, making it more accessible for general use cases. Triton, on the other hand, provides a lower-level approach for GPU programming, allowing for more fine-grained control and potentially better performance in specific scenarios. The choice between the two depends on the project requirements, desired level of abstraction, and performance needs.

Convert Figma logo designs to code with AI

Visual Copilot

Introducing Visual Copilot: A new AI model to turn Figma designs to high quality code using your components.

Try Visual Copilot

README

Triton logo
DocumentationNightly Wheels
DocumentationWheels

Triton

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. See also these third-party Triton puzzles, which can all be run using the Triton interpreter -- no GPU required.

Quick Installation

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

pip install triton

Binary wheels are available for CPython 3.9-3.13.

Enabling Blackwell Support

The main branch now features support for NVIDIA Blackwell GPUs using 5th generation tensor cores. To enable this, you will need two additional steps:

  1. Build a pre-release PyTorch from source with CUDA 12.8
  2. Build triton from the latest source

First, to build pytorch you need to have CUDA 12.8 installed locally. If not, follow the instructions for your platform

# Clone and checkout pytorch 2.6 release candidate
git clone https://github.com/pytorch/pytorch
cd pytorch
git checkout v2.6.0-rc9
git submodule sync
git submodule update --init --recursive -j 8

# Install build dependencies (assumes you already have a system compiler)
pip install -r requirements.txt
pip install mkl-static mkl-include wheel

# Build PyTorch (will take a long time)
export CUDA_HOME=/usr/local/cuda-12.8
export CUDA_PATH=$CUDA_HOME
export TORCH_CUDA_ARCH_LIST=Blackwell
python setup.py develop

# Optional, package build into a wheel to install on other machines.
python setup.py bdist_wheel
ls dist  # Wheel should be output in this directory

Note that if you use the domain libraries (torchvision, torchtext, torchaudio, etc.) these will need to be built from source as well, otherwise their custom PyTorch extensions will not work.

Finally, follow the instructions below to install triton from source.

Install from source

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

pip install ninja cmake wheel pybind11 # build-time dependencies
pip 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 dependencies
pip install -e python

Building with a custom LLVM

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.

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

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

  2. git checkout LLVM at this revision. Optionally, make additional modifications 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;lld" -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

Tips for building

  • Set TRITON_BUILD_WITH_CLANG_LLD=true as an environment variable to use clang and lld. lld in particular results in faster builds.

  • Set TRITON_BUILD_WITH_CCACHE=true to build with ccache.

  • Set TRITON_HOME=/some/path to change the location of the .triton directory where Triton's cache is located and downloads are stored during the build. By default, this is the user's home directory. It can be changed anytime.

  • If you're running out of memory when building Triton, specify the MAX_JOBS environment variable (to the pip install -e python command) to limit the number of jobs.

  • Pass --no-build-isolation to pip install to make nop builds faster. Without this, every invocation of pip install uses a different symlink to cmake, 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, but instead use setup.py). Teach vscode how to compile Triton as follows.

    • Do a local build. Run command pip install -e python
    • Get the full path to the compile_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 the C/C++ extension, then open the command palette (Shift + Command + P on Mac, or Shift + Ctrl + P on Windows/Linux) and open C/C++: Edit Configurations (UI).
    • Open "Advanced Settings" and paste the full path to compile_commands.json into the "Compile Commands" textbox.

Running tests

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

# One-time setup.  Note this will reinstall local Triton because torch
# overwrites it with the public version.
$ make dev-install

# To run all tests (requires a GPU)
$ make test

# Or, to run tests without a gpu
$ make test-nogpu

Tips for hacking

For detailed instructions on how to debug Triton's frontend, please refer to this tutorial. 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 all kernels. Use MLIR_ENABLE_DUMP=kernelName to dump for a specific kernel only.

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

  • MLIR_DUMP_PATH specifies where MLIR_ENABLE_DUMP will dump to. If unset will dump to stderr.

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

  • TRITON_REPRODUCER_PATH=<reproducer_path> will generate an MLIR reproducer file at <reproducer_path> before each MLIR compiler stage. If any of the stages fail, <reproducer_path> will be a local MLIR reproducer captured right before the failing pass.

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

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

    An alternative way to reduce output noisiness is running with LLVM_IR_ENABLE_DUMP=1, extract the IR before the LLVM pass of interest, and then run LLVM's opt standalone, perhaps passing -debug-only=foo on the command line.

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

  • TRITON_ENABLE_ASAN=1 invokes the LLVM address sanitizer for memory leak and out of bounds access detection. Currently only supported on the AMD backend. This must be run using the ASAN libraries documented here.

    When enabling the address sanitizer it is recommended to disable various memory caching strategies both within the ROCm stack and PyTorch. This will give the address sanitizer the best chance at finding the memory fault where it originates. See this test for more details.

  • USE_IR_LOC={ttir,ttgir} reparses the IR such that the location information will 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 mapping from the IR to llir/ptx. When used with performance tools, it can provide a breakdown on IR instructions.

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

  • DISABLE_LLVM_OPT will disable llvm optimizations for make_llir and make_ptx if its value is true when parsing as Bool. Otherwise, it will be parsed as a list of flags to disable llvm optimizations. One usage case is DISABLE_LLVM_OPT="disable-lsr" Loop strength reduction is known to cause up to 10% performance changes for certain 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_DIAGNOSTICS=<comma-separated> controls diagnostic emission in MLIR. Options are: warnings, remarks, stacktraces, operations. Use comma-separated values to customize output. For example, MLIR_ENABLE_DIAGNOSTICS=remarks,operations enables remarks and IR operations, while MLIR_ENABLE_DIAGNOSTICS=warnings,stacktraces enables warnings with stacktraces. By default, only errors are shown. Setting warnings includes errors and warnings; remarks includes errors, warnings, and remarks.

  • MLIR_ENABLE_REMARK is deprecated. Please use MLIR_ENABLE_DIAGNOSTICS=remarks.

  • TRITON_KERNEL_DUMP enables the dumping of the IR from each compilation stage and the final ptx/amdgcn.

  • TRITON_DUMP_DIR specifies the directory to save the dumped IR and ptx/amdgcn when TRITON_KERNEL_DUMP is set to 1.

  • TRITON_KERNEL_OVERRIDE enables the override of the compiled kernel with a user-specified IR/ptx/amdgcn at the beginning of each compilation stage.

  • TRITON_OVERRIDE_DIR specifies the directory from which to load the IR/ptx/amdgcn files when TRITON_KERNEL_OVERRIDE is set to 1.

  • TRITON_F32_DEFAULT sets the default input precision of tl.dot when using 32-bit floats, which can be either ieee, tf32, or tf32x3.

  • TRITON_FRONT_END_DEBUGGING=1 disables exception wrapping when an error occurs in the compiler frontend, allowing the full stack trace to be seen.

Kernel Override Steps

export TRITON_ALWAYS_COMPILE=1
export TRITON_KERNEL_DUMP=1
export TRITON_DUMP_DIR=<dump_dir>
export TRITON_KERNEL_OVERRIDE=1
export TRITON_OVERRIDE_DIR=<override_dir>
# Step 1: Run the kernel once to dump kernel's IRs and ptx/amdgcn in $TRITON_DUMP_DIR
# Step 2: Copy $TRITON_DUMP_DIR/<kernel_hash> to $TRITON_OVERRIDE_DIR
# Step 3: Delete the stages that you do not want to override and modify the stage you do want to override
# Step 4: Run the kernel again to see the overridden result

Changelog

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)

Contributing

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.

Compatibility

Supported Platforms:

  • Linux

Supported Hardware:

  • NVIDIA GPUs (Compute Capability 8.0+)
  • AMD GPUs (ROCm 6.2+)
  • Under development: CPUs

Development Container (Dev Container)

Dev Containers for the Triton project are available from the triton-dev-containers repository

Key Benefits:

  • Consistency: All developers can work with the same development environment, ensuring uniform behavior across different systems.
  • Isolation: The container prevents potential conflicts with software installed on your local machine.
  • Portability: Easily share the development environment with team members, minimizing onboarding time and setup issues.

How to Use the Dev Container:

For detailed instructions on how to use the dev containers please see the dev container user guide