Performance

Measure on the GPU

Time kernels with enceladus.testing.do_bench(fn), which uses GPU timestamps and returns the median in milliseconds. The function must launch on enceladus.Tensor arguments, because a launch that synchronizes can't be timed. The GPU runs slowly for about 50 ms after it idles, so compare kernels in the same process.

Size blocks and SIMD groups

• Memory-bound kernels reach full bandwidth with blocks of 1,024-4,096 elements and 4-8 SIMD groups. On an M4 Pro, measured copy bandwidth is 238 GB/s.
• Keep each tile within about 128 registers per thread. Registers per thread are roughly elements * words_per_element / (32 * num_warps).
• dot_warps=(WM, WN) arranges the SIMD groups as a grid over each tl.dot result. Flash attention runs fastest with dot_warps=(num_warps, 1).

Feed tl.dot from descriptors

When a tl.dot operand is a descriptor load that nothing else uses, tl.dot reads fragments straight from device memory, with no threadgroup memory and no barriers. This path ran 7-15% faster than staging. Epilogues on the accumulator cost little: a fused bias and GELU runs within about 2-3% of a plain matmul.

Choose a tl.dot backend

Set dot_backend as a launch option, a Config field, or an argument to warmup and explain:

A loop is eligible for mpp when the accumulator starts from tl.zeros, is updated only by acc = tl.dot(a, b, acc), both operands are descriptor loads, the output tile is 16-128 in each dimension, and the result reaches a single descriptor store through elementwise ops only. kernel.explain names the reason for any fallback.

Autotune

@enceladus.autotune compiles every candidate in parallel, times each one, and keeps the fastest per key values and argument dtypes. Results persist in ~/.cache/enceladus/autotune/, keyed by kernel source, compiler version, and GPU. The grid must be a function of meta:

from enceladus.configs import matmul_configs


@enceladus.autotune(configs=matmul_configs("float16"), key=["M", "N", "K"])
@enceladus.jit
def matmul_kernel(a_ptr, b_ptr, c_ptr, M, N, K, stride_am, stride_bk, stride_cm,
              BM: tl.constexpr, BN: tl.constexpr, BK: tl.constexpr):
pid_n, pid_m = tl.program_id(0), tl.program_id(1)
a = tl.make_tensor_descriptor(a_ptr, [M, K], [stride_am, 1], [BM, BK])
b = tl.make_tensor_descriptor(b_ptr, [K, N], [stride_bk, 1], [BK, BN])
c = tl.make_tensor_descriptor(c_ptr, [M, N], [stride_cm, 1], [BM, BN])
acc = tl.zeros((BM, BN), dtype=tl.float32)
for k in range(0, K, BK):
    acc = tl.dot(a.load([pid_m * BM, k]), b.load([k, pid_n * BN]), acc)
c.store([pid_m * BM, pid_n * BN], acc.to(c.dtype))


grid = lambda meta: (enceladus.cdiv(N, meta["BN"]), enceladus.cdiv(M, meta["BM"]))
matmul_kernel[grid](a, b, c, M, N, K, K, N, N)
print("chose", matmul_kernel.config_for(a, b, c, M, N, K, K, N, N))

With ENCELADUS_PRINT_AUTOTUNING=1, the output is similar to the following:

enceladus: autotuning matmul_kernel [1024|1024|512|float16,float16,float16] chose enceladus.Config({'BM': 64, 'BN': 32, 'BK': 32}, num_warps=4, dot_warps=(4, 1), dot_backend='mpp') (0.186 ms)

Launch asynchronously

An asynchronous launch costs about 3.4 µs of host time, and batching 64 launches per command buffer brings a small launch to about 1 µs of GPU cost. A synchronized launch costs 70-100 µs. For many small kernels, allocate with enceladus.empty, zeros, or randn, or wrap arrays with enceladus.from_numpy, so that launches never wait.

Other tips:

• Prefer tl.exp2, with log2(e) folded into the scale, over tl.exp.
• Don't use math_mode="fast" for a softmax masked with -inf.
• Call kernel.warmup(*args, **constexprs) to compile ahead of the first launch.