Quickstart

Install Enceladus, and then run a vector add, a row softmax, a matrix multiplication, and flash attention on your Mac's GPU.

Install Enceladus

To use Enceladus, you need the following:

A Mac with Apple silicon (M1 or later).
macOS 15 or later.
Python 3.11 or later.
uv, the Python package and project manager.

To add Enceladus to your uv project, run the following command:

shell

uv add enceladus

If you plan to pass PyTorch tensors or MLX arrays to kernels, install the matching extra, for example uv add "enceladus[torch]" or uv add "enceladus[mlx]".

To check that Enceladus finds your GPU, run the following command:

shell

uv run python -c "import enceladus; print(enceladus.get_device())"

The output is similar to the following:

<enceladus.Device Apple M4 Pro (applegpu_g16s)>

To build from source, clone the repository, run uv sync, and then run the tests with uv run pytest -q.

Run four kernels

Each of the following programs is complete. Save one to a file and run it with uv run python FILE.py, where FILE is the name you chose.

import numpy as np

import enceladus
import enceladus.language as tl


@enceladus.jit
def add_kernel(x_ptr, y_ptr, out_ptr, n, BLOCK: tl.constexpr):
pid = tl.program_id(0)
offs = pid * BLOCK + tl.arange(0, BLOCK)
mask = offs < n
x = tl.load(x_ptr + offs, mask=mask)
y = tl.load(y_ptr + offs, mask=mask)
tl.store(out_ptr + offs, x + y, mask=mask)


n = 98_432
x = enceladus.randn(n, seed=0)
y = enceladus.randn(n, seed=1)
out = enceladus.empty_like(x)
grid = (enceladus.cdiv(n, 1024),)
add_kernel[grid](x, y, out, n, BLOCK=1024)

np.testing.assert_allclose(out.numpy(), x.numpy() + y.numpy())
print("vector add matches NumPy")

import numpy as np

import enceladus
import enceladus.language as tl


@enceladus.jit
def softmax_kernel(out_ptr, in_ptr, stride_in, stride_out, n_cols,
               BLOCK: tl.constexpr):
row = tl.program_id(0)
cols = tl.arange(0, BLOCK)
mask = cols < n_cols
x = tl.load(in_ptr + row * stride_in + cols, mask=mask, other=-float("inf"))
x = x.to(tl.float32)
x = x - tl.max(x, axis=0)
num = tl.exp(x)
tl.store(out_ptr + row * stride_out + cols, num / tl.sum(num, axis=0), mask=mask)


x = np.random.default_rng(0).standard_normal((512, 1000), dtype=np.float32)
out = np.empty_like(x)
n_rows, n_cols = x.shape
softmax_kernel[(n_rows,)](
out, x, enceladus.element_strides(x)[0], enceladus.element_strides(out)[0], n_cols,
BLOCK=enceladus.next_power_of_2(n_cols),
)

expected = np.exp(x - x.max(axis=1, keepdims=True))
expected /= expected.sum(axis=1, keepdims=True)
np.testing.assert_allclose(out, expected, rtol=1e-5, atol=1e-6)
print("softmax matches NumPy")

import numpy as np

import enceladus
import enceladus.language as tl


@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))


M, N, K = 1000, 800, 600
rng = np.random.default_rng(0)
a = rng.standard_normal((M, K), dtype=np.float32)
b = rng.standard_normal((K, N), dtype=np.float32)
c = np.empty((M, N), np.float32)
grid = (enceladus.cdiv(N, 64), enceladus.cdiv(M, 64))
matmul_kernel[grid](a, b, c, M, N, K, K, N, N, BM=64, BN=64, BK=32, num_warps=4)

np.testing.assert_allclose(c, a @ b, rtol=1e-4, atol=1e-3)
print("matmul matches NumPy")

# The kernel from examples/08_flash_attention.py: online softmax, float32 accumulator.
@enceladus.jit
def attention_kernel(q_ptr, k_ptr, v_ptr, o_ptr, sm_scale, N_CTX, stride_h, stride_m,
                 HEAD_DIM: tl.constexpr, BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr,
                 CAUSAL: tl.constexpr):
start_m = tl.program_id(0)
base = tl.program_id(1) * stride_h
q_desc = tl.make_tensor_descriptor(q_ptr + base, [N_CTX, HEAD_DIM], [stride_m, 1],
                                   [BLOCK_M, HEAD_DIM])
k_desc = tl.make_tensor_descriptor(k_ptr + base, [N_CTX, HEAD_DIM], [stride_m, 1],
                                   [BLOCK_N, HEAD_DIM])
v_desc = tl.make_tensor_descriptor(v_ptr + base, [N_CTX, HEAD_DIM], [stride_m, 1],
                                   [BLOCK_N, HEAD_DIM])
o_desc = tl.make_tensor_descriptor(o_ptr + base, [N_CTX, HEAD_DIM], [stride_m, 1],
                                   [BLOCK_M, HEAD_DIM])
offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M)
offs_n = tl.arange(0, BLOCK_N)
m_i = tl.full([BLOCK_M], float("-inf"), tl.float32)
l_i = tl.zeros([BLOCK_M], tl.float32)
acc = tl.zeros([BLOCK_M, HEAD_DIM], tl.float32)
qk_scale = sm_scale * 1.4426950408889634  # log2(e), so the softmax can use exp2
q = q_desc.load([start_m * BLOCK_M, 0])
hi = N_CTX
if CAUSAL:
    hi = tl.minimum((start_m + 1) * BLOCK_M, N_CTX)
for start_n in range(0, hi, BLOCK_N):
    qk = tl.dot(q, tl.trans(k_desc.load([start_n, 0]))) * qk_scale
    cols = start_n + offs_n
    mask = cols[None, :] < N_CTX
    if CAUSAL:
        mask = mask & (offs_m[:, None] >= cols[None, :])
    qk = tl.where(mask, qk, float("-inf"))
    m_ij = tl.maximum(m_i, tl.max(qk, 1))
    p = tl.exp2(qk - m_ij[:, None])
    alpha = tl.exp2(m_i - m_ij)
    l_i = l_i * alpha + tl.sum(p, 1)
    acc = acc * alpha[:, None]
    acc = tl.dot(p.to(q_desc.dtype), v_desc.load([start_n, 0]), acc)
    m_i = m_ij
acc = acc / l_i[:, None]
o_desc.store([start_m * BLOCK_M, 0], acc.to(o_desc.dtype))


# Launch: one program per BLOCK_M query rows of each (batch, head) pair.
# dot_warps=(num_warps, 1) gives each SIMD group whole rows, so row reductions stay local.
grid = (enceladus.cdiv(n, 32), z * h)
attention_kernel[grid](q, k, v, o, scale, n, n * d, d, HEAD_DIM=d, BLOCK_M=32,
                   BLOCK_N=32, CAUSAL=causal, num_warps=4, dot_warps=(4, 1))

The grid launches cdiv(98_432, 1024) = 97 programs. The last one covers only 128 valid elements, and mask keeps it in bounds.

How the vector add works

The vector add shows the core ideas:

tl.program_id(0) returns the index of the current program. The grid launches enough programs to cover all n elements.
tl.arange(0, BLOCK) creates a tile of 1,024 indices. x_ptr + offs is a tile of pointers.
mask keeps the last program from reading or writing past the end of the arrays.
BLOCK: tl.constexpr makes the block size a compile-time constant. Enceladus compiles one version of the kernel for each value you pass.

The launch returns before the GPU finishes when every array argument is an enceladus.Tensor. out.numpy() waits for the result. A launch with NumPy arrays waits for the GPU before it returns, so the arrays hold the result right away.

Run a kernel on the CPU

Enceladus includes an interpreter that runs kernels on the CPU with NumPy, one program at a time. In the interpreter, you can use print() and pdb inside a kernel. To run the vector add in the interpreter, run the following command:

shell

ENCELADUS_INTERPRET=1 uv run python add.py

The interpreter is much slower than the GPU. Use it to debug a kernel and to check the compiled kernel's results.

What's next

To learn how programs, tiles, and SIMD groups fit together, see Programming model.
To look up a tl function, see Language reference.
To make kernels faster, see Performance.
To use Enceladus with PyTorch or MLX, see Framework interop.
To port an existing Triton kernel, see Porting from Triton.