Skand Tandon
Standalone · Kernel

Triton: the compiler that pretends to be a library

The @triton.jit decorator does not decorate a function. It parses the function's AST, runs it through an MLIR pipeline, and emits a GPU binary. The Python function never runs as Python. This is how Triton compiles code, what it decides on your behalf, and where the abstraction falls short.

2026-08-20 kernel compiler 20 min read

Triton is a compiler with a Python frontend. The @triton.jit decorator does not decorate a function. It parses the function's AST, runs it through an MLIR pipeline, and emits a GPU binary. The Python function never runs as Python.

This matters because it sets what Triton can and cannot do. It tiles a GEMM for you, stages data through shared memory, maps blocks to warps, and picks tensor core instructions. You never write a line of PTX or manage a thread index. But you also cannot reach past the abstraction when it gets in the way.

The programming model: blocks, not threads

In CUDA you write code for one thread and the hardware runs it across thousands. Triton flips this. You write code for a block of data and the compiler splits it across threads.

import triton
import triton.language as tl

@triton.jit
def softmax_kernel(
    input_ptr, output_ptr,
    n_cols,
    BLOCK_SIZE: tl.constexpr,
):
    # 1. Each "program" handles one row. No thread indexing.
    row_idx = tl.program_id(0)
    col_offsets = tl.arange(0, BLOCK_SIZE)
    mask = col_offsets < n_cols

    # 2. Load an entire row as a block. The compiler picks
    #    coalescing, vector width, and predication.
    row = tl.load(input_ptr + row_idx * n_cols + col_offsets, mask=mask, other=-float('inf'))

    # 3. Block-level reductions. The compiler turns these into
    #    warp shuffles and shared memory reductions.
    row_max = tl.max(row, axis=0)
    numerator = tl.exp(row - row_max)
    denominator = tl.sum(numerator, axis=0)
    result = numerator / denominator

    tl.store(output_ptr + row_idx * n_cols + col_offsets, result, mask=mask)

The key operations:

You never write threadIdx.x. There is no __syncthreads(). There is no shared memory declaration. The compiler handles all of it. That is both the selling point and the ceiling.

The compilation pipeline

Triton lowers code through four IRs: Python AST → Triton IR (TTIR) → Triton GPU IR (TTGIR) → LLVM IR → PTX → cubin.

Stage 1: Python AST → Triton IR (TTIR). On first call with concrete arguments, Triton parses the Python AST and traces it into TTIR, a hardware-independent MLIR dialect (tt namespace). Operations stay abstract here: tl.load becomes tt.load, tl.dot becomes tt.dot. Standard compiler passes run: constant folding, CSE, dead code removal. TTIR knows nothing about threads, warps, or shared memory. It works on tensors whose shapes come from the constexpr parameters.

Stage 2: Triton IR → Triton GPU IR (TTGIR). This is where the compiler makes its hard calls. TTGIR adds GPU-specific structure:

This stage is where most of Triton's value lives. It is also where most of its performance bugs come from. A few real examples from the Triton issue tracker:

The pattern is consistent: TTGIR layout decisions are the single biggest lever on Triton kernel performance, and the compiler does not always get them right. Diagnosing it means reading TTGIR dumps.

Stage 3: TTGIR → LLVM IR. The GPU-specific MLIR gets lowered to plain LLVM IR. By now the tile ops have been broken into per-thread scalar ops. LLVM handles register allocation, instruction selection, and scheduling. The target is nvptx64.

Stage 4: LLVM IR → PTX → cubin. LLVM's NVPTX backend emits PTX. Triton then calls ptxas to produce the cubin. The binary gets cached on disk, keyed by a hash of the source, the constexpr values, and the target architecture. Later calls with the same inputs skip the whole pipeline and load the cached binary.

What the IR actually looks like

Take a vector add kernel and trace it through TTIR and TTGIR.

The Python source:

@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)

TTIR (simplified, from TRITON_KERNEL_DUMP=1):

// Everything is tensor-typed. No threads, no warps.
module {
  tt.func @add_kernel(%x: !tt.ptr<f32>, %y: !tt.ptr<f32>,
                       %out: !tt.ptr<f32>, %n: i32) {
    %pid  = tt.get_program_id {axis = 0} : i32
    %offs = tt.make_range {start = 0, end = 1024} : tensor<1024xi32>
    %base = tt.splat %pid : tensor<1024xi32>
    %idx  = arith.addi %base, %offs : tensor<1024xi32>
    %mask = arith.cmpi slt, %idx, %n_splat : tensor<1024xi1>
    %xv   = tt.load %x_ptrs, %mask : tensor<1024xf32>   // ← block load
    %yv   = tt.load %y_ptrs, %mask : tensor<1024xf32>
    %sum  = arith.addf %xv, %yv : tensor<1024xf32>
    tt.store %out_ptrs, %sum, %mask : tensor<1024xf32>   // ← block store
    tt.return
  }
}

Notice: no shared memory, no thread indices, no layout annotations. TTIR is pure math on tensors. TTGIR is where the layout decisions, shared memory allocation, and thread mapping appear.

Mental model

Triton is a chef who does all the knife work for you: it slices the data into blocks, stages it in the right pans (shared memory), and picks the right burner (tensor core instruction). You just write the recipe. But if the chef's default plating (layout) is wrong, the dish comes out 2-5x slower, and you have to read the kitchen log (TTGIR) to fix it.

Where Triton fits in the ML systems stack

Triton sits between hand-written CUDA and a full compiler like XLA. It is the productivity layer for custom kernels: the performance of a tuned kernel, the ergonomics of Python, and the portability of a compiler target. It is why vLLM, PyTorch's Inductor, and JAX all emit Triton for their custom kernels.

For inference engineering, Triton is the tool you reach for when you need a kernel that does not exist yet, or when the existing one is not fast enough on your hardware. The tradeoff is the layout system: the compiler's decisions are mostly right, occasionally wrong, and always worth checking with a TTGIR dump when the numbers do not add up.

Triton is a compiler that pretends to be a library. The abstraction is the point, and the layout system is the price.

Back to the blog