Skip to content

1. Offsets and masks: which values belong to me?

Time: 45–60 minutes. Prerequisite: lesson 0 and integer division. Edit: offsets.py.

Work a tiny case before writing code

Let N=10 values and BLOCK=4 values per Triton program. You need ceil(10/4)=3 programs, because two programs cover only eight values.

Program ID Local positions Global offsets Valid load/store mask
0 0,1,2,3 0,1,2,3 True,True,True,True
1 0,1,2,3 4,5,6,7 True,True,True,True
2 0,1,2,3 8,9,10,11 True,True,False,False

Two formulas give every entry:

offset = program_id * BLOCK + local_position
valid  = offset < N

Index N is already out of bounds: a vector of length 10 has indices 0 through 9. The mask does not remove an offset from the block. It tells a memory operation which lanes are allowed to access the pointer and which use a substitute value.

Your implementation

First implement ceil_div. Integer division discards the fractional part, so 10 // 4 is only 2. Think about how much to add before // to round upward. Handle N=0 without creating a fictitious program.

Then implement block_offsets, returning both lists. Preserve their length: the last program still has four logical positions even though only two are valid.

python -m learning.explain indexing
python -m learning.check offsets
python -m learning.check offsets --hint

Try N=1, N=4, N=5 and N=0 on paper. Change BLOCK to 8. For N=10, the number of programs changes from 3 to 2; the number of useful output elements stays 10.

How it appears in the real kernel

Read _add in triton_kernels.py:

offset = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)
x = tl.load(X + offset, offset < N, other=0)
tl.store(OUT + offset, x, offset < N)

X + offset is pointer arithmetic in elements. Triton knows the pointer's element type. You do not multiply by four yourself for FP32. The load substitutes zero for masked lanes; the masked store writes nothing for those lanes.

Common bugs and their symptoms

Mistake What happens Revealing input
N // BLOCK grid size Tail outputs never get written N=257, BLOCK=256
Use <= N in mask One invalid element is accessed N=257, BLOCK=256 (with N=256 no lane has offset N)
Forget program_id * BLOCK All programs write the first block N=1024
Mask the store but not the load Reads can still go out of bounds N=257

Exit check: draw the offsets and masks for N=17, BLOCK=8. Explain which program owns element 16 and why each valid output has exactly one owner. Then continue to strides.