2 / 11 Kernels (CUDA / Triton / Pallas)

What pid Indexes in a Triton Kernel

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

grid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE']), )

A CUDA developer reads this and asks where the per-thread code went. Which reading is correct?

A

mask guards against numerical overflow in the addition, and the ragged tail is handled by a separate cleanup kernel Triton emits automatically.

B

Because Triton hides thread indexing, it cannot express shared-memory tiling, so kernels like FlashAttention have to be written in CUDA or Pallas instead.

C

pid identifies one program instance, and each instance owns a whole BLOCK_SIZE-element tile: offsets is a vector, not a scalar. The compiler assigns threads within the tile, coalesces the loads, and manages shared memory, which in CUDA you would write by hand.

D

pid is the thread index, the direct analogue of blockIdx.x * blockDim.x + threadIdx.x, so each program instance handles exactly one element.

E

BLOCK_SIZE: tl.constexpr marks it as a runtime argument, so changing it reuses the same compiled binary with a different tile size.