Vol. XVI · Kernels and Efficiency Plate 5 of 8compilers, torch.compile, Triton, tiling, autotuning, tile shape, wave quantisation, graph captureCE

Compilers and Kernel Languages.

Write the multiply as tiles, let the compiler try many tile shapes, and keep the fastest — a different one for each shape.

A matrix multiply on a GPU is computed in tiles: each thread block owns a bm × bn tile of the output, steps along K in slices of bk, and keeps its slices of A and B in shared memory (plate XIII.2). Larger tiles reuse each loaded value more and move fewer bytes, but need more shared memory and produce fewer tiles — and tiles run in waves of one per multiprocessor, so 133 tiles on 132 multiprocessors take two waves, the second almost empty.

Kernel languages such as OpenAI’s Triton let the multiply be written once with the tile sizes as parameters; autotuning runs the candidates on the actual shape and keeps the fastest. Graph compilers such as PyTorch’s torch.compile capture the model, fuse element-wise chains (plate XVI.2) and pick or generate kernels like these. The workshop models a tiled multiply on an H100 and shows the time of every tile shape for four problem shapes.

Below: left, the modelled time of each tile shape bm × bn (bk = 64) for the chosen problem; darker cells are slower. Right, the best, the worst and a fixed 128 × 128 choice.

What you will learn
Workshop
A model of a tiled BF16 multiply on an H100 (132 SMs, 228 KB shared memory, 989 TFLOP/s, 3.35 TB/s): HBM traffic reduced by g, the square root of the tiles sharing panels in one wave through L2; compute time divided by the wave efficiency and by a tile efficiency that falls for tiles smaller than 128 × 128; time is the larger of the two. Not measured; it captures the trade-offs, not a particular kernel.

A different tile for every shape

For a large square multiply the model prefers big tiles — 64 × 256 at 1,146 µs — and the worst choice is 1.83 times slower; a fixed 128 × 128 does as well here (1.00×). For the small 1 000 × 1 000 × 8 192 problem the answer changes: 128 × 128 tiles make only 64 tiles for 132 multiprocessors, so half of them idle, and the best is 32 × 256, which makes enough tiles to fill a wave (97% wave efficiency); the fixed choice is 1.50 times slower. For a 16-token decode every choice is memory-bound and they all tie at 35.2 µs — the roofline floor, set by reading the weights.

That is why kernels are autotuned per shape and why compilers specialise on shapes: the same source, compiled for different sizes, produces different tile choices. Triton made this accessible by letting the tile be the unit of programming — loads, dot products and stores on blocks — while the compiler handles the threads, shared memory and tensor-core instructions. torch.compile builds on it, generating fused Triton kernels for the element-wise parts of a model and calling tuned libraries for the large multiplies.

Real kernels add choices the model omits — the pipeline depth of asynchronous copies, split-K for skinny problems, the order in which tiles are scheduled — which is why autotuning measures rather than predicts.

Terms

Tile
The block of the output one thread block computes, bm × bn.
Wave quantisation
Losing efficiency when the tile count is not a multiple of the number of multiprocessors.
Autotuning
Running candidate kernel configurations and keeping the fastest.
Graph capture
Recording a model’s operations so a compiler can optimise them together.
A tiled multiply’s best tile shape depends on the problem shape; autotuning tries them and keeps the fastest, and compilers specialise per shape.

Sources

  1. Tillet, Kung and Cox, “Triton: An Intermediate Language and Compiler for Tiled Neural Network Computations”, MAPL 2019.
  2. Ansel et al., “PyTorch 2: Faster Machine Learning Through Dynamic Python Bytecode Transformation and Graph Compilation”, ASPLOS 2024.
  3. NVIDIA CUTLASS documentation (tile shapes, wave quantisation).