Metadata-Version: 2.4
Name: computeml
Version: 0.1.1
Summary: ML inference and training on any GPU via Vulkan compute shaders
Author: Jorge Cardona
License: MIT
Project-URL: Homepage, https://gitlab.com/cardona-oss/compute-ml
Project-URL: Issues, https://gitlab.com/cardona-oss/compute-ml/-/tree/main/docs/issues
Classifier: License :: OSI Approved :: MIT License
Classifier: Operating System :: POSIX :: Linux
Classifier: Programming Language :: Python :: 3
Requires-Python: >=3.9
Description-Content-Type: text/markdown
Requires-Dist: numpy>=1.22
Provides-Extra: test
Requires-Dist: pytest; extra == "test"
Requires-Dist: torch; extra == "test"
Requires-Dist: transformers; extra == "test"

# computeml

Python API for [ComputeML](https://gitlab.com/cardona-oss/compute-ml): 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

```bash
pip install computeml
```

```bash
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

```python
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`](https://gitlab.com/cardona-oss/compute-ml/-/blob/main/python/computeml/ops.py).

## Models

```python
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:

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