Performance: Warps, Stages & Autotune

1The three knobs

Every Triton kernel you’ve written so far has three performance dials you never touched:

There is no “correct” value for any of them — it depends on the GPU, the sizes, the dtype. The whole craft is: pick a sane default, then measure. That’s what autotune automates.

2num_warps: people per kitchen

A program is executed by a team of threads. Triton hides the team from you — except for its size. num_warps=4 means 4 warps (128 threads) work on your one tile; num_warps=8 means 256.

Practical rule of thumb: tiles of 16–64 elements per dimension want 4 warps; 64–128 want 8. The compiler will tell you if you’re way off — but it can’t tell you what’s fastest. Measure.

3num_stages: software pipelining

Your matmul K-loop does the same dance every iteration: load the next A/B chunks from HBM into shared memory, compute with them. Loading is slow (memory latency); computing is fast. If the loop waits for each load before computing, it idles half the time.

Software pipelining fixes this the way a chef preps while cooking: while the tensor cores compute with chunk k, the memory system is already fetching chunks k+1 … k+num_stages. The loads are issued ahead, so the compute never waits:

kernel[grid](..., num_stages=3)   # 3 chunks in flight: computing k, prefetching k+1, k+2

num_stages is how deep that prefetch pipeline goes. Deeper = less waiting, but more registers and shared memory held for the in-flight chunks — too deep and you spill, and spilling is HBM traffic again. Typical sweet spot on a T4: 2–4.

4@triton.autotune: let the GPU decide

Since no one can predict the best knob values, don’t: benchmark them all at first launch and keep the winner.

@triton.autotune(
    configs=[
        triton.Config({"BLOCK_M": 16, "BLOCK_N": 16, "BLOCK_K": 16}, num_warps=4, num_stages=2),
        triton.Config({"BLOCK_M": 64, "BLOCK_N": 64, "BLOCK_K": 32}, num_warps=8, num_stages=3),
        triton.Config({"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 64}, num_warps=8, num_stages=4),
        # ... usually 4–8 configs
    ],
    key=["M", "N", "K"],               # re-tune per (M, N, K) — cached after the first call
)
@triton.jit
def matmul_kernel(...):
    ...                                 # unchanged: BLOCK_* come from the winning config

grid = lambda meta: (triton.cdiv(M, meta["BLOCK_M"]), triton.cdiv(N, meta["BLOCK_N"]))
matmul_kernel[grid](a, b, c, M, N, K, ...)   # first call benchmarks, rest use the cache

Three things to notice:

• The grid becomes a lambda: it reads the tile size from the config, because the grid depends on BLOCK_M which autotune is about to choose.
• BLOCK_M/N/K disappear from the launch — the winning config injects them.
• The benchmark happens once per key (per unique M, N, K), then the best config is cached. That first-call cost is why you don’t list 243 configs — each one takes a real timing run.

5do_bench & the launch-bound truth

Timing kernels naively with time.time() lies: the first call includes compilation, the GPU pipelines calls, and the clock is noisy. The right tool:

import triton.testing as tt

fn()                                 # warmup: compile + fill caches
us = tt.do_bench(fn, rep=5) * 1e3    # median of several runs, in µs

do_bench clears the L2 cache between runs (so you measure cold-memory reality), warms up, and reports a stable median. Rule: never trust a number you didn’t get from do_bench.

6Playground

7Homework — seriously, do this

Run lesson6_perf.py on Colab (it’s the first lesson where the local box can’t show you the interesting part), then try each of these. Hints are collapsible.

@triton.autotune(
    configs=[
        triton.Config({"BLOCK_K": 64}, num_warps=4, num_stages=2),
        triton.Config({"BLOCK_K": 128}, num_warps=4, num_stages=3),
        triton.Config({"BLOCK_K": 256}, num_warps=8, num_stages=3),
    ],
    key=["K"],
)
@triton.jit
def gemv_kernel(x_ptr, w_ptr, out_ptr, K, N,
                BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr):
    pid = tl.program_id(0)
    offs_n = pid * BLOCK_N + tl.arange(0, BLOCK_N)
    acc = tl.zeros((BLOCK_N,), dtype=tl.float32)
    for k in range(0, tl.cdiv(K, BLOCK_K)):
        offs_k = k * BLOCK_K + tl.arange(0, BLOCK_K)
        x_v = tl.load(x_ptr + offs_k, mask=offs_k < K, other=0.0)
        w_t = tl.load(w_ptr + offs_n[:, None] * K + offs_k[None, :],
                      mask=offs_k[None, :] < K, other=0.0)
        acc += tl.sum(x_v[None, :] * w_t, axis=1)
    tl.store(out_ptr + offs_n, acc, mask=offs_n < N)

8Your learning journey