computeml

Python API for ComputeML: ML inference and training on any GPU through Vulkan compute shaders. 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).
• A Vulkan driver: NVIDIA, AMD, or Intel GPU drivers include one. With no GPU, Mesa lavapipe (apt install mesa-vulkan-drivers) runs on the CPU.

Usage

import computeml as cml
import numpy as np

ctx = cml.Context()                 # or Context(backend="vulkan")
print(ctx.backend)                  # "Vulkan"

a = ctx.from_numpy(np.random.rand(512, 512).astype(np.float32))
b = ctx.from_numpy(np.random.rand(512, 512).astype(np.float32))
c = cml.ops.matmul(a, b)            # recorded on GPU
print(c.numpy())                    # flush + download

# Quantized weights
w_q8 = b.quantize(cml.Dtype.Q8_0)
c2 = cml.ops.matmul(a, w_q8)        # dequantizes in-register on GPU

Ops are recorded into a command buffer and submitted on ctx.flush(). Tensor.numpy() flushes implicitly. Available ops: element-wise math, matmul (fp32/fp16/Q8/Q4), softmax, layernorm/rmsnorm, fused linear+activation, flash attention (multi-head, causal), RoPE, embeddings, reductions, see computeml/ops.py.

Models

from computeml.models import BertEncoder

enc = BertEncoder(ctx, "models/bert-tiny")   # weights from scripts/export_bert_tiny.py
hidden = enc.forward([101, 7592, 2088, 102]) # (seq, hidden) np.ndarray

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