Common operations, declarations, call flows

OpenAI Triton, the kernel DSL — not the inference server. We build the mental model from first principles: a Triton kernel is a program that owns a tile, and the compiler — not you — decides how threads, warps, vectorization, and shared memory realize that tile on the GPU.

1. The tile programming model from first principles

In CUDA you write the code for one thread. The thread owns one scalar (or a small handful), you compute its global index from blockIdx, blockDim, and threadIdx, and the hardware runs tens of thousands of copies of that scalar program. You — the author — are responsible for the thread-to-data mapping, for coalescing, for staging data through shared memory, and for the warp-level choreography.

In Triton you write the code for one program. A program owns a whole tile: a 1-D vector of BLOCK elements, or a 2-D block of BLOCK_M × BLOCK_N, etc. You express the math over the entire tile at once — x + y, tl.sum(x, axis=1), tl.dot(a, b) — as if it were NumPy. The Triton compiler then decides:

• How many threads/warps execute the tile and how the tile's lanes map onto them.
• How loads vectorize and whether accesses coalesce.
• When to stage operands through shared memory or registers.
• How to software-pipeline loads against compute.

2. A minimal complete kernel: vector add

The "hello world" of Triton. One program owns one BLOCK-sized slice of the vectors; the grid covers all slices.

import torch
import triton
import triton.language as tl

@triton.jit
def add_kernel(
    x_ptr,            # *fp32, first input
    y_ptr,            # *fp32, second input
    out_ptr,          # *fp32, output
    n_elements,       # int, total length
    BLOCK: tl.constexpr,   # compile-time tile size
):
    # 1. tile owner: which BLOCK-slice is mine?
    pid = tl.program_id(axis=0)

    # 2. offsets: the element indices this program touches
    offs = pid * BLOCK + tl.arange(0, BLOCK)

    # 3. mask: the last program may run past the end
    mask = offs < n_elements

    # 4. pointers + masked loads (pointers are tensors of addresses)
    x = tl.load(x_ptr + offs, mask=mask, other=0.0)
    y = tl.load(y_ptr + offs, mask=mask, other=0.0)

    # 5. compute over the whole tile at once
    out = x + y

    # 6. masked store
    tl.store(out_ptr + offs, out, mask=mask)


def add(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
    assert x.is_cuda and y.is_cuda and x.shape == y.shape
    out = torch.empty_like(x)              # allocate output
    n = x.numel()

    # grid: a callable taking the meta-params, returns # of program instances
    grid = lambda meta: (triton.cdiv(n, meta["BLOCK"]),)

    # launch: kernel[grid](args..., constexpr kwargs, num_warps=...)
    add_kernel[grid](x, y, out, n, BLOCK=1024, num_warps=4)
    return out

Notes that matter at interview time:

• triton.cdiv(n, BLOCK) is ceiling division — it guarantees the grid covers the tail, which is exactly why the mask exists.
• BLOCK is passed as a keyword because it is a tl.constexpr compile-time constant (see §4).
• num_warps is a launch-time hint about execution resources per program — not a tile dimension (see §11).
• We never wrote a thread index. The compiler maps the 1024-lane tile onto the 4 warps.

3. The call flow

What actually happens between add(x, y) and the result tensor:

4. constexpr and JIT specialization

An argument annotated : tl.constexpr is a compile-time constant. The compiler bakes its value into the generated code — so it can unroll loops over it, size tl.arange, and pick vector widths. The cost: each distinct value of a constexpr (and each distinct combination of argument dtypes / pointer-alignment classes) is a separate specialization that must be compiled.

# BLOCK=1024 and BLOCK=2048 compile to two different kernels.
add_kernel[grid](x, y, out, n, BLOCK=1024)   # specialization A
add_kernel[grid](x, y, out, n, BLOCK=2048)   # specialization B (separate compile)

# num_warps / num_stages also key the specialization:
add_kernel[grid](x, y, out, n, BLOCK=1024, num_warps=8)   # yet another variant

5. Vocabulary you must be able to define

6. Pointer arithmetic with strides

Triton does not assume your tensors are contiguous. You pass the strides explicitly and compute addresses yourself. This is what lets one kernel work on transposed views, slices, and non-contiguous layouts without a copy.

Consider scaling each row of a matrix X (M × N) by a per-row scalar. One program owns one row.

@triton.jit
def row_scale_kernel(
    x_ptr, scale_ptr, out_ptr,
    M, N,
    stride_xm, stride_xn,      # element strides of X
    stride_om, stride_on,      # element strides of OUT
    BLOCK_N: tl.constexpr,
):
    row = tl.program_id(axis=0)            # this program owns row `row`
    if row >= M:
        return

    s = tl.load(scale_ptr + row)           # scalar for this row

    # sweep the columns in BLOCK_N-wide chunks
    for start in range(0, N, BLOCK_N):
        offs_n = start + tl.arange(0, BLOCK_N)
        mask = offs_n < N

        # address = base + row*row_stride + col*col_stride
        x_ptrs   = x_ptr   + row * stride_xm + offs_n * stride_xn
        out_ptrs = out_ptr + row * stride_om + offs_n * stride_on

        x = tl.load(x_ptrs, mask=mask, other=0.0)
        tl.store(out_ptrs, x * s, mask=mask)

7. 2-D pointer blocks via broadcasting

To own a 2-D tile, build two 1-D offset vectors and broadcast them into a grid of addresses with None (newaxis) indexing — exactly like NumPy.

@triton.jit
def copy_block_kernel(
    x_ptr, out_ptr,
    M, N,
    stride_m, stride_n,
    BLOCK_M: tl.constexpr,
    BLOCK_N: tl.constexpr,
):
    pid_m = tl.program_id(axis=0)
    pid_n = tl.program_id(axis=1)

    offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)   # (BLOCK_M,)
    offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)   # (BLOCK_N,)

    # broadcast to a 2-D tile of addresses: (BLOCK_M, BLOCK_N)
    ptrs = x_ptr + offs_m[:, None] * stride_m + offs_n[None, :] * stride_n

    # 2-D mask: combine two boundary checks with bitwise &
    mask = (offs_m[:, None] < M) & (offs_n[None, :] < N)

    block = tl.load(ptrs, mask=mask, other=0.0)

    out_ptrs = out_ptr + offs_m[:, None] * stride_m + offs_n[None, :] * stride_n
    tl.store(out_ptrs, block, mask=mask)

8. Masks and the other value

Any lane that is out of bounds, on the wrong side of a causal frontier, or otherwise invalid must be masked — both on load (so you do not read garbage / out-of-bounds memory) and on store (so you do not clobber memory you do not own). On load, other= supplies the value substituted for masked-off lanes.

The fill value must be the identity of the downstream reduction, or it will corrupt the result:

9. Reductions inside a tile

Reductions over a tile use tl.sum, tl.max, tl.min (and friends) with an axis argument, mirroring NumPy. For a 2-D tile shaped (BLOCK_M, BLOCK_N):

• axis=0 reduces down the rows → a (BLOCK_N,) vector (one value per column).
• axis=1 reduces across the columns → a (BLOCK_M,) vector (one value per row).

@triton.jit
def row_sum_kernel(x_ptr, out_ptr, M, N,
                   stride_m, stride_n,
                   BLOCK_N: tl.constexpr):
    row = tl.program_id(axis=0)
    offs_n = tl.arange(0, BLOCK_N)
    mask = offs_n < N
    x = tl.load(x_ptr + row * stride_m + offs_n * stride_n,
               mask=mask, other=0.0)            # 0.0 is the identity for sum
    acc = tl.sum(x.to(tl.float32), axis=0)      # reduce in fp32
    tl.store(out_ptr + row, acc)

10. dtype handling

Loads from a bf16/fp16 tensor return bf16/fp16 tiles. Cast to fp32 before any reduction or accumulation, then cast back to the output dtype on store.

x = tl.load(x_ptr + offs, mask=mask, other=0.0)   # bf16 tile
acc = tl.sum(x.to(tl.float32), axis=0)            # accumulate in fp32
tl.store(out_ptr + row, acc.to(x.dtype))          # store back as bf16

11. Choosing num_warps and num_stages

• num_warps — execution resources per program instance. More warps give the scheduler more in-flight work to hide memory latency, which helps for large tiles and memory-bound kernels. Too many warps for a small tile wastes registers and hurts occupancy. Typical: 4 or 8.
• num_stages — software-pipeline depth. The compiler prefetches the next stage's loads while the current stage computes, overlapping memory and math. It matters for loops with heavy compute per iteration (e.g. tl.dot in a matmul or attention inner loop). For tiny elementwise kernels there is nothing to overlap, so it is irrelevant — leave it at the default.

12. PyTorch autograd integration

To make a Triton kernel differentiable, wrap forward and backward launches in a torch.autograd.Function. Save what backward needs in ctx; return one gradient per forward input.

class FusedAdd(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x, y):
        out = torch.empty_like(x)
        n = x.numel()
        grid = lambda meta: (triton.cdiv(n, meta["BLOCK"]),)
        add_kernel[grid](x, y, out, n, BLOCK=1024)
        # nothing input-dependent needed for d(x+y); ctx would hold tensors otherwise
        return out

    @staticmethod
    def backward(ctx, grad_out):
        # d(x+y)/dx = 1, d(x+y)/dy = 1  ->  both grads pass through unchanged
        return grad_out, grad_out

# usage
z = FusedAdd.apply(x, y)

backward returns a gradient for each positional input to forward (here x and y), in order; return None for non-tensor or non-differentiable inputs.

Interview checklist

1. State what one program owns. A 1-D slice? A row? A 2-D BLOCK_M × BLOCK_N block? Everything else follows from this.
2. Derive the pointer tensors from tl.program_id and tl.arange, multiplying offsets by the passed-in strides (do not assume contiguity); broadcast with [:, None] / [None, :] for 2-D.
3. Mask every boundary on both load and store, and pick other to be the identity of the downstream reduction (0 for sum, −inf for max).
4. Cast to fp32 for any reduction/accumulation; store back in the output dtype.
5. Choose BLOCK / num_warps / num_stages from tile size and arithmetic intensity (and autotune).
6. Say whether it reduces within one program or needs a multi-pass launch — there is no global barrier inside a kernel.