Kernels & Compilers advanced 9 min read 7 flashcards

Autotuning GPU Kernels

Autotuning systematically searches a discrete configuration space of tile sizes, warp counts, and pipeline stages to find the fastest kernel for a given GPU and problem shape, replacing manual heuristics with empirical benchmarking.

A single matrix-multiplication kernel can run at 40% peak FLOPS or 95% peak FLOPS on the same GPU, with the only difference being the choice of three integers: the tile height, tile width, and the K-dimension block size. Getting those wrong does not crash the program; it just silently hemorrhages throughput. That gap is the problem autotuning solves.

Why static heuristics break down

GPU performance is not a smooth function of problem size. It is lumpy because of discrete hardware resources: shared-memory banks, register file partitions, warp schedulers, L2 cache sets. A tile size of 128x128 with 4 pipeline stages may be optimal for a 4096x4096 matrix on an A100 but suboptimal on an H100 with more SRAM, and catastrophically slow on a workstation RTX 4090 with different register file pressure.

Library authors (cuBLAS, FlashAttention) invest person-years writing hand-tuned schedules for a fixed set of shapes and hardware generations. That approach does not transfer. Every new GPU SKU, every new operator shape (e.g., the 4096x2048 projections in a 7B LLM vs. the 1x4096 decode step), and every change to mixed-precision strategy potentially invalidates the heuristic.

Autotuning frames this as a search problem:

  • Search space: a finite set of candidate configurations, each a tuple (BLOCK_M, BLOCK_N, BLOCK_K, num_warps, num_stages, ...).
  • Objective: wall-clock latency (or throughput) measured on the actual target hardware.
  • Policy: exhaustive grid search (Triton's default), Bayesian optimisation, evolutionary search, or learned cost models.

The key insight is that measurement beats prediction: a microsecond benchmark run during a one-time compile step is cheaper than a performance engineer.

The configuration space and its knobs

Understanding what you are searching over matters before understanding how to search.

Parameter Effect Typical range
BLOCK_M, BLOCK_N Tile footprint in output matrix; controls reuse per load 32-256 (powers of 2)
BLOCK_K Accumulation depth per inner loop; balances arithmetic intensity vs. shared-mem 16-128
num_warps Warps per thread block; affects occupancy and register pressure 2-16
num_stages Software pipeline depth; hides global-memory latency 1-7
num_ctas Number of thread-block clusters (Hopper+) 1-4

These interact non-linearly. A large BLOCK_M * BLOCK_N tile increases arithmetic intensity (good) but also increases the register footprint per thread, which caps the number of thread blocks that can live simultaneously on a Streaming Multiprocessor (SM). That cap is occupancy, and low occupancy means the warp scheduler has fewer warps to hide memory latency behind. The optimal point depends on the ratio of arithmetic-to-memory latency for the specific GPU's memory subsystem.

A rough heuristic for tile size selection:

target_reuse = arithmetic_intensity(op) / peak_compute_throughput
               / peak_memory_bandwidth
# if target_reuse > 1: compute bound, maximise tile; else: memory bound, smaller tiles + fuse

But this is just the starting point. Autotuning replaces the derivation with measurement.

Triton's @triton.autotune decorator

Triton makes the search explicit and declarative. You annotate a kernel with a list of triton.Config objects and a set of key dimensions. At first call, Triton benchmarks every config and caches the winner keyed to the shape:

@triton.autotune(
    configs=[
        triton.Config({'BLOCK_M': 128, 'BLOCK_N': 256, 'BLOCK_K': 64,
                       'GROUP_SIZE_M': 8}, num_stages=3, num_warps=8),
        triton.Config({'BLOCK_M': 64,  'BLOCK_N': 128, 'BLOCK_K': 32,
                       'GROUP_SIZE_M': 8}, num_stages=4, num_warps=4),
        # ... more configs ...
    ],
    key=['M', 'N', 'K'],   # re-tune when these change
)
@triton.jit
def matmul_kernel(a_ptr, b_ptr, c_ptr, M, N, K, ...):
    ...

The key argument is important. If you specify key=['M', 'N', 'K'], Triton re-runs the search whenever those values change. In a serving workload with many different prompt lengths, this means autotuning happens at each new shape, which can add cold-start latency of several seconds on first encounter. Libraries like vLLM pre-warm a fixed set of shapes at server startup to avoid this.

The official Triton matrix-multiplication tutorial demonstrates this pattern with 16+ configs for CUDA and a separate set for HIP/ROCm, showing that the winning config differs across GPU families. See Triton matrix multiplication tutorial.

torch.compile and automatic autotuning

torch.compile (PyTorch 2.x) adds autotuning lower down the stack. When you call:

model = torch.compile(model)

the compilation pipeline runs through TorchDynamo (graph capture), AOTAutograd (backward graph), and then Inductor, which is the default code-generation backend. Inductor lowers ATen ops to Triton kernels and runs its own autotuning loop over a set of candidate tile sizes, pipelining depths, and vectorisation widths.

This is "autotuning by default" for PyTorch users who do not write Triton manually. The trade-off: Inductor's search space is smaller and more conservative than a hand-written Triton kernel's space, because Inductor must handle arbitrary graphs automatically. A hand-crafted Triton kernel with a large config list will often outperform an Inductor-generated one on the critical path, but Inductor wins on coverage and engineering cost.

Key practical notes on torch.compile autotuning:

  • First-call compilation is slow (10s-200s on large models) because it JIT-compiles and benchmarks.
  • Set TORCHINDUCTOR_MAX_AUTOTUNE=1 to enable the full search; the default is a faster but shallower heuristic pass.
  • Compiled artifacts can be cached with TORCHINDUCTOR_CACHE_DIR to avoid re-tuning on every process restart.

See the torch.compile tutorial for baseline usage.

When it falls down

Autotuning is not a silver bullet. It fails or becomes counterproductive in several well-documented situations:

Shape variability at runtime. If your model encounters thousands of distinct (M, N, K) shapes during inference (e.g., variable-length sequences without padding or bucketing), autotuning incurs both cold-start latency and cache explosion. The solution is shape bucketing: snap inputs to a small set of canonical shapes and pre-tune only those.

Small batch / decode regime. For LLM token generation at batch size 1 or small batches, the arithmetic-intensity bottleneck is so severe that most autotuning configurations collapse to nearly identical throughput. The kernel is memory-bound regardless of tiling, and the search overhead exceeds any gain. Custom decode kernels (e.g., PagedAttention, continuous batching) address this at a higher level.

Register spilling from over-ambitious tiles. A configuration with BLOCK_M=256, BLOCK_N=256 may look attractive because of high reuse, but if it exhausts the register file, threads spill to slow local memory and performance collapses. Triton's profiler and NVIDIA's ncu (Nsight Compute) will show this as high l1tex__t_sectors_pipe_lsu_mem_local_op_ld.sum. The fix is to add register-pressure constraints to the search space, e.g., filtering configs where BLOCK_M * BLOCK_N * dtype_bytes / num_warps / 32 > register_file_per_warp.

Cross-op interactions invalidate single-kernel tuning. Autotuning one kernel in isolation may not account for how it interacts with adjacent ops in a fused graph. A tile size that maximises throughput for a standalone GEMM may thrash L2 cache when preceded by a bandwidth-heavy attention softmax.

Autotuning overhead in CI. If every CI run re-tunes from scratch, compile times balloon. The correct posture is: tune once on representative shapes, commit the cache artifact, and only re-tune on hardware or model changes.

Further reading

Check yourself

7 flashcards for this concept

Click a card to reveal the answer.

Drill the whole track