computeml — Python bindings

Python API for ComputeML: ML inference and training on any GPU via Vulkan compute shaders. Pure-Python ctypes binding over libcomputeml.so — no compilation step, only NumPy required.

Setup

# 1. Build the shared library (embeds all SPIR-V shaders)
cmake -B build -S . && cmake --build build --target computeml_shared

# 2. Install the package (editable, from repo root)
pip install -e python/

The binding finds the library via $COMPUTEML_LIB, a bundled copy next to the package, or <repo>/build/libcomputeml.so.

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

Tests

python -m pytest python/tests/ -v

test_ops.py checks GPU ops against NumPy references; test_bert.py checks BERT-tiny end-to-end numerical equivalence against HuggingFace PyTorch.