2D Tiles & Matrix Multiply

1Why 2D? The grid becomes a grid

Last lesson your grid was a line: grid = (4,), and tl.program_id(axis=0) told each program its position in that line. But a matrix is a rectangle — so this lesson the launch grid becomes a rectangle too:

grid = (triton.cdiv(M, BLOCK_M), triton.cdiv(N, BLOCK_N))   # (row-blocks, col-blocks)

pid_m = tl.program_id(axis=0)     # which ROW-block of C am I?
pid_n = tl.program_id(axis=1)     # which COL-block of C am I?

Two numbers, two axes. axis=0 indexes the first grid dimension (rows), axis=1 the second (columns). Each program now owns a 2D tile of the output — a BLOCK_M × BLOCK_N rectangle — instead of a 1D slice.

The only real difference from lesson 1: instead of one index vector, you build two — one for the rows you own, one for the columns:

offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)   # (BLOCK_M,)  row indices of my tile
offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)   # (BLOCK_N,)  col indices of my tile

This is the whole Part A kernel from lesson2_matmul.py — it fills every cell of C with its own linear index i*N + j, so you can literally see which program owned which cell:

@triton.jit
def tile_fill_kernel(out_ptr, M, N, BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr):
    pid_m = tl.program_id(axis=0)                     # which row-block am I?
    pid_n = tl.program_id(axis=1)                     # which col-block am I?

    offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)  # (BLOCK_M,) row indices
    offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)  # (BLOCK_N,) col indices

    # The star of the lesson, in one line:
    #   column vector (BLOCK_M, 1)  +  row vector (1, BLOCK_N)
    #   -> broadcast into a (BLOCK_M, BLOCK_N) tile of linear indices
    linear = offs_m[:, None] * N + offs_n[None, :]

    mask = (offs_m[:, None] < M) & (offs_n[None, :] < N)
    tl.store(out_ptr + linear, linear.to(tl.float32), mask=mask)

2Broadcasting: build a tile in one line

Here is the trick that unlocks everything in Triton. We have a column vector and a row vector, and we want a whole 2D tile:

column = tl.arange(0, BLOCK_M)[:, None]     # shape (BLOCK_M, 1)   — a column
row    = tl.arange(0, BLOCK_N)[None, :]     # shape (1, BLOCK_N)   — a row
tile   = column + row                       # shape (BLOCK_M, BLOCK_N) — the whole 2D tile!

[:, None] adds a new axis at the end (turns a vector into a column); [None, :] adds a new axis at the front (turns a vector into a row). When you add a (BLOCK_M, 1) and a (1, BLOCK_N) tensor, Triton broadcasts: the column is repeated across all N positions, the row is repeated across all M positions, and every cell gets i + j — its row index plus its column index.

This is exactly what the matmul kernel will use, in a slightly disguised form: offs_m[:, None] * stride + offs_k[None, :] — a column of row-indices plus a row of column-indices = a 2D tile of pointer offsets. Learn it here, on plain numbers:

3Strides: how matrices actually live in memory

GPU memory is a long 1D strip of addresses — there are no 2D boxes in it. A 2D matrix is laid out row by row (this is called row-major order), and the address of element (i, j) is:

address(i, j) = base + i * stride_0 + j * stride_1

stride_0 = distance in memory between row i and row i+1   (= K for an M×K matrix)
stride_1 = distance between column j and column j+1       (= 1, always adjacent)

For an M×K matrix, stride_0 = K and stride_1 = 1: row 0 starts at element 0, row 1 at element K, row 2 at element 2K…

Why the kernel takes strides, not shapes

The kernel signature from lesson2_matmul.py is deliberately verbose — that verbosity is the point:

def matmul_kernel(a_ptr, b_ptr, c_ptr,
                  M, N, K,                          # real sizes
                  stride_am, stride_ak,             # how A steps through memory
                  stride_bk, stride_bn,             # how B steps through memory
                  stride_cm, stride_cn, ...):       # how C steps through memory

The kernel never needs to know that A is an M×K matrix — it needs to know how to walk it. Moving from offs_m to offs_m + 1 (one row down) means jumping stride_am elements; moving along offs_k (one column right) means jumping stride_ak elements. That’s why the pointer tiles are built the way they are:

# a tile of pointers: my BLOCK_M rows × BLOCK_K columns, positioned by strides
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

Look at a_ptrs: a column of row-indices (offs_m[:, None]) times the row stride, plus a row of column-indices (offs_k[None, :]) times the column stride — the §2 broadcasting trick, in its real job. One line, and you have a (BLOCK_M, BLOCK_K) tile of addresses.

4The matmul kernel, line by line

Now the whole thing. This is matmul_kernel from lesson2_matmul.py — the classic tiled GEMM in ~25 lines. Read it once through, then read the anatomy table below it:

@triton.jit
def matmul_kernel(a_ptr, b_ptr, c_ptr,
                  M, N, K,                          # real sizes
                  stride_am, stride_ak,             # how A steps through memory
                  stride_bk, stride_bn,             # how B steps through memory
                  stride_cm, stride_cn,             # how C steps through memory
                  N_ITERS: tl.constexpr,            # how many K-chunks (see note)
                  BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr):
    pid_m = tl.program_id(axis=0)                   # row-block / col-block
    pid_n = tl.program_id(axis=1)

    # rows/cols of C owned by this program, plus the K-chunk index vector
    offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)   # (BLOCK_M,)
    offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)   # (BLOCK_N,)
    offs_k = tl.arange(0, BLOCK_K)                     # (BLOCK_K,)

    # pointer tiles we slide along K in the loop below
    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, N_ITERS):
        # ragged K edge? load 0.0 for the out-of-range columns so the
        # partial sums stay exactly right
        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)                      # acc += a @ b  (tensor cores)
        a_ptrs += BLOCK_K * stride_ak                # slide both windows to the
        b_ptrs += BLOCK_K * stride_bk                # next K-chunk

    c_ptrs = c_ptr + offs_m[:, None] * stride_cm + offs_n[None, :] * stride_cn
    mask = (offs_m[:, None] < M) & (offs_n[None, :] < N)
    tl.store(c_ptrs, acc, mask=mask)

And the launch, for reference — the grid is 2D now, and the strides are passed straight from the torch tensors:

grid = (triton.cdiv(M, BLOCK_M), triton.cdiv(N, BLOCK_N))
matmul_kernel[grid](
    a, b, c, M, N, K,
    a.stride(0), a.stride(1),                    # row stride = K, col stride = 1
    b.stride(0), b.stride(1),
    c.stride(0), c.stride(1),
    N_ITERS=triton.cdiv(K, BLOCK_K),
    BLOCK_M=16, BLOCK_N=16, BLOCK_K=16,
)

5Why tiling is fast: three reasons

The naive way to compute C = A @ B would have each of the M·N output elements walk a full row of A and a full column of B — every element of A is read N times, every element of B is read M times. That’s 2·M·N·K memory reads, and memory — not math — is the bottleneck on a GPU.

Tiling fixes this in three compounding ways:

One more thing worth saying out loud: acc = tl.dot(a, b, acc) accumulates in place — the acc argument is the running sum. That is exactly the acc += a @ b of the math, and it means the K-loop can run as many chunks as needed, even when K is huge, without ever holding more than one chunk of A and B at a time.

6The K-loop, frame by frame

Everything above collapses into one tiny picture. Here is a full matmul with BLOCK_M = BLOCK_N = 2, BLOCK_K = 1 — one program, one K-chunk per step. Watch what gets loaded (the hot column of A and the hot row of B), what gets added to the accumulator, and how acc evolves into A @ B:

7Homework — seriously, do this

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

@triton.jit
def transpose_kernel(in_ptr, out_ptr, M, N,
                     stride_im, stride_in, stride_om, stride_on,
                     BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr):
    pid_m = tl.program_id(axis=0)
    pid_n = tl.program_id(axis=1)
    offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
    offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
    # read the (n, m) element by swapping which index multiplies which stride
    in_ptrs = in_ptr + offs_n[:, None] * stride_im + offs_m[None, :] * stride_in
    mask = (offs_n[:, None] < N) & (offs_m[None, :] < M)
    val = tl.load(in_ptrs, mask=mask)
    out_ptrs = out_ptr + offs_m[:, None] * stride_om + offs_n[None, :] * stride_on
    tl.store(out_ptrs, val, mask=mask)

8Your learning journey