Triton
Triton is, in its own words, “a language and compiler for writing highly efficient custom Deep-Learning primitives,” aimed at giving programmers “higher productivity than CUDA” and “higher flexibility than other existing DSLs.” It grew out of a MAPL 2019 paper by Tillet, Kung, and Cox, and the premise stated in its own documentation is that blocked programs can produce compute kernels competitive with hand-written CUDA, while staying far more flexible than prior compiler frameworks like Halide or TVM.
The distinction Triton draws against CUDA is precise. CUDA, in its own framing, is “Scalar Program, Blocked Threads”: a programmer writes one thread’s worth of scalar arithmetic, and the hardware replicates it across a block of threads, each computing one output element of a matrix multiply. Triton inverts this to “Blocked Program, Scalar Threads”: the program itself loops over tiles of the output, stepping by block sizes MB, NB, and KB, and every statement inside that loop already operates on whole blocks of A, B, and the accumulator at once. There is no per-thread code to write, because the block is the unit the language exposes.
That inversion only pays off because of what the compiler does with it. Triton relies on block-level data-flow analysis to turn a blocked program into a scheduled one, and that analysis is what drives automatic memory coalescing, thread swizzling, prefetching, vectorization, tensor-core-aware instruction selection, shared-memory allocation and synchronization, and asynchronous copy scheduling. Those are exactly the decisions a CUDA programmer makes by hand at the thread-block level described in Foundations; Triton moves them into the compiler and asks the programmer to reason about tiles instead.
Triton’s own first tutorial makes the style concrete with the same computation Foundations used to introduce CUDA: vector addition. Where vecAdd gave every thread exactly one element, the Triton kernel is written per program, and each program handles BLOCK_SIZE elements at once. The tutorial’s comments call out the two moves that define the model: the kernel asks tl.program_id which of the parallel program instances it is, and offsets “is a list of pointers,” a whole block of addresses that tl.load reads in one statement:
@triton.jit
def add_kernel(x_ptr, y_ptr, output_ptr, n_elements,
BLOCK_SIZE: tl.constexpr):
pid = tl.program_id(axis=0)
block_start = pid * BLOCK_SIZE
offsets = block_start + tl.arange(0, BLOCK_SIZE)
mask = offsets < n_elements
x = tl.load(x_ptr + offsets, mask=mask)
y = tl.load(y_ptr + offsets, mask=mask)
output = x + y
tl.store(output_ptr + offsets, output, mask=mask)
def add(x: torch.Tensor, y: torch.Tensor):
output = torch.empty_like(x)
n_elements = output.numel()
grid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE']), )
add_kernel[grid](x, y, output, n_elements, BLOCK_SIZE=1024)
return outputNothing in the kernel names a thread. The mask guards the loads and stores against out-of-bounds accesses when the vector length is not a multiple of the block size, which is the blocked model’s replacement for the per-thread bounds check of Foundations. The launch side is what the tutorial calls an SPMD grid, “analogous to CUDA launch grids”: here a 1D grid of cdiv(n_elements, BLOCK_SIZE) program instances, with each torch.Tensor argument implicitly converted into a pointer to its first element. Everything below the program, the warps, the shared memory, the vector widths, is the compiler’s problem.
That same tile abstraction is why Triton and CUTLASS keep meeting in the same place. NVIDIA has since built a backend, Triton-to-TileIR, that lets Triton programs compile directly to CUDA Tile IR (covered next) instead of PTX, preserving Triton’s tile semantics rather than lowering them to per-thread SIMT instructions first. See further reading for tutorials and the original paper.
Past vector addition: the matmul#
Vector addition hides the part of the model that costs anything, because a block of pointers into a one-dimensional array is just an arange. The matrix multiplication tutorial writes an FP16 matmul it claims “achieves performance on par with cuBLAS or rocBLAS,” and it names the four things that separate it from the first example: block-level matrix multiplications, multi-dimensional pointer arithmetic, program re-ordering for improved L2 cache hit rate, and automatic performance tuning. The algorithm underneath is the doubly nested parallel loop over output tiles from the programming guide, with each iteration of that loop “performed by a dedicated Triton program instance.”
The pointer arithmetic is the mechanical part. For a row-major 2D tensor the tutorial states the address rule as &X[i, j] = X + i*stride_xi + j*stride_xj, so a whole tile of pointers is built by broadcasting one arange down the rows, another across the columns, and scaling each by the matching stride. An extra modulo handles M or N not being a multiple of the block size, padding “with some useless values, which will not contribute to the results,” while the K dimension is handled by masking the loads instead. The kernel also carries a block of tl.assume calls the tutorial explains as integer bound hints “to guide integer analysis in the backend to optimize load/store offset address calculation”:
@triton.autotune(
configs=get_autotune_config(),
key=['M', 'N', 'K'],
)
@triton.jit
def matmul_kernel(
a_ptr, b_ptr, c_ptr,
M, N, K,
stride_am, stride_ak,
stride_bk, stride_bn,
stride_cm, stride_cn,
BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr,
BLOCK_SIZE_K: tl.constexpr, GROUP_SIZE_M: tl.constexpr,
ACTIVATION: tl.constexpr
):
pid = tl.program_id(axis=0)
num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)
num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
num_pid_in_group = GROUP_SIZE_M * num_pid_n
group_id = pid // num_pid_in_group
first_pid_m = group_id * GROUP_SIZE_M
group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)
pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m)
pid_n = (pid % num_pid_in_group) // group_size_m
offs_am = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M
offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N
offs_k = tl.arange(0, BLOCK_SIZE_K)
a_ptrs = a_ptr + (offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak)
b_ptrs = b_ptr + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn)
accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
a = tl.load(a_ptrs, mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, other=0.0)
b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k * BLOCK_SIZE_K, other=0.0)
accumulator = tl.dot(a, b, accumulator)
a_ptrs += BLOCK_SIZE_K * stride_ak
b_ptrs += BLOCK_SIZE_K * stride_bk
if ACTIVATION == "leaky_relu":
accumulator = leaky_relu(accumulator)
c = accumulator.to(tl.float16)
offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
c_ptrs = c_ptr + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :]
c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N)
tl.store(c_ptrs, c, mask=c_mask)The K loop is the whole model in six statements. accumulator is a BLOCK_SIZE_M by BLOCK_SIZE_N tile the tutorial deliberately keeps in fp32 “for higher accuracy” and converts back to fp16 only after the loop ends. tl.dot(a, b, accumulator) multiplies two loaded tiles and accumulates into it in one statement, which is where a tensor core gets used without the program ever naming one, and the two pointer blocks then advance by BLOCK_SIZE_K times the appropriate stride. Because the accumulator is still in registers at fp32 when the loop exits, arbitrary activations fuse there for free, which is exactly the customization the tutorial opened by arguing that vendor kernel libraries “cannot be easily customized to accommodate the needs of modern deep learning workloads.”
The one scheduling decision Triton leaves to the author is which program computes which tile, and the tutorial is blunt that the obvious answer is wrong: a simple row-major ordering of pid into (pid_m, pid_n) “is just not going to cut it.” Grouping GROUP_SIZE_M rows of output blocks before moving to the next column changes what the L2 sees. In the tutorial’s 9-by-9-block example, row-major ordering loads 90 blocks into SRAM to produce the first 9 output blocks while grouped ordering loads 54, and it reports the reordering improving the kernel “by more than 10% on some hardware architecture (e.g., 220 to 245 TFLOPS on A100).” The compiler owns the mapping onto threads; the mapping from program id onto work is still arithmetic the programmer writes, and it is worth double-digit percentages.
What the autotuner searches#
The listing above opens with a decorator, not a kernel. Triton documents @triton.autotune as consuming two things: a list of triton.Config objects “that define different configurations of meta-parameters (e.g., BLOCK_SIZE_M) and compilation options (e.g., num_warps) to try,” and “an auto-tuning key whose change in values will trigger evaluation of all the provided configs.” The matmul kernel passes key=['M', 'N', 'K'], so every new problem shape reruns the search and every repeat of a shape already seen reuses its winner.
What a Config holds is broader than tile sizes. Beside the kwargs of meta-parameters it carries num_warps, where the docs note that “if num_warps=8, then each kernel instance will be automatically parallelized to cooperatively execute using 8 * 32 = 256 threads”; num_stages, “the number of stages that the compiler should use when software-pipelining loops,” described as “mostly useful for matrix multiplication workloads on SM80+ GPUs”; num_ctas, the “number of blocks in a block cluster” and SM90+ only; and maxnreg, a cap on registers per thread corresponding to the PTX .maxnreg directive. Those are the occupancy and latency-hiding levers of Foundations exactly, moved out of the source and into a list of candidates.
The tutorial’s own candidate list shows how wide that is. Its CUDA configs range over block sizes of 32 to 256 in M and N and 32 to 128 in K, num_stages of 3 to 5, num_warps of 2 to 8, a fixed GROUP_SIZE_M of 8, and a separate group of entries the source marks as good for fp8 inputs. The HIP list beside it is a different set entirely: every entry at num_warps=8 and num_stages=2, GROUP_SIZE_M of 4 or 6, and an extra matrix_instr_nonkdim of 16 that has no NVIDIA counterpart. One kernel, two backends, two candidate lists, and no single set of constants that is correct for both.
Searching is not free, and the API is shaped around the cost. Because “when all the configurations are evaluated, the kernel will run multiple times,” anything the kernel updates gets updated once per trial, so reset_to_zero names arguments to zero before each config and restore_value names arguments to put back afterwards. prune_configs_by accepts a perf_model and a top_k so only the most promising configs are actually benchmarked, or an early_config_prune function that filters the list before any of them run. TRITON_PRINT_AUTOTUNING=1 prints the winning config and the time the search took, per kernel.
This replaces hand-tuning rather than automating it. In chapter 1’s CUDA the block size, the pipeline depth, and the register budget are constants in the source, and retuning them for a new shape or a new device means editing and recompiling. The Triton compiler already owns coalescing, swizzling, shared-memory allocation, and asynchronous copies; what it cannot know at compile time is which tile shape and pipeline depth make those decisions pay off for a particular problem on a particular GPU. The autotuner closes that gap by measuring instead of predicting, and key reopens it the moment the thing it measured against changes.
Inspecting a kernel you did not schedule#
Once the compiler owns the schedule, reading the Python tells you less about what ran than reading CUDA did, so Triton ships a set of ways to look. The language itself has four debugging operators, in two pairs: static_print and static_assert “are intended for compile-time debugging,” while device_print and device_assert “are used for runtime debugging.” The runtime assert is off by default and “executes only when TRITON_DEBUG is set to 1,” the other three run regardless.
The bigger tool is the interpreter. Setting TRITON_INTERPRET=1 “causes all Triton kernels to bypass compilation and be simulated by the interpreter using numpy equivalents of Triton operations,” and the interpreter “processes each Triton program instance sequentially, executing operations one at a time.” That buys ordinary Python debugging: print(tensor) for a whole tile, print(tensor.handle.data[idx]) for one value, and pdb either attached to the run or invoked from a pdb.set_trace() inside the kernel body. The documented limits are worth knowing before relying on it: the interpreter does not support bfloat16, so a bf16 tensor has to be cast to float32 first, and it does not support indirect memory access patterns where a loaded value is itself used as a pointer.
Below the language, the compiler is a pipeline of intermediate representations, and each stage has a knob. MLIR_ENABLE_DUMP=1 “dumps the IR before every MLIR pass Triton runs, for all kernels,” narrowable to one kernel by name and redirectable from stderr to a file with MLIR_DUMP_PATH. LLVM_IR_ENABLE_DUMP=1 does the same for the LLVM passes below it, and TRITON_KERNEL_DUMP “enables the dumping of the IR from each compilation stage and the final ptx/amdgcn” into the directory named by TRITON_DUMP_DIR. Because the compilation cache can make all of this silently produce nothing, the documented first move when a dump fails to appear is to clear ~/.triton/cache, and TRITON_ALWAYS_COMPILE=1 forces compilation regardless of a cache hit.
Two of those knobs turn inspection into intervention. TRITON_KERNEL_OVERRIDE “enables the override of the compiled kernel with a user-specified IR/ptx/amdgcn at the beginning of each compilation stage,” reading from TRITON_OVERRIDE_DIR: dump once, copy the kernel’s hash directory across, edit the stage in question, rerun. USE_IR_LOC={ttir,ttgir} “reparses the IR such that the location information will be the line number of the IR file with that particular extension, instead of line number of the python file,” which “can provide a direct mapping from the IR to llir/ptx” and, used with performance tools, “a breakdown on IR instructions.” That is how a Triton kernel gets attributed at instruction granularity, and it is the bridge into the profilers later in this chapter. For the failures a profiler never reports, the guide points outward instead: prepend compute-sanitizer on NVIDIA GPUs, use LLVM’s AddressSanitizer on ROCm, and reach for triton-viz, described as agnostic to the underlying GPU, to visualize memory access.