Part IV — Back end

Autotuning & cost models

Lesson 09 handed you a code generator that will faithfully emit any schedule you point it at — pick a tile shape, a num_warps, a num_stages, and out comes a working Triton or PTX kernel. But it answered how to emit, not what to emit. The schedule space it can target is combinatorial, and the fastest point in it is different for every shape and every chip. Hand-picking does not scale past a handful of kernels. This lesson is the missing decision procedure: define a search space, then search it — either by benchmarking real candidates on the device, or by predicting their latency with a learned cost model so you can prune the space before you ever launch a kernel.

1 · The space is combinatorial — you cannot enumerate it

A search space is the set of legal schedules the back end could emit for one operator on one target. Start with a single tiled matmul. The schedule knobs from lesson 08, made concrete for a GPU kernel:

Multiply it out. Six choices of BLOCK_M × six of BLOCK_N × four of BLOCK_K × four num_warps × four num_stages is 6 · 6 · 4 · 4 · 4 = 2,304 configurations for one operator on one shape. Add unroll and layout and it is tens of thousands. A real model has dozens of distinct GEMMs, convolutions, and fused reductions, each with several shapes seen in practice. Benchmarking one config means: generate the code, compile it (often hundreds of milliseconds to seconds), launch it many times to get a stable median, and clean up. Call that ≈0.5 s amortized per candidate. Brute-forcing 2,304 candidates is ≈19 minutes for a single kernel/shape pair. For a whole model that is days. And the space is not smooth: latency as a function of tile size has cliffs (a tile one element too big spills registers to local memory and collapses; one too small underutilizes the tensor cores), so you cannot just hill-climb naively from a guess. This is the wall: codegen can emit any point, but the right point is buried in a jagged, high-dimensional space.

2 · Templates + tuning — AutoTVM, the human writes the shape of the search

AutoTVM is the first disciplined answer: split the work between human and machine. A human writes a schedule template for an operator class — the loop structure, where the tiling happens, which axes get split — but leaves the actual values as tunable knobs (TVM calls them cfg.define_split, cfg.define_knob). The template fixes the structure of the search space; the autotuner fills the values.

human template (TVM):              machine fills the knobs:
  schedule conv2d:                   tile_f  ∈ {1,2,4,8,16,32}
    split output by tile_f           tile_y  ∈ {1,2,4,7}
    split spatial by tile_y, tile_x  tile_x  ∈ {1,2,4,7}
    reorder, vectorize, unroll       num_threads, unroll ∈ {...}
    cache_read to shared             →  search picks one assignment

The tuner then explores the value space guided by a cost model trained on the fly: it measures a batch of configs on the device, fits a model (XGBoost over hand-engineered schedule features in classic AutoTVM) to predict latency, uses the model to propose promising next configs (simulated annealing over the predicted surface), measures those, and repeats. The win is that the human supplies the inductive bias — the template rules out illegal or absurd schedules up front — so the machine searches a far smaller, all-legal space. The cost: someone has to write and maintain a template per operator class, and a template that is too rigid hides the global optimum. That maintenance burden is exactly what the next idea removes.

3 · Auto-schedulers + learned cost models — generate the space, predict the latency

Ansor (TVM's auto-scheduler, also called Auto-scheduler) drops the human template entirely. From the operator's mathematical definition it generates the search space itself — it derives a hierarchy of possible loop structures (sketches) and then samples concrete schedules (annotations) within them. This widens coverage dramatically: Ansor routinely finds schedules a human template author never wrote, often beating hand-tuned libraries.

But generating a bigger space makes the measurement problem worse, not better — now there are even more candidates. The lever that makes it tractable is the cost model: a learned function that predicts a schedule's latency from cheap-to-extract features — tile sizes, estimated bytes moved, arithmetic intensity, vectorization width, occupancy proxies — without launching the kernel. A prediction costs microseconds; a real benchmark costs ~0.5 s. So the loop becomes: generate thousands of candidates, score all of them with the cost model in milliseconds, then measure only the top few dozen on the device to correct the model's errors and feed those true latencies back as training data. You measure tens of kernels instead of thousands and still land near the optimum.

The same idea lives inside production compilers under different names. XLA uses an analytical/learned latency-estimating cost model to drive fusion and layout decisions and to rank autotuned GEMM/convolution algorithms, so most of its choices are made by prediction with a thin layer of real autotuning on top for the hot kernels. The frontier — "ML for ML compilers" — is making these cost models good enough that on-device measurement becomes the exception, because measurement is the only slow part of the loop.

4 · Triton's pragmatic autotune — a config list, measured once, cached by key

The most common autotuner an ML engineer actually touches is the simplest one on the table. Triton's @triton.autotune takes you at your word: you hand it a short list of autotune configs (each a tuple of meta-parameters — block sizes, num_warps, num_stages), and on the first call for a given input it benchmarks every config in the list and remembers the winner. There is no generated space and no learned model — just a curated handful of points you believe span the good region, measured for real.

@triton.autotune(
    configs=[
        triton.Config({'BLOCK_M': 128, 'BLOCK_N': 256, 'BLOCK_K': 64}, num_warps=8, num_stages=3),
        triton.Config({'BLOCK_M': 64,  'BLOCK_N': 256, 'BLOCK_K': 32}, num_warps=4, num_stages=4),
        triton.Config({'BLOCK_M': 128, 'BLOCK_N': 64,  'BLOCK_K': 32}, num_warps=4, num_stages=5),
        triton.Config({'BLOCK_M': 64,  'BLOCK_N': 64,  'BLOCK_K': 32}, num_warps=2, num_stages=2),
    ],
    key=['M', 'N', 'K'],          # re-tune only when these change; otherwise reuse the cached winner
)
@triton.jit
def matmul_kernel(A, B, C, M, N, K, BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr):
    ...

The decisive line is key=['M', 'N', 'K']. The cache key is the set of properties that, if unchanged, mean the previously-found winner is still valid: here the problem dimensions (and implicitly the dtypes and device). On the first call for a new (M, N, K) Triton pays the full benchmark sweep — the compile-time cost — and stores the best config; every subsequent call with the same key skips straight to the cached kernel for free. Pick the key too narrow (omit dtype) and you serve a kernel tuned for the wrong precision; pick it too wide (include the exact pointer values) and you re-tune on every call and never amortize. The key is the amortization contract — and, as section 5 shows, it is also the exact place dynamic shapes break.

For the full mechanics of building a config list, sizing the sweep, and wiring autotune into a real kernel pipeline, see triton_kernels · 13 the autotune pipeline; this lesson is the why, that one is the how.

5 · The economics — pay once, amortize a million times

Autotuning only makes sense because of a property unique to ML workloads: the same kernel on the same shapes runs astronomically often. A served model executes its core GEMMs on every token of every request, billions of times over the deployment's life; a training run does millions of identical steps. So the trade is starkly asymmetric. Let search cost T_search seconds once, and let the tuned kernel save Δt per invocation over an untuned default across N runs. Amortization says you win whenever:

Plug in numbers. Suppose tuning a kernel costs T_search ≈ 60 s and the winner is Δt ≈ 50 µs faster per call. Break-even is N = 60 / 50×10⁻⁶ = 1.2 million calls — which a served model crosses in minutes. After that, every call is pure profit. This is why heavy compile-time search is rational here and would be insane for, say, a script you run once.

The break-even math also dictates when you tune:

Either way the winner is written to a cache keyed by shape, dtype, and device, so the cost is paid once per distinct key. And there is the crack. The entire economic argument rests on the same key recurring millions of times. That holds for training (fixed batch, fixed sequence length, fixed shapes every step) and for static-shape inference. It does not hold when the batch size, the sequence length, or a ragged input varies request to request: every new shape is a fresh key, a cache miss, and a cold-start re-tune. A serving system that sees a hundred distinct sequence lengths gets a hundred compile storms and amortizes nothing. The cache key that made autotuning pay is exactly what dynamic shapes shatter — which is the subject of the next lesson.

6 · Drive it: search strategies on a hidden latency surface

The widget below is a 2-D slice of the schedule space: tile_m × tile_n, each a grid of configs, with a hidden "true" latency surface (a smooth basin plus a jagged cliff where tiles spill registers, marked in red). The ★ is the true optimum. Pick a strategy and a measurement budget. Exhaustive measures every cell — it always finds the ★ but spends the whole budget. Random-k samples cells blindly. Cost-model-guided scores every cell with a cheap (noisy) predictor first, then spends its budget measuring only the predicted-best cells — watch it land near the ★ with a fraction of the measurements. The KPIs make the trade explicit: how many real measurements you spent, the best latency you found, the search time, and what percent of the true optimum that is.

The shape of the result is the lesson. Exhaustive is a flat line at 100% of optimum and maximum search time — correct and unaffordable. Random-k's quality rises with budget but jumps around. Cost-model-guided beats random at every budget and nearly matches exhaustive's quality at a tenth of its measurements: search buys speed; the cost model buys cheaper search.

Failure modes & checklist

Checkpoint

Where this points next

You now have the full back end: lesson 08 chose a schedule, lesson 09 emitted it as code, and lesson 10 searched the schedule space — cheaply, via cost models — and cached the winner so the search amortizes over millions of identical runs. The whole pipeline is finally complete for a fixed problem. But notice the load-bearing word in that sentence: fixed. Fusion grouped a graph of known buffer sizes; memory planning packed lifetimes of known extents; layout assignment optimized concrete tensors; and autotuning just keyed its entire cache on the shape. Every middle-end and back-end pass quietly assumed static shapes. Production inference does not oblige: batch size varies request to request, sequence length grows every decode step as the KV cache extends, MoE routing produces ragged tensors, and data-dependent control flow changes the graph itself. Each new shape is a fresh cache key, a cache miss, and a fresh compile — a compile storm that amortizes nothing. The assumption that made the whole stack tractable is the one reality breaks first. That is lesson 11, Dynamic shapes & control flow.

Interview prompts

• Why does autotuning make economic sense for an ML kernel but not for general code? (§5 — the same kernel on the same shapes runs millions of times, so a one-time search cost T_search amortizes whenever N·Δt > T_search; for a served model that break-even is minutes, after which every call is profit. A run-once program never recoups the search.)
• What's the difference between AutoTVM and Ansor? (§2–3 — AutoTVM needs a human-written template that fixes the schedule structure and exposes tunable knobs the machine fills; Ansor generates the search space from the op definition itself, so it needs no template and finds schedules a human wouldn't write, at the cost of longer search.)
• What does a learned cost model buy you, and why still measure on device? (§3 — a prediction costs microseconds vs ~0.5 s for a real benchmark, so the model prunes thousands of candidates to a top-k; you still measure the top-k because the model has error, especially near latency cliffs, and those true latencies retrain it.)
• Explain Triton's @autotune and the role of the key argument. (§4 — it benchmarks a user-supplied config list on first call and caches the winner; key=['M','N','K'] says re-tune only when those change, otherwise reuse the cached kernel. The key is the amortization contract — too narrow serves a wrong kernel, too wide re-tunes every call.)
• AOT vs JIT autotuning — when would you pick each? (§5 — AOT/offline when shapes are known ahead: ship a tuning log so serve-time compile is free; JIT/online when shapes are unpredictable: tune on first encounter, accepting a cold-start spike per new key.)
• Your autotuned serving kernel shows compile time in steady-state latency. What's wrong? (§Failure modes — the cache key includes something that varies per call (e.g. a varying batch/seq dim or a pointer), so it never hits and re-tunes constantly; fix the key to shape+dtype+device, or pre-warm the cache.)
• Why does autotuning's whole premise crack under dynamic shapes? (§5 / Where next — the cache amortizes only when the same key recurs millions of times; variable batch/sequence length makes every shape a new key, a cache miss, and a fresh compile, so you pay search repeatedly and amortize nothing — forcing the dynamic-shape handling of lesson 11.)