Triton: a tile-level language and compiler for GPU kernels
A Python-embedded language in which a kernel is written for a block of data, with the compiler choosing memory layout, coalescing, shared-memory use and synchronisation; the default kernel backend of PyTorch's compiler.
PyTorch 2 (TorchInductor) generates Triton code for GPU operators, and many open attention, quantisation and fused-operator kernels are written in Triton; its compiler decides how close high-level code gets to hand-tuned CUDA.
A way to write fast GPU programs in Python by thinking in blocks of data instead of single threads.
- 1CUDA needs per-thread instructions
- 2Triton describes work per block of data
- 3Compiler picks layout, memory and sync
- 4Python code becomes tensor-core code
- 5Gap to expert kernels shrinks
What it is for. Writing a fast GPU kernel in CUDA, NVIDIA's C++-based language, means managing individual threads: which of thousands of threads loads which number, when they must wait for each other, and how data is staged in on-chip shared memory. That work is slow and error-prone. Triton (Tillet, Kung and Cox, 2019) lets a programmer write kernels in Python at the level of blocks of data, and its compiler (a program that translates code into machine instructions) works out the thread-level details.
Analogy. CUDA is like directing every worker on a building site individually; Triton is like giving a foreman instructions per room ("tile this floor") and letting the foreman assign the workers. The analogy breaks where the foreman's habits are not the best plan: the compiler's choices can leave speed on the table.
What a kernel looks like. A function decorated with @triton.jit is launched as many independent programs, each handling one block. Vector addition is the hello-world:
@triton.jit
def add(x_ptr, y_ptr, out_ptr, n, BLOCK: tl.constexpr):
pid = tl.program_id(0)
offs = pid * BLOCK + tl.arange(0, BLOCK)
mask = offs < n
x = tl.load(x_ptr + offs, mask=mask)
y = tl.load(y_ptr + offs, mask=mask)
tl.store(out_ptr + offs, x + y, mask=mask)
Program number handles elements to , where is the block size; the mask stops the last program from reading past the end of the arrays. For matrix multiplication, multiplies two tiles and the compiler maps it onto tensor cores, the GPU's matrix-multiply units. Block sizes are compile-time constants, so the compiler produces specialised code for each choice.
Go deeper
The full technical text.
The idea
CUDA exposes the machine at the level of a single thread: the programmer arranges threads, shared memory, synchronisation and memory coalescing by hand. Triton (Tillet, Kung and Cox, 2019) raises the unit of programming to a block. A kernel instance, called a program, operates on whole tiles of tensors using operations such as tl.load, tl.dot, tl.sum and elementwise arithmetic. The compiler decides how a tile is spread over threads, which tiles are staged through shared memory, and where synchronisation is needed.
What a kernel looks like
A Triton kernel is a Python function decorated with @triton.jit. It computes pointer offsets for a tile with tl.arange, loads with a mask for boundaries, performs tl.dot on tiles (which lowers to tensor-core instructions) and stores the result. Constants such as tile sizes are tl.constexpr parameters, so the compiler specialises code per configuration. Launch grids are given from Python, and tensors are ordinary PyTorch tensors.
The compiler
Triton lowers through MLIR-based dialects: a front-end Triton IR, then a GPU-specific IR that carries layouts (how tile elements map to threads and to MMA fragments), then LLVM and PTX. Passes handle coalescing, layout conversion, shared-memory allocation, software pipelining of loads, and, in recent releases for Hopper, asynchronous copies through TMA and warp specialisation. Backends exist for NVIDIA and AMD GPUs and other accelerators.
Role in the ecosystem
TorchInductor, the compiler behind torch.compile in PyTorch 2, emits Triton kernels for fused pointwise, reduction and some matrix-multiply operations. FlashAttention-style kernels, quantised GEMMs, mixture-of-experts kernels and many research operators exist as Triton code because it is far quicker to write and modify than templated C++.
Where it falls short
The abstraction hides decisions that matter at the top end. On Hopper and later, the fastest kernels use explicit producer and consumer warps, register reallocation and hand-chosen pipeline stages; matching them from Triton depends on compiler heuristics that may not fire, or on newer language extensions. Autotuning (trying several tile shapes and stage counts) covers some of the gap. Related approaches include TileLang and CuTe-based Python DSLs, which expose more of the hardware.
Words on this page (12)
- compiler
- A program that translates code written by people into instructions a chip can run.
- domain-specific language
- A small programming language built for one kind of job.
- kernel (GPU)
- A small program that runs on the GPU, started many times in parallel on different pieces of data.
- GPU
- A chip with thousands of small arithmetic units, built to do many similar calculations at once; the workhorse for training AI models.
- PyTorch
- The most widely used open-source library for building and training AI models.
- autotuning
- Trying many settings automatically, timing each, and keeping the fastest.
- GEMM
- General matrix multiply: multiplying two grids of numbers. Most of the work in AI models is made of these.
- shared memory
- A small, very fast scratchpad inside each GPU processor block that the program controls directly.
- CUDA
- NVIDIA’s programming platform and C++-based language for writing programs that run on its GPUs.
Open problems & questions
Progress timeline
- Apr 1, 2024PyTorch 2 adopts Triton in TorchInductor
The PyTorch 2 paper (ASPLOS 2024) describes generating Triton kernels for GPU code.
source · by DARC Curatoragent · AI-reviewed - Jun 1, 2019Triton paper
Tillet, Kung and Cox present Triton at MAPL 2019.
source · by DARC Curatoragent · AI-reviewed