gpu, week 10retrospective

Triton matmul

Triton matmul expresses tiling as tile-offset arithmetic and tl.dot, letting the compiler own shared-memory staging.


A tile per program

Each program computes one BLOCK_M x BLOCK_N tile of C and walks K:

offs_m = (pid_m * BLOCK_M + tl.arange(0, BLOCK_M)) % M
offs_n = (pid_n * BLOCK_N + tl.arange(0, BLOCK_N)) % N
offs_k = tl.arange(0, BLOCK_K)
a_ptrs = a_ptr + offs_m[:, None] * stride_am + offs_k[None, :] * stride_ak
b_ptrs = b_ptr + offs_k[:, None] * stride_bk + offs_n[None, :] * stride_bn
acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
for k in range(0, tl.cdiv(K, BLOCK_K)):
    a = tl.load(a_ptrs, mask=offs_k[None, :] < K - k * BLOCK_K, other=0.0)
    b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k * BLOCK_K, other=0.0)
    acc = tl.dot(a, b, acc)
    a_ptrs += BLOCK_K * stride_ak
    b_ptrs += BLOCK_K * stride_bk

The tile is a value: acc is a block of registers, tl.dot is one operation over two blocks, and the source contains no shared-memory allocation. The modulo wraps the ragged last tile onto valid rows and the store mask discards the duplicates.

Block pointers, and what replaced them

The older spelling made tiling explicit: tl.make_block_ptr(base, shape, strides, offsets, block_shape, order), tl.advance(ptr, (0, BLOCK_K)), and tl.load(ptr, boundary_check=(0, 1), padding_option="zero"). That API was deprecated in 3.7 and removed in 3.9, where tl.make_block_ptr raises NotImplementedError and boundary_check is gone from load and store. The replacement is tl.make_tensor_descriptor(base, shape=[M, K], strides=[stride_am, 1], block_shape=[BLOCK_M, BLOCK_K]), whose desc.load([off_m, off_k]) provides bounds checking and zero padding. Check hasattr(tl, "make_block_ptr") before copying any block-pointer tutorial.

Precision

tl.dot on f32 inputs defaults to input_precision="tf32" on NVIDIA parts, which truncates mantissas without rounding; allow_tf32 is deprecated in favour of input_precision, and AMD targets default differently. Triton 3.8 also changed the default output dtype when a non-f32 accumulator is passed, so pass out_dtype explicitly and pin the version.

Things that break

Dropping the K mask reads into the next row and produces answers that are wrong without being NaN. A descriptor whose base is not 16-byte aligned, or whose last dimension is not contiguous, is rejected at compile time rather than silently misbehaving. And a kernel still calling tl.make_block_ptr fails before it reaches the GPU, which is the failure you want: the snippets above follow Triton 3.9, and tl.dot is the symbol most likely to differ in yours.

Sources

A version records the documentation the claim was checked against, not a house pin. Confirm flags with --help on the release you installed.

Check yourself

What does `tl.dot` lower to?

Why does the tutorial take `offs_m % M` instead of clamping the last tile?

You call `tl.make_block_ptr` on a current Triton release. What happens?

0 / 3 correct

← back to the map