The tiled matmul of §11.5 took two dozen lines of CUDA before any of the refinements that make it fast, and a competitive one runs to thousands of lines of templates. Most of those lines are not about the math. They are about which thread loads which element, how tiles are laid out in shared memory to avoid bank conflicts, when to synchronize, and how to feed the tensor cores. What if a compiler handled all of that, and you wrote only the tile-level algorithm?

That is Triton (Tillet, Kung & Cox, 2019). In Triton you do not program threads. You program blocks. Each instance of your kernel handles a whole tile of data, written as operations on small arrays, and the compiler decides how the block's threads share the work, coalesces the loads, stages tiles through shared memory, and lowers tl.dot onto tensor cores. A matmul that is competitive with cuBLAS fits on one screen of Python. Triton is also what torch.compile generates for GPUs (§11.7), and many of the fused kernels used in training and serving, from attention variants to fused cross-entropy losses, are written in it.

This section teaches the model with three kernels of growing difficulty, a vector add, a fused softmax, and an autotuned tiled matmul, each emulated in NumPy so you can run the logic here. It ends with a guide to when writing a kernel is worth it at all.

The block programming model

A Triton kernel is a Python function decorated with @triton.jit, launched on a grid of programs. Each program is one instance of the kernel, roughly one CUDA thread block, and it asks tl.program_id(axis) which one it is. Instead of a scalar index, it builds a block of indices with tl.arange(0, BLOCK), where BLOCK is a compile-time constant (tl.constexpr) and a power of two. Pointer arithmetic on blocks gives blocks of pointers, and tl.load and tl.store move whole blocks between memory and the program's registers, with a mask to switch off out-of-range lanes.

In plain words, where CUDA asks "what does thread $i$ do?", Triton asks "what does this program do to its tile?".

Here is what each layer takes care of.

How does the compiler do it? It turns the decorated function into its own intermediate representation, assigns every block a layout that says which thread and register hold which element, chooses shared memory layouts and swizzles for the operands of tl.dot, and lowers the result through MLIR and LLVM to PTX, the same assembly nvcc produces, compiling once per combination of constexpr values and argument types. What you keep is the part that needs algorithmic judgment, how to split the problem into tiles and in what order to visit them. What you give up is control over the lowest level, which matters only for the last few percent on the newest hardware.

A vector add

The Triton version of the vector add from §11.3 makes the model concrete. Each program handles BLOCK consecutive elements, and the mask does the bounds check that CUDA wrote as an if.

The grid is a function of the kernel's meta-parameters, so that it stays right when BLOCK changes (the autotuner below relies on this). Tensors are passed directly and become pointers to their first element, which is why the wrapper insists on contiguous inputs. For debugging, setting the environment variable TRITON_INTERPRET=1 runs kernels on the CPU with NumPy semantics, so ordinary print and breakpoints work.

The same semantics are easy to emulate. Below, a program is a Python function of its pid, arange makes the block of offsets, and masked loads fill switched-off lanes with a default value, exactly as tl.load(..., mask=..., other=...) does.

Two-dimensional blocks

Most real kernels work on 2D tiles, and the central Triton idiom is building a 2D block of offsets by broadcasting a column of row indices against a row of column indices. For a row-major matrix with row stride stride, the tile of rows offs_m and columns offs_n lives at offs_m[:, None] * stride + offs_n[None, :], and its mask is (offs_m[:, None] < M) & (offs_n[None, :] < N). Every kernel in the rest of this section uses that line.

import numpy as np

def bias_relu_program(pid_m, pid_n, X, bias, out, M, N, BM, BN):
    offs = pid_m * BM + np.arange(BM * BN)           # placeholder: treats the tile as 1-D
    mask = offs < M * N
    tl_store(out, offs, np.maximum(tl_load(X, offs, mask), 0), mask)

def launch_bias_relu(X, bias, M, N, BM, BN):

A fused softmax

Softmax is memory-bound (§11.2), so its speed is set by the bytes it moves, and the way to make it faster is to move fewer. Suppose an $M\times N$ input is processed by separate kernels, one per step of the stable softmax of §4.3. Count the elements each one moves.

A fused kernel reads each row once into registers, does all five steps there, and writes the result once, $MN$ reads and $MN$ writes. That is about four times fewer bytes, and since the kernel is memory-bound, close to four times faster. One program per row is the natural decomposition.

Two details carry the correctness. Padded lanes load $-\infty$, so they never win the maximum and contribute $e^{-\infty} = 0$ to the sum. And tl.max and tl.sum are reductions across the block, which the compiler implements with the warp shuffles and shared memory of §11.4.

Rows that do not fit, the online softmax

The kernel above holds a whole row in registers. That works for rows of a few thousand elements, but a language model's output row has a vocabulary's worth of logits, 128,256 for Llama 3, and does not fit. The fix is the online-softmax operator of §11.4 (Milakov & Gimelshein, 2018). Stream the row through in tiles, keeping a running maximum $m$ and a running sum $\ell = \sum_j e^{x_j - m}$ of exponentials relative to it. Each tile contributes its own maximum $\tilde m$ and sum $\tilde\ell = \sum_{\text{tile}} e^{x_j - \tilde m}$, and the merge is

In plain words, when a bigger maximum shows up, shrink everything summed so far by the factor it would have been shrunk by had that maximum been known all along. Every exponent is at most zero, so nothing overflows. §12.2 derives this rule and proves it exact, and it is the heart of FlashAttention. Here we need only the kernel's structure. One pass over the tiles computes $(m, \ell)$, and a second pass reads the row again and writes $e^{x_j - m}/\ell$, three transfers of the row instead of two, and still far fewer than the eight of the unfused version. In Triton, the first pass is a loop inside one program, with tl.max and tl.sum applied to each tile.

import numpy as np

def online_stats(tiles):
    x = np.concatenate(list(tiles), axis=1)          # placeholder: materializes the whole row
    m = x.max(1)
    return m, np.exp(x - m[:, None]).sum(1)

def softmax_tiled(x, block):

A tiled matmul, with grouped ordering and autotuning

Now the main event, $\mathbf{C} = \mathbf{A}\mathbf{B}$ with $\mathbf{A} \in \R^{M\times K}$ and $\mathbf{B} \in \R^{K\times N}$. Each program computes one BM by BN tile of $\mathbf{C}$, looping over $K$ in steps of BK, the same structure as the shared-memory kernel of §11.5. But the loads are whole 2D blocks, the multiply is tl.dot, and the shared memory, barriers, register tiling, and tensor core instructions are all generated.

Passing strides rather than assuming row-major makes the kernel work on transposed views for free, which is how the same kernel serves $\mathbf{A}\mathbf{B}$, $\mathbf{A}^\top\mathbf{B}$, and $\mathbf{A}\mathbf{B}^\top$ in a backward pass. The accumulator is float32 even for bfloat16 inputs, and only the final store casts down, which keeps long sums accurate (the summation error of §11.4, and the formats of §8.5).

Why the program order matters

The first few lines of the kernel remap the program id, and they are worth a closer look. Programs are launched roughly in order of pid, and the ones running at the same time share the 50 MB L2 cache. With plain row-major order, pid_m = pid // n_tiles_n, the first wave of programs (132 of them on an H100) covers whole rows of tiles. For an $8192\times8192$ output in $128\times128$ tiles, that is two rows and a bit, so the wave needs 3 row panels of $\mathbf{A}$ but all 64 column panels of $\mathbf{B}$. With grouped order, consecutive programs walk down a group of GROUP_M rows of tiles before moving right, so the same wave covers about an $8\times 17$ block of tiles and needs only 8 panels of $\mathbf{A}$ and 17 of $\mathbf{B}$. Fewer distinct panels means more L2 hits and less HBM traffic for the same kernel.

With row-major order the first wave touches about 140 MB of operands, far more than L2 holds, so panels of $\mathbf{B}$ are evicted before the next row of tiles can reuse them. With groups of 8 it touches about 50 MB. The Triton tutorial on matrix multiplication reports that this remapping alone took its kernel from about 220 to 245 TFLOP/s on an A100, more than 10 percent.

def grouped_tile(pid, n_m, n_n, group_m):
    return pid // n_n, pid % n_n                      # placeholder: plain row-major order

Autotuning

The decorator @triton.autotune lists candidate configurations, each a choice of tile sizes, num_warps (how many warps run one program), and num_stages (how many tiles of $\mathbf{A}$ and $\mathbf{B}$ the compiler keeps in flight in shared memory, so loads for later iterations overlap the math of the current one). The first time the kernel is called with a new value of the key arguments, Triton benchmarks every configuration and caches the fastest.

The best choice depends on the shape and the GPU, which is why tuning beats any fixed rule. The constraints from §11.1 and §11.5 still decide which configurations are even sensible.

• Shared memory. Each stage holds a BM by BK tile of $\mathbf{A}$ and a BK by BN tile of $\mathbf{B}$, so a configuration needs about num_stages * (BM + BN) * BK * s bytes, with $s$ bytes per element. At BM = BN = 128, BK = 64, and 4 stages in bfloat16 that is 128 KB, within the H100's 228 KB but over the limit of many consumer GPUs.
• Registers. The BM by BN float32 accumulator is spread across the program's threads, so a $128\times256$ tile over 8 warps (256 threads) is 128 accumulators per thread before anything else.
• Intensity. Per program, the tile's intensity against L2 or HBM is $2\,b_Mb_N/(s(b_M + b_N))$ with $b_M$ = BM and $b_N$ = BN (§11.5), about 85 FLOPs per byte for $128\times256$ in bfloat16, so larger tiles are better until they no longer fit. Even 85 is below the H100's ridge point of 295, so a fast matmul also relies on concurrent programs finding their operand panels in L2, which is exactly what the grouped ordering arranges.
• Waves. Larger tiles mean fewer programs, and a skinny problem may not fill the GPU's SMs at all, or may waste most of its last wave.

The explorer computes all four for a configuration and a problem shape.

Tuning has costs. Benchmarking eight configurations at each new shape can take seconds, so dynamic shapes (variable sequence lengths) need a small key or bucketing. Timings on a busy or thermally throttled GPU are noisy, so tune on an idle machine. And a configuration tuned on an H100 can be poor or invalid on an A100, so keep the search space broad rather than shipping one winner. Setting TRITON_PRINT_AUTOTUNING=1 shows what was chosen.

When should you write a kernel?

Almost never first. The honest order of operations is a ladder, and you climb it only when a profile (§11.7) shows the time is going somewhere the current rung cannot reach.

1. Use the library. cuBLAS for matmuls, cuDNN and scaled_dot_product_attention (which dispatches to FlashAttention) for attention. You will not beat them on their own problems.
2. Use torch.compile. It fuses chains of elementwise and reduction operations into generated Triton kernels automatically, which covers most memory-bound waste.
3. Write Triton when there is a fusion the compiler misses or a computation no library offers. Typical cases are an attention variant with a custom mask or score modification (often expressible in FlexAttention, itself Triton underneath), a loss fused with the final projection so the $B\times T\times V$ logits never reach HBM, quantized matmuls that dequantize weights inside the tile loop, and long chains of memory-bound operations with awkward shapes.
4. Write CUDA, CUTLASS, CuTe, or ThunderKittens when you need hardware features Triton does not expose well, such as explicit warp specialization, TMA multicast across clusters, or the last 10 to 30 percent on the newest GPU. FlashAttention-3 is written this way.

A custom kernel is also a maintenance commitment. It needs a backward pass (often harder than the forward), tests against a reference in float64 on awkward shapes, handling of non-contiguous inputs, and re-tuning on new hardware. The payoff has to be real, measured in the end-to-end step time, not in one kernel's microbenchmark.

In PyTorch

The CPU shows the fusion effect too. The block computes softmax as five separate tensor operations, each producing a full intermediate tensor, then with PyTorch's single fused softmax, checks they agree, and times both.

Read more

• Tillet, Kung & Cox (2019). Triton: An Intermediate Language and Compiler for Tiled Neural Network Computations. The paper that introduced the block-level programming model and its compiler.
• Triton. Tutorials. Vector add, fused softmax, matrix multiplication with grouped ordering, LayerNorm, and fused attention, each with benchmarks.
• Milakov & Gimelshein (2018). Online normalizer calculation for softmax. The running maximum and sum that make one-pass softmax statistics possible.
• OpenAI (2021). Introducing Triton: Open-source GPU programming for neural networks. The motivation for Triton and a comparison with CUDA, from its authors.
• Hsu et al. (2024). Liger Kernel: Efficient Triton Kernels for LLM Training. A collection of fused Triton kernels (RMSNorm, RoPE, SwiGLU, fused linear cross-entropy) with measured memory and speed gains.
• Spector et al. (2024). ThunderKittens: Simple, Fast, and Adorable AI Kernels. A tile-based CUDA library for when you need control below Triton's level.

Watch next

▶

▶

▶