Test a custom kernel#

The problem#

You wrote a kernel: a Triton or CUDA matmul, a Metal shader, a C++ reduction. You need to know that it is correct for many shapes and inputs, not only that it looks close to PyTorch on one random matrix. Comparing with the framework's own result can hide a shared mistake, and comparing with a tolerance picked by eye fails or passes for the wrong reasons.

Why the usual approach falls short#

The framework's output is one correct answer among many, so a different correct answer looks like a difference. A kernel that splits the sum into blocks, as split-K kernels do, adds in another order than PyTorch and can legitimately differ from it. What you need is the set of answers a correct implementation can give.

The solution#

Run your kernel, copy the output to the CPU, and check it against ExactRef for the same inputs. Here a "kernel" written with PyTorch tiles stands in for yours: it splits the inner dimension into blocks of 32 and adds the block products, as a split-K kernel does.

Code
import math
import random
import torch
import exactref as er

# Seeded, so every run uses the same values
rng = random.Random(0)

# A tensor of the given shape, filled with random values
def randn(*shape):
    # One value per element: normal, with mean 0 and standard deviation 1
    values = [rng.gauss(0, 1) for _ in range(math.prod(shape))]
    # torch.tensor makes float32, PyTorch's default for Python floats
    return torch.tensor(values).reshape(shape)

# The kernel under test: A @ B, adding one block's product at a time
def my_matmul(A, B, block=32):
    # Start from an output of zeros
    out = torch.zeros(A.shape[0], B.shape[1])
    for k in range(0, A.shape[1], block):
        # The product of A's columns k to k + block - 1 and the same rows of B
        out = out + A[:, k:k + block] @ B[k:k + block, :]
    return out

# Shapes (m, k, n): A is m x k and B is k x n
for m, k, n in [(1, 1, 1), (3, 5, 7), (64, 256, 32), (17, 300, 9)]:
    A, B = randn(m, k), randn(k, n)
    # ExactRef reads CPU memory, so a kernel run on a GPU copies its output back
    out = my_matmul(A, B).cpu()
    # Could a correct implementation that adds in any order produce out?
    covered = er.matmul(A, B).covers(out, standard="any_order").all()
    print((m, k, n), "covered:", covered)
Output
(1, 1, 1) covered: True
(3, 5, 7) covered: True
(64, 256, 32) covered: True
(17, 300, 9) covered: True

Use standard="any_order" for a kernel that adds in its own order: it accepts every order of additions that does not overflow. When your kernel fixes every step of its order, such as sequential sums inside blocks of a known size, the matching algorithm gives a tighter check and can emulate the kernel's output exactly.

Test more than random inputs: cancellation, a large inner dimension, values near the top of the format's range, subnormals, and infinities and NaN in the inputs. If your kernel accumulates in a narrower format, say so with compute= (Check low-precision kernels).