Writing a fast GPU kernel by hand means warp shuffles, shared memory swizzling, and a week in the profiler, which is why most deep learning teams would rather not. Triton, started by Philippe Tillet and described in the 2019 MAPL paper Triton: An Intermediate Language and Compiler for Tiled Neural Network Computations, took a different bet: let the programmer think in blocks of a program and in tiles of a tensor, and let the compiler figure out layouts, vectorization, and shared memory. That bet won. Today the MIT-licensed repository, developed at triton-lang with OpenAI among its stewards, is the kernel language underneath PyTorch’s torch.compile, a growing share of vLLM, and uncounted research projects, shipping wheels for CPython 3.10 through 3.14 that drive both NVIDIA GPUs of compute capability 8.0+ and AMD GPUs on ROCm 6.2+. As always in this series, what follows is an educational tour of published source code.

The repository is two compilers wearing one trench coat. The Python half is the friendly part: a jit decorator that captures your kernel function, a language of tl builtins, and a compile driver that walks the code from Python AST to typed IR. The C++ half is the muscle: an MLIR codebase under lib/ with custom TT dialects, conversion passes that lower tiles through increasingly GPU-shaped forms, and backends that finish the job through LLVM into PTX for NVIDIA or GCN for AMD. Between them lives a runtime layer that handles grids, caches, autotuning, and even a CPU interpreter that can run your kernel with Python breakpoints. Reading the overview from left to right follows one kernel from decorator to machine code.

Triton repository architecture overview

Overview of Triton: the jit decorator and tl builtins flow through semantic analysis into the CodeGenerator, the compile driver runs MLIR passes over custom TT dialects, and NVIDIA plus AMD backends emit hardware code while the runtime handles grids, tuning, and the interpreter.

Reading the overview from left to right:

  • @jit is the entry decorator, defined across overloads at lines 953 to 971 of a 1198-line file that also holds JITFunction at 638.
  • tl builtins form the language at 3901 lines: dot at 2253, load at 2425, store at 2485, arange at 1882.
  • The CodeGenerator at line 294 of 1792 lines compiles the Python AST into Triton IR through MLIR’s Python bindings.
  • compile() at line 226 of 522 lines orchestrates stages from an ASTSource at 52 to a CompiledKernel at 407.
  • The MLIR core lives under lib/Dialect, with conversions under lib/Conversion and targets under lib/Target.
  • The NVIDIA backend is a 669-line compile chain ending in PTX via LLVM and ptxas, mirrored by an AMD backend at 827 lines.
  • @autotune at line 424 of 527 lines wraps kernels to search block sizes and num_warps at launch time.
  • The interpreter is a 1694-line CPU emulation, with InterpretedFunction at 1655, that runs under TRITON_INTERPRET=1 without a GPU.

Why You Need This

The first reason is the productivity gap it closes. Handwritten CUDA rewards specialists; Triton rewards domain experts who know the math but not the memory hierarchy. You describe what one block of the program computes, the compiler decides how tiles map to hardware, and the same kernel source runs on two vendors’ GPUs. The proof is adoption: when PyTorch wants to generate fast elementwise and matmul fusion code dynamically, it emits Triton, because a generated kernel in Triton is a few hundred lines of template instead of a few thousand lines of device code with per-architecture special cases.

The second reason is that the compiler is genuinely good, and its quality lives in the middle you would otherwise hand-tune. The conversion passes under lib/Analysis and lib/Conversion reason about axis information, aliasing, and data layouts, picking shared-memory stage counts, vector widths, and instruction shapes for you. Version 2.0’s rewrite onto MLIR, noted in the repository’s own changelog, is what enabled kernels with back-to-back matmuls like flash attention to schedule efficiently. You are not just avoiding CUDA boilerplate; you are inheriting a layout engine tuned across thousands of real kernels.

The third reason is observability, which custom GPU projects rarely survive without. Every stage of compilation is inspectable: MLIR_ENABLE_DUMP prints IR before every pass, LLVM_IR_ENABLE_DUMP does the same for the backend, TRITON_KERNEL_DUMP plus TRITON_KERNEL_OVERRIDE let you capture a compiled kernel and patch individual stages by hand, and knobs.py exposes hundreds of environment-driven settings in 599 documented lines. Dumps can target a single kernel with MLIR_ENABLE_DUMP=kernelName, TRITON_F32_DEFAULT adjusts fp32 rounding for kernels ported from CUDA defaults, and TRITON_REPRODUCER_PATH captures an MLIR reproducer right before a failing pass. Add autotuning that prints its winning configuration, and a CPU interpreter that turns kernel debugging into ordinary Python debugging, and you have a compiler you can actually learn from, not just use.

Triton detailed architecture

The detailed map: the language frontend with jit machinery and tl builtins, the Python compiler core with caching and warmup pools, the C++ MLIR core, both hardware backends, the runtime with tuning and the interpreter, and the tooling layer.

Start at the decorator. @jit wraps your function in a JITFunction at line 638, whose run method at 736 handles argument binding, specialization keys, and launch. The DependenciesFinder at line 39 walks the function’s globals at decoration time so a kernel can reference other Python objects, and KernelParam at 305 classifies each parameter, separating tl.constexpr values that participate in compilation from runtime scalars. The grid is Python all the way: KernelInterface.getitem at line 370 turns kernelgrid into a launch, so your grid lambda receives the compiled kernel’s constants and can compute block counts from them.

The language layer is bigger than it looks. core.py spans 3901 lines because every builtin is overloaded across dtypes and carries a _semantic hook that connects it to the type system in semantic.py, 2048 lines of broadcasting and promotion rules that decide, for example, how a pointer plus a range becomes an address vector. On top sit standard.py at 663 lines, whose softmax and reduction helpers are themselves written in Triton, and math.py at 261. The CodeGenerator at line 294 of code_generator.py then walks the AST, resolving each builtin through the semantic layer and emitting ops in the TT dialect.

The compile driver at compile() line 226 turns that IR into machine code through a stage list contributed by the active backend, timing itself with CompileTimer at 187 and returning a CompiledKernel at 407 whose AsmDict at 390 holds every intermediate: ttir, ttgir, llir, ptx or amdgcn, and the final binary. Results land in the cache of cache.py at 323 lines keyed by a hash of source, options, and environment, and a 624-line warmup pool can pre-compile kernels ahead of first use. The C++ heart does the heavy lifting: dialects under lib/Dialect, layout-aware lowering under lib/Conversion, and target emission under lib/Target.

The backends finish the story. NVIDIA runs 669 lines of stage flow through LLVM to PTX, then shells out to ptxas, with its driver at 406 lines managing modules and streams; AMD’s 827-line compiler targets GCN through ROCm. Both implement the small BaseBackend contract at line 93 of backends/compiler.py, and driver.py picks one at runtime in 49 lines, overridable with TRITON_DEFAULT_BACKEND. Around them, @autotune at line 424 of autotuner.py benchmarks Config entries at 343 across the grid, heuristics at 506 computes them cheaply, and the interpreter re-implements the language on numpy so the same kernel source runs on a laptop CPU, which is exactly how the popular Triton-Puzzles work.

From Install to a Running Kernel

The stable path is one pip command, with wheels for CPython 3.10 through 3.14:

pip install triton

From source, the README’s recipe builds against a pinned LLVM recorded in cmake/llvm-info.json:

git clone https://github.com/triton-lang/triton.git
cd triton
pip install -r python/requirements.txt
pip install -e .

And the canonical hello world, a masked vector add in the shape every Triton kernel follows:

import torch
import triton
import triton.language as tl

@triton.jit
def add_kernel(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)

n = 1024
x = torch.randn(n, device="cuda")
y = torch.randn(n, device="cuda")
out = torch.empty_like(x)
grid = lambda meta: (triton.cdiv(n, meta["BLOCK"]),)
add_kernel[grid](x, y, out, n, BLOCK=1024)

Eleven lines of kernel plus a grid lambda replace the launch configuration, bounds checks, and vectorization logic you would write by hand in CUDA, and run unchanged on AMD hardware with the interpreter available as a debug escape hatch.

Honest limits: the platform list is short, Linux only, with CPU support marked under development, so Windows and macOS users are compiling inside containers. The compiler decides layouts for you, which is the entire point, but when autotuning lands on a slow configuration there is no manual layout escape hatch short of kernel restructuring. Compile latency is real for dynamic shapes, which is exactly why the warmup pool and cache exist. The frontend rejects genuinely dynamic Python, since kernels are traced rather than interpreted, and error messages through four IR levels can take you from Python line to ptxas failure in one confusing stack. And the MLIR core is C++ with its own build system, so kernel authors rarely read it, even though that is where the performance lives. Still, the bet has paid: Triton made writing a GPU kernel something a Python programmer does on a Tuesday, and made the compiler the place where hardware-specific brilliance accumulates for everyone.

Watch PyShine on YouTube

Contents