Triton (programming language)

Triton is an open-source domain-specific language and compiler for writing compute kernels for graphics processing units (GPUs). Triton is used in deep learning frameworks including vLLM.

Kernels are written as decorated Python functions that operate on blocks ("tiles") of data, and it differs from lower level languages like CUDA in that much of the scheduling within a block is handled automatically.

Triton began as the doctoral research of Philippe Tillet at Harvard University and was further developed at OpenAI, which released version 1.0 in 2021. It is now the default GPU code-generation target for PyTorch's torch.compile, and is the subject of a book by Harshwardhan Fartale published by Manning.

Triton supports Nvidia and AMD GPUs, and Intel maintains an out-of-tree backend for its own GPUs. CPU support is under development.

Example

The code below illustrates how a softmax kernel could be written in Triton.

import torch
import triton
import triton.language as tl

@triton.jit
def softmax_kernel(out_ptr, in_ptr, in_stride, out_stride, n_cols,
                   BLOCK: tl.constexpr):
    row = tl.program_id(0)
    cols = tl.arange(0, BLOCK)
    mask = cols < n_cols
    # Padding lanes load -inf, so exp() maps them to 0
    x = tl.load(in_ptr + row * in_stride + cols, mask=mask, other=-float("inf"))
    x = x - tl.max(x, axis=0)          # subtract row max for numerical stability
    num = tl.exp(x)
    y = num / tl.sum(num, axis=0)
    tl.store(out_ptr + row * out_stride + cols, y, mask=mask)

def softmax(x: torch.Tensor) -> torch.Tensor:
    n_rows, n_cols = x.shape
    out = torch.empty_like(x)
    BLOCK = triton.next_power_of_2(n_cols)
    softmax_kernel[(n_rows,)](out, x, x.stride(0), out.stride(0), n_cols,
                              BLOCK=BLOCK)
    return out

Writing the kernel this way produces a single GPU kernel that computes the maximum, exponentials, sum and division in one pass, reading the input from memory once. A naive PyTorch implementation built from separate tensor operations launches each step as its own kernel, reading and writing intermediate results to memory between them. In practice, PyTorch's built-in torch.softmax avoids this by calling a dedicated hand-written CUDA kernel.

In CUDA, a row-wise softmax requires explicit synchronisation between the threads that cooperatively compute each row's maximum and sum; in Triton the compiler handles the synchronisation and shared-memory management.

References