computeml

Run ML inference and training on any GPU from Python, through Vulkan compute shaders. You work with NumPy arrays. The wheel ships the compiled library, so you only need NumPy and a Vulkan driver.

Install

pip install computeml

uv add computeml

Requirements:

• Linux x86_64 (glibc 2.27 or newer) and Python 3.10 or newer.
• A Vulkan driver. NVIDIA, AMD and Intel GPU drivers include one. With no GPU, Mesa lavapipe runs on the CPU: apt install mesa-vulkan-drivers.

Quickstart

import numpy as np
import computeml as cml

ctx = cml.Context()
print(ctx.device_name)

a = ctx.asarray(np.random.rand(512, 256))
b = ctx.asarray(np.random.rand(256, 128))
c = cml.ops.relu(a @ b) * 2.0
print(c.shape, c.numpy()[0, :4])

ctx.asarray copies an array to the GPU. Ops queue work on the GPU. .numpy() waits for the work and copies the result back.

Coming from NumPy, PyTorch or CuPy

Differences you will meet:

• No broadcasting. Shapes must match element for element, or the op raises ValueError.
• float32 only for inputs and outputs. asarray converts ints and float64 to float32.
• tensor + number is not supported yet. tensor * number and tensor / number work.
• np.asarray(tensor) raises TypeError. Call .numpy(), so every GPU-to-CPU copy is visible.
• Every layer takes the Context first. There is no default device.

Layers

import numpy as np
import computeml as cml

ctx = cml.Context()
model = cml.nn.Sequential(
    cml.nn.Linear(ctx, 16, 32, activation="gelu"),
    cml.nn.Linear(ctx, 32, 8),
    cml.nn.LayerNorm(ctx, 8),
)
x = ctx.asarray(np.random.rand(4, 16))
print(model(x).numpy().shape)

Linear weights use the PyTorch layout (out, in). Load trained weights with cml.nn.Linear.from_numpy(ctx, weight, bias).

Training

import numpy as np
import computeml as cml

ctx = cml.Context()
X = np.random.rand(128, 4)
y = (X.sum(axis=1) > 2).astype(np.float32)

model = cml.nn.Sequential(
    cml.nn.Linear(ctx, 4, 16, activation="relu"),
    cml.nn.Linear(ctx, 16, 2),
).requires_grad_()
opt = cml.optim.AdamW(model.parameters(), lr=0.05)

x, t = ctx.asarray(X), ctx.asarray(y)
for _ in range(100):
    opt.zero_grad()
    loss = cml.ops.cross_entropy(model(x), t)
    loss.backward()
    opt.step()
print(float(loss))

Gradients are computed on the GPU by the C library's autograd tape. See the training guide for which ops can be trained.

Fused ops

One kernel does the work of several ops, so the data is read from GPU memory once.

import numpy as np
import computeml as cml

ctx = cml.Context()
x = ctx.asarray(np.random.rand(8, 64))
w = ctx.asarray(np.random.rand(128, 64))
b = ctx.asarray(np.random.rand(128))
y = cml.ops.linear(x, w, b, activation="gelu")  # gelu(x @ w.T + b), one kernel
print(y.shape)

cml.ops also has softmax, layer and RMS norm, flash attention (multi-head, causal), RoPE, embeddings, permute, concat and reductions. See ops.py.

Capture and replay

ctx.capture records a function's GPU work on the first call. Later calls resubmit the recorded work with new input values, with no Python per op.

import numpy as np
import computeml as cml

ctx = cml.Context()
w = ctx.asarray(np.random.rand(64, 64))

@ctx.capture
def step(x):
    return cml.ops.relu(x @ w)

for _ in range(3):
    y = step(np.random.rand(32, 64))
    print(y.numpy().sum())
print(step.record_count)  # 1: recorded once, replayed twice

The output is the same tensor on every call, so read it before the next call. Any other op on the context in between makes the next call record again.

Quantized weights

import numpy as np
import computeml as cml

ctx = cml.Context()
x = ctx.asarray(np.random.rand(16, 256))
w = ctx.asarray(np.random.rand(256, 64)).quantize("Q8_0")
print((x @ w).numpy().shape)  # w is dequantized inside the matmul kernel

Errors and logs

• cml.NoDeviceError: no Vulkan driver or GPU. The message says what to install.
• ImportError on import computeml: the Vulkan loader libvulkan.so.1 is missing.
• ValueError / TypeError: a wrong shape or dtype, raised before any GPU work.
• Set CML_LOG=1 to print the chosen GPU and kernel class to stderr.
• Set CML_GPU_INDEX=1 to use the second GPU.

Build from source

Set COMPUTEML_LIB to use your own libcomputeml.so. To build it and run the tests from a checkout of the repo:

cmake -B build -S . && cmake --build build --target computeml_shared
cd python && uv run --with pytest python -m pytest -q

A BERT-tiny example lives in examples/python/bert_tiny.py.