Triton: programming the GPU
in tiles, not threads

A modern NVIDIA GPU has tens of thousands of threads running at once. CUDA — the standard way to program them — asks you to write code from the point of view of one of those threads. Every thread reads its own ID, figures out which slice of memory it owns, and you, the programmer, are responsible for making the choices that decide whether the kernel runs at 5% or 95% of peak throughput: how to group threads into warps, which loads will coalesce into a single memory transaction, when to stage data in the on-chip scratchpad (__shared__), how to avoid bank conflicts.

None of that work is about your algorithm. It's bookkeeping that exists because the model is per-thread. Drag the slider below to grow the vector and watch the per-thread bookkeeping multiply — every cell is a thread that has to compute its own index.

Triton flips the unit of programming. Where CUDA gives each thread a scalar ID and asks it to fetch one element, Triton gives each program (a block of threads) a single ID and asks it to operate on a whole array — a tile — at once.

Inside a Triton kernel, when you write x + y, those aren't two numbers — they're two tiles, say of shape (BLOCK_SIZE,). Loads, stores, reductions, even matrix multiplications operate on tiles. The compiler is the one that decides how those tile operations get mapped down to individual threads, warp shuffles, and shared memory. You write the math; it writes the choreography.

Two ingredients make this work. First, every program asks the runtime which tile it owns:

pid = tl.program_id(0)              # which tile am I?
offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)

offsets is a vector of BLOCK_SIZE indices — the row of cells this program is responsible for. Second, because the last tile may run off the end of the data, a boolean mask is passed to every load and store:

mask = offsets < N
x = tl.load(x_ptr + offsets, mask=mask)

The masked load reads valid entries and pretends the rest are zero. That's enough to handle any N. Play with the widget below — change N and BLOCK_SIZE and watch the tile partition reshape. The faint cells in the last tile are the ones the mask filters out.

A Triton kernel is launched by handing it a grid — the number of programs to run, possibly across multiple axes. Each program reads its coordinate via tl.program_id(axis). The runtime schedules them on the GPU's streaming multiprocessors in whatever order it pleases, in parallel, with no ordering guarantees.

For a 1D problem like vector add, the grid is one-dimensional: grid = (cdiv(N, BLOCK_SIZE),). For a 2D problem like matmul, the grid is two-dimensional — (cdiv(M, BLOCK_M), cdiv(N, BLOCK_N)) — and each program produces one tile of the output. The visualization below shows programs firing concurrently and finishing on their own schedule.

Everything in the steps above fits in twenty lines of code. Here is the complete kernel for z = x + y. Step through it with the widget — each press highlights one line of the kernel and shows the memory region it touches.

@triton.jit
def add_kernel(x_ptr, y_ptr, z_ptr, N, BLOCK_SIZE: tl.constexpr):
    pid     = tl.program_id(0)
    offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    mask    = offsets < N
    x       = tl.load(x_ptr + offsets, mask=mask)
    y       = tl.load(y_ptr + offsets, mask=mask)
    z       = x + y
    tl.store(z_ptr + offsets, z, mask=mask)

# Launching it:
N = 100_000
grid = lambda meta: (triton.cdiv(N, meta['BLOCK_SIZE']),)
add_kernel[grid](x, y, z, N, BLOCK_SIZE=1024)

The kernel never mentions threads, warps, shared memory, or coalescing. It says: this program owns tile pid; load its slice from x and y, add them, store the result. The compiler does the rest.

Softmax over a row is four operations: subtract the max for numerical stability, exponentiate, sum, divide. In PyTorch each of these can run as a separate kernel, and the intermediate results travel through the GPU's main memory (HBM) between kernels. HBM bandwidth is the bottleneck of nearly every deep-learning op — a round-trip costs orders of magnitude more than the arithmetic itself.

In Triton you write the whole softmax as one kernel. The intermediate tile lives in registers / on-chip SRAM the entire time. The row is loaded once, the result is stored once, and the four math operations happen in between for free.

# Input X has shape (n_rows, N).  Launch with grid = (n_rows,) — one program per row.
# BLOCK is a compile-time power-of-two ≥ N, so a whole row fits in one tile.
@triton.jit
def softmax_row(x_ptr, y_ptr, stride, N, BLOCK: tl.constexpr):
    row     = tl.program_id(0)                       # scalar : which row this program owns
    cols    = tl.arange(0, BLOCK)                     # tile   : (BLOCK,)  = [0, 1, …, BLOCK-1]
    mask    = cols < N                              # tile   : (BLOCK,)  bool — True where col is in range
    x       = tl.load(x_ptr + row*stride + cols,        # tile   : (BLOCK,)  one full row in SRAM
                       mask=mask, other=-float('inf'))   #          out-of-range lanes → -∞ so max ignores them
    x       = x - tl.max(x, axis=0)                  # tile   : (BLOCK,)  — broadcast subtract of a scalar max
    num     = tl.exp(x)                              # tile   : (BLOCK,)  numerator of softmax
    y       = num / tl.sum(num, axis=0)                # tile   : (BLOCK,)  — divide by scalar sum
    tl.store(y_ptr + row*stride + cols, y, mask=mask) # write the row back to HBM (mask drops the tail)

Before stepping through the code, one piece of background that the walkthrough leans on. A matrix in memory is not actually a two-dimensional thing — HBM is a single flat array of bytes. The 2D view X[row, col] is a fiction we maintain by storing the rows back-to-back in one long vector, then computing the address as row · stride + col, where stride is the number of elements between the start of consecutive rows (usually just the row width, in elements). This is called row-major layout and it's the standard in PyTorch, NumPy, and Triton. It matters here for two reasons. First, it's why the kernel can compute addresses with a single multiply-add — no nested index, just one pointer. Second, the rows being contiguous means lanes 0..BLOCK-1 of one program touch BLOCK consecutive bytes, which is what lets the GPU coalesce them into one wide memory transaction. The flat layout and the load throughput are the same fact.

Walk through the kernel below. The widget after that races the timelines to make the bandwidth win concrete.

Matmul is the canonical GPU kernel — the one where every bandwidth-vs-compute decision shows up. Before the code, the symbols:

• A, B, C — pointers to the three matrices in HBM. We are computing C = A @ B.
• M, N, K — the matrix shapes: A is M × K, B is K × N, so C is M × N. The recipe is \( C_{ij} = \sum_{k} A_{ik}\,B_{kj} \).
• K is the contraction axis — the dimension shared by both inputs that gets summed away. M and N are the outer dimensions that survive into the output.
• BM, BN, BK — compile-time tile sizes. Each program writes one BM × BN tile of C and scans across K in chunks of BK, accumulating BK-wide outer products into a register-resident accumulator.
• stride_am, stride_ak, etc. — bytes (or elements) to step in memory to advance one row vs one column of each matrix. Same role as row*stride in the softmax kernel, generalized to 2D.

The Triton recipe is the same idea as before, in two dimensions: launch a 2D grid where each program owns one BM × BN tile of C, and inside the program, march across K a tile at a time, accumulating partial products into a register-resident accumulator.

@triton.jit
def matmul_kernel(A, B, C, M, N, K,
                  stride_am, stride_ak, stride_bk, stride_bn, stride_cm, stride_cn,
                  BM: tl.constexpr, BN: tl.constexpr, BK: tl.constexpr):
    pid_m = tl.program_id(0)
    pid_n = tl.program_id(1)

    offs_m = pid_m*BM + tl.arange(0, BM)
    offs_n = pid_n*BN + tl.arange(0, BN)
    offs_k = tl.arange(0, BK)

    acc = tl.zeros((BM, BN), tl.float32)
    for k in range(0, K, BK):
        a = tl.load(A + offs_m[:, None]*stride_am + (k+offs_k)[None, :]*stride_ak)
        b = tl.load(B + (k+offs_k)[:, None]*stride_bk + offs_n[None, :]*stride_bn)
        acc += tl.dot(a, b)            # this lowers to tensor cores

    tl.store(C + offs_m[:, None]*stride_cm + offs_n[None, :]*stride_cn, acc)

The two parts to internalize. (1) Each program writes one tile of C; the grid covers the output. (2) The K-loop streams narrow tiles of A and B through the program, feeding the accumulator. tl.dot lowers to NVIDIA's tensor cores — the same instructions cuBLAS uses. Drag the sliders and watch the program that owns the highlighted output tile do its work.