Triton autotuning

The sweep

@triton.autotune(
    configs=[
        triton.Config({'BLOCK_M': 128, 'BLOCK_N': 128, 'BLOCK_K': 32}, num_warps=8, num_stages=3),
        triton.Config({'BLOCK_M': 64, 'BLOCK_N': 64, 'BLOCK_K': 32}, num_warps=4, num_stages=4),
        triton.Config({'BLOCK_M': 128, 'BLOCK_N': 64, 'BLOCK_K': 64}, num_warps=4, num_stages=2),
    ],
    key=['M', 'N', 'K'],
)
@triton.jit
def matmul_kernel(...):
    ...

key names runtime arguments, not constexprs. Every distinct (M, N, K) triggers a fresh sweep, because a tile that fills the SMs at 4096 is a waste at 512. Each config is compiled, run on your actual data, timed, and the winner is cached for that key.

What the knobs mean

BLOCK_M, BLOCK_N and BLOCK_K set the tile and therefore the arithmetic intensity. num_warps sets how many threads cooperate on one program, and Config defaults it to 4, so 128 threads. num_stages is the software-pipelining depth of the K loop, useful mainly for matmul on SM80 and newer; num_ctas is cluster size on SM90 and newer. The K loop must stay a range for the pipeliner to overlap iterations, since tl.static_range unrolls it at compile time instead, and one loop can override the config with tl.range(..., num_stages=...). Larger tiles mean more registers per thread and fewer resident blocks, and past a point the compiler spills to local memory, so the tile that looked best on paper loses.

Record the winner

Autotuning is a measurement, so it needs provenance. TRITON_PRINT_AUTOTUNING=1 prints the winning config and the tuning time, cache_results=True keeps timings on disk, and your results file should store the config beside the shape and the Triton version. A tuned number without its config is not reproducible, because the next release can choose differently.

Things that break

A sweep on a noisy rig can pick a config that won by less than the band, so run it twice and check the winner is stable. A config tuned at one dtype does not transfer, since tf32 and fp16 use different tensor-core shapes, so either add the dtype to the key or tune per dtype. And a kernel that writes an output buffer needs reset_to_zero, or the correctness check that follows tuning compares against a buffer every config has already touched. These snippets follow the Triton 3.9 docs; confirm the triton.Config keyword names against your installed version, since num_stages defaults and the decorator arguments have changed.