Chapter 4: Your First Kernel: Vector Addition

The skeleton from the previous topic starts with a puzzle. Every program in the grid runs exactly the same code, so how does each one know which part of the data to process? If two programs computed the same offsets, they would duplicate work and race on the same output; if a gap opened between them, part of the output would never be written. Triton resolves this with a single integer. tl.program_id(axis) returns the program's index along a grid axis, and from it each program derives a contiguous range of global indices: start at pid * BLOCK_SIZE, then add the vector tl.arange(0, BLOCK_SIZE). The result, offs, is itself a tile of indices, and every later load, arithmetic op and store acts on all of them at once. Because consecutive programs own consecutive, non-overlapping ranges, the tiling is disjoint and complete, and neighbouring lanes touch neighbouring addresses, which is what makes the loads coalesce. This topic works through the arithmetic for a 1D grid, where program 10 with 256-element tiles owns indices 2560 to 2815. It then extends the idea to 2D tiles with broadcasting, the addressing pattern matmul uses later.

The previous topic gave each program a tile of offsets, but because the grid rounds up, the last tile almost always extends past the end of the data. For 1,000 elements in 256-element tiles, the fourth program's offsets run to 1023, and 24 of its lanes point at memory the tensor does not own. Reading there returns garbage, and writing there corrupts whatever the allocator placed next. The worst part is that it usually fails silently. In the probe on an RTX 2060, an unmasked fill with n = 1,000 wrote into allocator padding and appeared harmless, while n = 1,100 overwrote the first 128 elements of a neighbouring tensor without any error. tl.load and tl.store are where tiles cross between DRAM and registers, and their mask argument is what makes that crossing safe. A masked-off load substitutes the other value; a masked-off store writes nothing. This topic covers both calls, the mask = offs < n idiom, why other matters for reductions, and the cost of masking. On an exact multiple of the block size, a masked add ran within 0.5 percent of an unmasked one. The next topic turns to the block size itself.

With the kernel correct, the previous topics leave one free parameter: how many elements each program should own. BLOCK_SIZE is a tl.constexpr, so every value compiles a different kernel, and the choice trades opposing costs. Tiny tiles mean millions of programs, each paying scheduling and indexing overhead for very little data. Huge tiles mean each thread must keep more values in registers, which reduces how many programs fit on a streaming multiprocessor and, past the 255-register limit, forces spills to slow local memory. The usual folklore says to stay small to maximise occupancy, but measurements tell a more useful story. On an RTX 2060 adding 50 million floats, BLOCK_SIZE 32 reached only 132 GB/s; every size from 64 to 16,384 landed between 271 and 282 GB/s; and 65,536 spilled 1,326 registers and collapsed to 61 GB/s. The plateau is wide because a memory-bound kernel needs enough bytes in flight, which deep per-thread work supplies as well as extra warps do. This topic explains both cliffs, the register arithmetic behind them, and how num_warps interacts with tile size. Chapter 5 then automates the search with autotuning.

The kernel from the previous topics is only half of the program. Something on the host has to allocate the output, decide how many programs to launch, pass the arguments in the right form and tell you when the work is finished. That last point trips up nearly every first benchmark. A Triton launch is asynchronous: add_kernel[grid](...) places the kernel on a CUDA stream and returns as soon as it is queued. Timing that call with a CPU clock measures the launch, not the GPU work. On an RTX 2060, a 50-million-element add read 48 microseconds without synchronization and 2,157 microseconds with it, a 45 times error. Physics provides a sanity check: the add moves 600 MB, which takes at least 1.79 ms at the card's 336 GB/s peak, so any shorter time is wrong. This topic covers the host side end to end: torch.empty_like for the output, a grid written as lambda meta: (triton.cdiv(n, meta['BLOCK_SIZE']),) so it follows the tile size, tensors passed as pointers, and correct timing with torch.cuda.synchronize or triton.testing.do_bench. Chapter 5 builds on this timing to benchmark and autotune kernels.