First Triton vector add

The kernel

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

grid = lambda meta: (triton.cdiv(n, meta['BLOCK']),)
add_kernel[grid](x, y, out, n, BLOCK=1024)

Read the three numbers

BLOCK is a tl.constexpr, so it is baked into the compiled kernel as a shape: 1024 elements per program. The grid is cdiv(n, BLOCK) programs, the ceiling division that keeps the tail. num_warps is absent from the source and defaults to 4, so 128 threads cover 1024 lanes at eight elements each, which is what lets the loads vectorise. Changing BLOCK or num_warps recompiles the kernel.

Measure in GB/s

The kernel moves 12 * n bytes for n flops. Time it and divide:

ms = triton.testing.do_bench(lambda: add_kernel[grid](x, y, out, n, BLOCK=1024))
gbps = 3 * n * 4 * 1e-9 / (ms * 1e-3)

Three streams, not two: two reads and one write. do_bench returns milliseconds, so the unit conversion is part of the formula. Expect the result to sit just under the copy bandwidth from the roofline node, because this kernel has no reuse to exploit. A large gap is a question about masks or vectorisation before it is a question about the GPU.

Things that break

Allocating the output inside the timed lambda measures the caching allocator, not the kernel. A mask that is always true because n happens to divide BLOCK hides the tail bug until the first odd size, which is the size the test should use. And a kernel launched with a grid built from a stale BLOCK reads the wrong region entirely, which shows up as a mismatch in the first few elements rather than as a crash. The snippets follow Triton 3.9; check triton.cdiv and the meta['BLOCK'] grid convention against your installed version first.