The Triton Mental Model

1Why GPUs exist — and why it matters

Your CPU has maybe 8 very smart cores. Your GPU has thousands of dumb-but-fast cores. The entire story of Triton is about knowing which jobs suit which.

“Peeling 10,000 potatoes” is what we call data-parallel work: the same operation applied to a huge amount of data. Adding 1 to every element of a million-element array. Matrix multiplication. Attention. This is ~99% of deep learning’s compute — which is why GPUs train models.

Triton is how you write your own recipe cards for the line cooks.

2The paradigm shift: write one, run many

This is THE idea. Read it twice:

In normal Python you write code that processes one thing at a time:

for i in range(1000):        # the CPU does step i, then step i+1, then...
    out[i] = x[i] + y[i]     # ...one element per loop iteration

In Triton you write a kernel — a function decorated with @triton.jit:

@triton.jit
def add_kernel(x_ptr, y_ptr, output_ptr, n_elements, BLOCK_SIZE: tl.constexpr):
    pid = tl.program_id(axis=0)              # "which copy am I?"
    block_start = pid * BLOCK_SIZE
    offsets = block_start + tl.arange(0, BLOCK_SIZE)
    mask = offsets < n_elements
    x = tl.load(x_ptr + offsets, mask=mask)
    y = tl.load(y_ptr + offsets, mask=mask)
    tl.store(output_ptr + offsets, x + y, mask=mask)

…and then you launch it with a grid:

add_kernel[grid](x, y, out, N, BLOCK_SIZE=256)     # grid = (4,)

The grid is the magic number — it says how many copies of this function to run at once. The GPU photocopies your kernel grid[0] times (4 here) and runs them in parallel.

Every program computes its own pid, so every program knows its own chunk. Same code, different data. That’s the whole game. Play with it below — click a program to see what it owns:

3The second paradigm shift: tile code, not scalar code

Here’s what trips up almost every Triton beginner (me included, back in the day). Look at tl.arange(0, BLOCK_SIZE):

offsets = block_start + tl.arange(0, BLOCK_SIZE)

tl.arange(0, 256) is a vector of 256 indices: [0, 1, 2, ..., 255]. So offsets is also a vector — [768, 769, ..., 1023] for program 3.

That means every operation in your kernel works on a whole block of data at once, not one element:

x = tl.load(x_ptr + offsets)                    # loads 256 elements in ONE operation
tl.store(output_ptr + offsets, x + y, mask=mask)   # adds 256 pairs in ONE operation

And that’s exactly why Triton is fast: the compiler takes your tray-level operations and maps them onto the GPU’s vector hardware and its thousands of threads automatically. You never manage threads. In raw CUDA you’d write threadIdx.x, blockIdx.x, and worry about warps and shared memory — in Triton you just describe tiles of data, and the compiler handles the machinery.

4The anatomy of a kernel: the 5-step universal pattern

Every Triton kernel you’ll ever write follows this skeleton:

   1. "Which program am I?"    →  pid = tl.program_id(axis=0)
   2. "Which data do I own?"   →  offsets = pid * BLOCK + tl.arange(0, BLOCK)
   3. "Am I out of bounds?"    →  mask = offsets < n_elements
   4. "Read my data"           →  x = tl.load(ptr + offsets, mask=mask)
   5. "Write my result"        →  tl.store(ptr + offsets, result, mask=mask)

Step 5 can be a huge computation (a matrix multiply, an attention score) — but the skeleton is always the same.

Line by line, in plain English

5Why the mask? The ragged-edge problem

This is the one thing beginners usually get wrong. Suppose N = 1000 and BLOCK_SIZE = 256:

Element 1000 of a 1000-element array is out of bounds. If program 3 tried to read it, you’d get garbage (or a crash). The mask acts as a safety fence:

mask = offsets < n_elements     # [768..999] → True, [1000..1023] → False

And you pass that same mask to load and store, so out-of-bounds elements are simply skipped. Every real Triton kernel has a mask somewhere — because tile sizes almost never divide the data size evenly.

6tl.constexpr vs runtime arguments

A subtle but critical thing:

def add_kernel(x_ptr, y_ptr, output_ptr, n_elements, BLOCK_SIZE: tl.constexpr):

• n_elements is a runtime value: known when you launch the kernel, can be anything (1000, 77, 12345). Compiles once, runs with different values.
• BLOCK_SIZE: tl.constexpr is a compile-time value: known when the kernel is compiled.

Why does this matter? Because Triton is a metaprogramming language — values marked tl.constexpr can control the program’s structure, like the size of a tl.arange, which must be fixed at compile time. The compiler needs to know the tray size to generate the vector instructions.

The tl. prefix everywhere just means “this function comes from Triton’s language library” — it’s the standard name for all Triton operations, exactly like np. is NumPy.

7Playground: the grid calculator

Change N and BLOCK_SIZE and watch the grid re-shape itself. This is exactly the interaction from homework #1 and #2 — play with it first, then run it on the GPU.

8Homework — seriously, do this

Run lesson1_vector_add.py on Colab, then try each of these by modifying the code. Hints are collapsible — try it first, peek only if stuck.

@triton.jit
def scale_kernel(x_ptr, out_ptr, n_elements, scalar, BLOCK_SIZE: tl.constexpr):
    pid = tl.program_id(axis=0)
    offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    mask = offsets < n_elements
    x = tl.load(x_ptr + offsets, mask=mask)
    tl.store(out_ptr + offsets, x * scalar, mask=mask)

scale_kernel[(triton.cdiv(N, BLOCK_SIZE),)](x, out, N, 3.0, BLOCK_SIZE=256)

9Your learning journey