Custom elementwise kernel via Triton

The roofline has been my compass for a while. Today I finally wrote a kernel that sits on it, and the chart stopped being abstract.

I wrote my first Triton kernel. Not a fancy one, an elementwise add, but it's mine, and it taught me more about the memory hierarchy than a month of reading.

Why Triton, why elementwise

Writing raw CUDA is a slog: grid, block, thread indexing, shared memory, synchronization, and an hour of debugging for a kernel that does one addition. Triton abstracts the thread-level bookkeeping and lets you write in a Python-like dialect. The compiler figures out the grid, the block, the vectorization, the memory coalescing. You focus on the algorithm.

And elementwise is the perfect first kernel because it's purely memory-bound. It does one FLOP per byte moved. It's the roofline's purest case: no tiling, no shared memory, just "read, compute, write".

The kernel

import triton
import triton.language as tl

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

# launch
grid = (triton.cdiv(n, BLOCK),)
add_kernel[grid](x, y, out, n, BLOCK=1024)

That's it. Four lines of actual logic. The compiler turns it into a CUDA kernel with coalesced memory access, vectorized loads, and no wasted threads.

What I learned

• Block size matters. Too small, and you spend everything on launch overhead. Too big, and you waste threads on the tail. 1024 was the sweet spot in my microbenchmark.
• Masking costs. The mask handles the ragged tail of the array. It's necessary, but it's not free. Aligning your data to the block size removes the mask entirely.
• Bandwidth is the ceiling. My kernel hit about 90% of the theoretical memory bandwidth. That's the roofline, in practice, on my own code. The diagonal line isn't a metaphor.

Why this matters for inference

Real inference workloads are full of elementwise-adjacent ops: activation functions, scaling, quantization, RoPE. When you need a custom op that PyTorch doesn't have, or the fused version isn't fast enough, Triton is how you write it without a CUDA compiler and a week of your life.

And once you can write a kernel, the serving engines stop being magic. When vLLM or SGLang is slow, you can actually look at what it's doing and ask the right question.

The takeaway

The roofline isn't a theory. It's a measurement, and now I've taken it with my own hands.

Tomorrow: a PyTorch custom op with a real CUDA backend, one step deeper.