LAB / GPU PROGRAMMING

DAY 20 / 30 · TRITON & MODERN KERNEL AUTHORING · ~12 MIN READ

Triton's programming model

Thinking in blocks instead of threads; what the compiler now does for you.

01 · Three primitives

A program id, a range, and masked memory

Almost every Triton kernel is built from the same three moves. tl.program_id — which block am I? tl.arange — the offsets this block owns. tl.load/tl.store — move a whole tile, masked at the edges. Here is day 13's vector add again — first the CUDA you wrote, then the same program in Triton:

// CUDA: per-thread thinking
__global__ void add(const float* x,
    const float* y, float* o,
    int n){
    int i = blockIdx.x * blockDim.x
                + threadIdx.x;
    if (i < n) o[i] = x[i] + y[i];
}
# Triton: per-block thinking
@triton.jit
def add(x_ptr, y_ptr, o_ptr, n,
        BLOCK: tl.constexpr):
    pid = tl.program_id(0)
    offs = pid * BLOCK + tl.arange(0, BLOCK)
    mask = offs < n
    x = tl.load(x_ptr + offs, mask=mask)
    y = tl.load(y_ptr + offs, mask=mask)
    tl.store(o_ptr + offs, x + y, mask=mask)

Same program, one abstraction up. No thread index, no guard statement, no launch arithmetic in sight — the grid comes from the launch site: add[(triton.cdiv(n, 1024),)](x, y, o, n, BLOCK=1024), which is day 14's ceil-division, named honestly.

02 · Block-shaped values

Everything is a tensor the size of your tile

The mental shift that makes Triton read naturally: offs, mask, x, y are all block-shaped values — tensors of BLOCK elements living in registers, one slice per hardware thread, sliced by the compiler. Arithmetic is elementwise over the whole block: x + y is BLOCK additions. Comparisons produce block-shaped masks. The compiler's job — the one you did by hand in CUDA — is laying those BLOCK-wide values onto warps and lanes so the memory access coalesces and the math saturates.

Two details of tl.load worth locking in. The other= argument fills masked-out lanes with a value (0 for add, -inf for max reductions — day 21 leans on this). And BLOCK is declared tl.constexpr: it's baked into the compiled kernel as a compile-time constant, which is what makes each block size a distinct, fully-specialized binary — and what day 22's autotuner searches over.

03 · What the compiler does — and doesn't

It owns lowering; you still own the algorithm

What the compiler handles: mapping your block onto warps and lanes, choosing vectorized load widths so accesses coalesce (day 07), inserting shared memory stages and barriers when reuse patterns justify them (day 15's job, automated), allocating registers, pipelining loops. What it deliberately does not do: choose your tiling strategy, your block shape, or your algorithm's data flow. A Triton kernel with a bad blocking scheme compiles beautifully and runs slowly — the compiler is a lowering engine, not an algorithm designer.

Also worth knowing what's absent: there are no warp shuffles or barriers in the language, no shared-memory arrays you can name, and limited control over instruction scheduling. Day 19's trade, restated as a list. When a kernel needs those, it's telling you it belongs in CUDA.

04 · Reading conventions

The shape of every kernel you'll meet from here

Triton code in the wild follows a grammar you can now parse on sight: constexpr block dimensions in the signature (BLOCK_M, BLOCK_N, BLOCK_K for GEMM-shaped things), pointers advanced by block strides, masks constructed from offsets against dimension sizes, results accumulated in block-shaped locals across a loop. When day 24 describes FlashAttention and day 29 builds a fused RMSNorm, both will be exactly this grammar — outer loop over output tiles, inner reuse, masked edges, constexpr knobs.

CHEATSHEET · DAY 20
  • 01program_id, arange, masked load/store — the whole grammar. Every kernel starts the same way: which block am I, what offsets do I own, load the tile, compute, store the tile.
  • 02Values are block-shaped tensors in registers. offs, masks, and data are all BLOCK-wide; elementwise ops act on the whole block. The compiler maps blocks onto warps for you.
  • 03constexpr sizes make each block shape a specialized kernel. tl.constexpr bakes BLOCK into the binary — enabling both compiler optimizations and day 22's autotuning search.
  • 04The compiler lowers; it doesn't design. Coalescing, staging, and pipelining are handled. Your blocking strategy and data flow remain your job — where all the performance lives.

Next: day 21 — writing a fused softmax in Triton: a real, small, complete kernel, start to finish, with the numerical stability trick that makes it correct.

get in touch