Conceptual
Browse Instruct Curate
Login

Writing a First Triton Kernel: Offsets, Masks and a Blocked Vector Add

The canonical first Triton kernel adds two vectors. Each program computes its element indices as pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE), forms a mask offsets < n_elements so the final partial block never reads or writes past the end of the tensor, loads x and y with tl.load(ptr + offsets, mask=mask), adds them, and stores with tl.store under the same mask. The Python wrapper allocates the output, computes the grid as (triton.cdiv(n, BLOCK_SIZE),), and launches with kernel[grid](x, y, out, n, BLOCK_SIZE=1024). Tensors are passed as their data pointers, so strides and dtypes are the programmer's responsibility, and the pattern of pointer plus offsets plus mask is reused unchanged in every larger kernel.