Triton and compiler-generated kernels

In one paragraph

Triton is a Python-embedded language for writing GPU kernels at the level of blocks of data rather than individual threads. The compiler takes care of much of what makes CUDA hard: coalescing, staging data through shared memory, and scheduling. PyTorch's torch.compile generates Triton kernels automatically through TorchInductor, which covers most fusion opportunities without anyone writing a kernel.

How it works

• Programs and blocks. Each program instance works on a block of data, identified by tl.program_id. Loads and stores take masks for blocks that run past the end of the data.
• Autotuning. triton.autotune tries configurations such as block sizes, warp counts and pipeline stages, and caches the fastest per input shape.
• TorchInductor. The default torch.compile backend fuses elementwise and reduction operations into generated Triton kernels.
• Helion. A higher-level PyTorch kernel language that compiles to Triton, currently in beta.

Design choices and trade-offs

• Triton vs CUDA. Triton is far quicker to write and port across NVIDIA generations; exact warp specialisation, register control and some of the newest hardware features are easier from CUDA, CUTLASS or the CuTe DSL.
• Compiler first. Before writing a kernel, check what torch.compile already produces.

When not to use it

When the operation is already a single library call (a large GEMM in cuBLAS), Triton rarely beats it.

How we evaluate it

Against the library or compiler baseline, on the shapes the model actually uses, with correctness checked against a higher-precision reference.

Sources

• Triton programming guide (opens in a new tab)
• Helion (opens in a new tab)
• GPU kernel engineering, which includes a fused softmax in Triton