Triton

Definition

Triton (Tillet, Kung & Cox, MAPL 2019; open-sourced by OpenAI 2021) is a Python-embedded language and MLIR-based compiler for writing GPU kernels at block granularity — you author the program for one block (a program instance) operating on tiles (small dense sub-tensors), and the compiler synthesizes the thread-level SIMT code. It sits one altitude above CUDA’s per-thread model, automating the tedious, error-prone parts (coalescing, SMEM staging, sync, pipelining) while leaving you the parts that need judgment (grid, tile shapes, program structure).

How it works

Block-level programming model. There is no threadIdx. tl.program_id(axis) gives the block index; tl.load(ptr + tl.arange(0, BLOCK)) loads a whole tile (the compiler spreads it across the block’s threads and coalesces it). Core vocabulary: @triton.jit, tl.arange, tl.load/store(mask=, other=), tl.dot(a,b,acc) (compiles to mma/wgmma Tensor-Core instructions when dims are multiples of 16), tl.max/sum reductions, and BLOCK_SIZE: tl.constexpr (baked in, enabling unroll/vectorization). Masking is correctness: masked loads protect numerical correctness (fill with the reduction identity — -inf for max, 0 for sum), masked stores protect memory safety. The SMEM-tiling choreography of a hand CUDA GEMM collapses to acc = tl.dot(a, b, acc) in a plain Python K-loop, with SMEM staging and double-buffering implicit.

What the compiler owns: intra-block thread mapping, coalescing, SMEM allocation for tl.dot operands (via liveness analysis), and software pipelining (num_stages). What you keep: grid size, tile shapes, loop/program structure, and coarse knobs num_warps/num_stages (see autotuning).

MLIR lowering to PTX. A progressive, mostly-MLIR pipeline:
$\text{Python AST}\to \text{TTIR}\to \text{TTGIR}\to \text{LLVM IR}\to \text{PTX}(\to \text{ptxas}\to \text{cubin})$

• TTIR (Triton IR, MLIR): hardware-independent — inlining, CSE, DCE, LICM, unrolling.
• TTGIR (TritonGPU IR, MLIR): the interesting stage — layout encodings (#blocked, #shared, nvidia_mma) decide how each tile distributes across warps/lanes; here run coalescing, SMEM allocation, pipelining, and Tensor-Core (mma/wgmma) lowering, plus vendor passes (NVIDIA TMA, AMD LDS).
• LLVM → PTX on NVIDIA, or AMDGCN → hsaco on AMD — one source, both backends. (Triton’s own MLIR passes are ~4/5 of compile time, not ptxas.) Dump IRs with TRITON_KERNEL_DUMP=1.

What Triton cannot express: anything below the tile abstraction — bespoke warp specialization, custom async-copy schedules, hand-laid register allocation — which is where CUDA/CUTLASS remains necessary.

Why it matters

Triton delivers ~80–95% of hand-CUDA performance for most memory-bound and many compute-bound ops at a fraction of the authoring cost, and it is the substrate of the modern kernel ecosystem: the FlashAttention tutorial, torch.compile/TorchInductor’s generated kernels, and production libraries (Liger-Kernel, Unsloth) are all Triton. It is the default first stop when torch.compile can’t express a needed fusion.

Taught in

• 07-triton-and-modern-kernels — the full lesson: model, worked kernels, autotuning, MLIR lowering, decision framework.
• 06-memory-wall-and-flashattention — fused attention as a Triton kernel.

See also

• torch-compile — generates Triton automatically
• kernel-autotuning — @triton.autotune
• kernel-fusion — what Triton makes easy
• flash-attention · online-softmax
• cuda-execution-model — the SIMT model Triton abstracts
• gemm · tensor-core · memory-coalescing · shared-memory-tiling
• gpu-systems-for-llms