Triton from zero · Lesson 4 of 7

Block Pointers

Triton’s modern way to describe tiles: one object carries the whole “where is my tile and how do I walk it” story — and masked edges become a one-word flag instead of hand-written masks.

▶ How to run the code (Colab · T4 GPU)

Create a Colab notebook with a GPU runtime (Runtime ▸ Change runtime type ▸ T4 GPU), then in the first cell run:

!nvidia-smi
!pip install -q triton

Paste the code from this lesson (or the whole companion script) into the next cell. No GPU at hand? Run locally — the script falls back to Triton’s CPU interpreter automatically (source /home/tensor/.venv/bin/activate, then cd decode-lab/artifacts and python lesson4_block_ptr_matmul.py). You should see PART A/B/C PASSED at the end.

1The problem with pointer arithmetic

Lesson 2 built its tiles by hand, every time:

a_ptrs = a_ptr + offs_m[:, None] * stride_am + offs_k[None, :] * stride_ak

That works — but it repeats the same choreography in every kernel, it’s easy to swap two strides, and the ragged edges still need hand-built masks. Production Triton (and the official matmul tutorial) uses a higher-level abstraction instead: the block pointer. One object that bundles everything a tile needs to know about the matrix it lives in:

a = tl.make_block_ptr(a_ptr, (M, K), (K, 1), (pid_m * BLOCK_M, 0),
                      (BLOCK_M, BLOCK_K), (1, 0))

The map analogy

Pointer arithmetic is navigating with raw street addresses — correct, but you compute every step by hand. A block pointer is a map with a marker: it knows the whole city (shape), the street grid (strides), where the marker sits (offsets), and the size of the plot you own (block shape). Moving the marker is one call: tl.advance. You never compute an address again.

2`tl.make_block_ptr`: the six arguments

tl.make_block_ptr(base, shape, strides, offsets, block_shape, order)

argument	meaning	example (A, M×K)
`base`	the pointer to the start of the matrix	`a_ptr`
`shape`	the real size of the matrix — the block pointer knows the full extent, so it can police the edges	`(M, K)`
`strides`	elements between rows / columns (lesson 3’s `stride_am`, `stride_ak`, as a tuple)	`(K, 1)`
`offsets`	where the window starts, in elements — usually `pid * BLOCK`	`(pid_m * BLOCK_M, 0)`
`block_shape`	the tile size (your `BLOCK_M × BLOCK_K`)	`(BLOCK_M, BLOCK_K)`
`order`	which axis is contiguous in memory, fastest first. `(1, 0)` = row-major: last axis moves by 1. Tells the compiler how to walk the tile efficiently.	`(1, 0)`

Gotcha — strides vs shape vs order

shape is the logical size of the whole matrix; strides is its physical layout; order is the layout of your tile. Three different things. If you ever declare a block shape that doesn’t fit the shape’s axes, you’ll read garbage (Part C of the companion hits exactly this — deriving the right tuple is the exercise).

3`boundary_check`: masks for free

Remember lesson 1’s ragged-edge problem? With block pointers it stops being your problem:

tile = tl.load(x, boundary_check=(0, 1), padding_option="zero")

The block pointer knows the real shape, so boundary_check=(0, 1) says: “on both axes, if my window hangs off the matrix, pad the out-of-bounds slots with zero (or nan) instead of crashing.” The whole mask= + other= dance of lessons 1–3 collapses into one flag — and the same flag works on tl.store.

Gotcha — padding is zero or nan, that’s all

padding_option only knows "zero" and "nan". For a matmul, zero padding is exact (adding zeros changes nothing), so ragged K is free. But lesson 3’s softmax needed −inf padding — block pointers can’t do that, so softmax still uses hand-built masks. Knowing which tool does what is the real lesson: block pointers for loads/stores of tiles, manual masks when the padding value matters to the math.

4`tl.advance`: sliding windows

Lesson 2’s K-loop slid its windows with pointer arithmetic:

a_ptrs += BLOCK_K * stride_ak

With block pointers, sliding is a method call on the pointer itself — offsets in elements, so “move one chunk along K” is literally “advance by BLOCK_K on axis 0”:

a = tl.advance(a, (0, BLOCK_K))     # slide a's window one K-chunk forward
b = tl.advance(b, (BLOCK_K, 0))     # slide b's window one K-chunk down

Advancing does not change the block pointer’s shape, strides or order — only its offsets. That’s the whole trick of the K-loop: the same two tl.load calls, a sliding window, and the block pointer keeps all the bookkeeping.

5The matmul, rewritten

Compare with lesson 2’s kernel: no offs_m/offs_n/offs_k, no pointer arithmetic, no masks. Same math, ~10 fewer lines, and the ragged edges police themselves:

@triton.jit
def matmul_bp_kernel(a_ptr, b_ptr, c_ptr, M, N, K,
                     N_ITERS: tl.constexpr,
                     BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr):
    pid_m = tl.program_id(0)
    pid_n = tl.program_id(1)

    a = tl.make_block_ptr(a_ptr, (M, K), (K, 1), (pid_m * BLOCK_M, 0),
                          (BLOCK_M, BLOCK_K), (1, 0))
    b = tl.make_block_ptr(b_ptr, (K, N), (N, 1), (0, pid_n * BLOCK_N),
                          (BLOCK_K, BLOCK_N), (1, 0))
    c = tl.make_block_ptr(c_ptr, (M, N), (N, 1),
                          (pid_m * BLOCK_M, pid_n * BLOCK_N),
                          (BLOCK_M, BLOCK_N), (1, 0))

    acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
    for k in range(0, N_ITERS):
        a_tile = tl.load(a, boundary_check=(0, 1), padding_option="zero")
        b_tile = tl.load(b, boundary_check=(0, 1), padding_option="zero")
        acc = tl.dot(a_tile, b_tile, acc)
        a = tl.advance(a, (0, BLOCK_K))              # slide along K
        b = tl.advance(b, (BLOCK_K, 0))

    tl.store(c, acc, boundary_check=(0, 1))

line	what changed vs lesson 2
`a = tl.make_block_ptr(...)`	The whole “where” story in one object: base, real shape `(M, K)`, strides `(K, 1)`, start offset, tile size, row-major order.
`c = tl.make_block_ptr(...)`	The output tile — built once, stored once. No `c_ptrs` arithmetic.
`tl.load(a, boundary_check=(0, 1), padding_option="zero")`	Ragged K, ragged M, ragged N — all handled by one flag. Zero padding is exact for matmul.
`a = tl.advance(a, (0, BLOCK_K))`	Slide the window. The old `a_ptrs += BLOCK_K * stride_ak`, now self-contained.
`tl.store(c, acc, boundary_check=(0, 1))`	Ragged output edges masked automatically — the store equivalent of the load flag.

The launch is identical to lesson 2’s — same grid, same N_ITERS constexpr pattern (the interpreter still can’t do runtime loop bounds):

grid = (triton.cdiv(M, BLOCK_M), triton.cdiv(N, BLOCK_N))
matmul_bp_kernel[grid](
    a, b, c, M, N, K,
    N_ITERS=triton.cdiv(K, BLOCK_K),
    BLOCK_M=16, BLOCK_N=16, BLOCK_K=16,
)

6Block pointer playground

Move the window and watch all six arguments follow. The hatched cells are the boundary-checked slots — loaded as zeros, never read as real data:

7Homework — seriously, do this

Run lesson4_block_ptr_matmul.py on Colab, then try each of these by modifying the code. Hints are collapsible — try it first, peek only if stuck.

1 · Rewrite Part A as a 1D block pointer (vector, not matrix)

A row is a matrix with one axis: shape=(N,), strides=(1,), block_shape=(BLOCK,), order=(0,). Load with boundary_check=(0,). You’ve now replaced lesson 1’s entire mask machinery with one flag.

2 · Break Part C on purpose, then fix it from first principles

This is the exact bug I hit while writing the lesson: I declared B’s block pointer as shape=(N, K) with strides=(1, N) — wrong in two places, and the result was garbage with max_err ≈ 13. Derive it yourself: B^T viewed as a (K, N) matrix has element (k, j) at address j*K + k, so strides are (1, K), shape is (K, N), and the fastest axis is 0. Fix, re-run, watch max_err drop to ~1e-6.

3 · Add boundary_check=(1,) only on the N axis — then only on the M axis

Watch the tests still pass but for subtly different reasons: when only one axis is boundary-checked, the other axis’s out-of-bounds loads read garbage. This is why production kernels often require M, N divisible by their blocks and only check the K edge — checks cost a little performance, so you pay only for the raggedness you actually have.

4 · The real test: softmax with a block pointer + manual mask

Block pointers can’t pad −inf, so load the row-tile with padding_option="nan" is not enough — combine a block pointer load (zero padding) with tl.where(col_mask, x, -inf) for the max pass, exactly as lesson 3 did. You’re now fluent in both idioms and when to use each:

@triton.jit
def softmax_bp_kernel(x_ptr, out_ptr, M, N, BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr):
    pid_m = tl.program_id(0)
    x = tl.make_block_ptr(x_ptr, (M, N), (N, 1), (pid_m * BLOCK_M, 0),
                          (BLOCK_M, BLOCK_N), (1, 0))
    offs_n = tl.arange(0, BLOCK_N)
    col_mask = offs_n[None, :] < N
    t = tl.load(x, boundary_check=(0, 1), padding_option="zero")
    t = tl.where(col_mask, t, -float("inf"))
    m = tl.max(t, axis=1)
    p = tl.exp(t - m[:, None])
    l = tl.sum(p, axis=1)
    out = p / l[:, None]
    tl.store(x, out, boundary_check=(0, 1))

5 · Explain why order=(0, 1) in Part C is not optional

With B^T’s strides (1, K), the contiguous axis is axis 0 — the k-axis. order tells the compiler which axis moves by 1 in memory so it can vectorize the loads along the right direction. Claiming (1, 0) with those strides is a lie the compiler will believe, and your tile will be read in the wrong order (correct values, terrible performance — or a crash).

8Your learning journey

Concept	Status
Grids, tiles, broadcasting, strides (Lessons 1–2)	✔ covered
Reductions & the two-pass dance (Lesson 3)	✔ covered
`tl.make_block_ptr` — the six arguments	✔ covered
`boundary_check` / `padding_option` — masks for free	✔ covered
`tl.advance` — sliding windows in the K-loop	✔ covered
Transpose-read via shape/strides swap	✔ covered
Next: fusing kernels — one pass instead of many (decode-lab’s `swiglu.py`)	⏳ Lesson 5