Machine learning on the GPU you already have
A tensor library and neural-network framework in C++20 with a Python API. Compute shaders only — no CUDA, no ROCm, no vendor runtime.
Get started API reference GitHub
Train something in twenty lines
Install, pick a device, and run a training loop. The whole API is NumPy-shaped
and PyTorch-shaped on purpose — a state_dict from torch loads
without translation.
pip install vkmlimport vkml
from vkml import nn, optim
vkml.init_vulkan(0)
dev, why = vkml.best_device()
print(why) # using Vulkan device 0: AMD Radeon RX 5600M (discrete, ...)
model = nn.Sequential(nn.Linear(784, 128), nn.ReLU(), nn.Linear(128, 10)).to(dev)
opt = optim.Adam(model.parameters(), lr=1e-3)
for x, y in loader:
opt.zero_grad()
vkml.backward(nn.cross_entropy(model(x), y))
opt.step()Runs on any Vulkan GPU
Vulkan 1.3 compute is the only requirement. That covers AMD, Intel, NVIDIA, Apple through MoltenVK, and the integrated GPU in a laptop — the same binary, with no vendor toolkit to install and no per-vendor code path to maintain.
No vendor runtime
Twenty-four GLSL compute shaders compiled to SPIR-V at build time. There is no CUDA, no ROCm and no vendor SDK anywhere in the build.
Written against the guarantee, not the machine
Push-constant budgets, workgroup widths and subgroup ranges are checked
against what Vulkan 1.3 requires, not what the development GPU happens
to report. VKML_MIN_SPEC=1 makes any device report the guaranteed
floor so that is testable.
Deterministic by contract
Identical inputs give identical bytes on the same device, and across
drivers where the contract says so. f32→f16 narrowing is done in software
rather than left to OpFConvert, whose rounding mode SPIR-V leaves
implementation-defined.
Lazy by default
Operations build a graph; evaluation waits until something observes the result. Eager mode collapses that while debugging.
The CPU backend is the oracle, not a fallback
Every operator has a CPU implementation whose job is to be right, and the Vulkan one is checked against it. Both are checked against PyTorch. Where the two backends disagree, the difference is a documented tolerance with a stated kind — exact, ULP, relative or backward — not a fudge factor.
Tested against two references
1,434 Python tests compare against PyTorch and against the CPU oracle, including strided inputs, empty tensors, broadcasting, mixed precision and extreme magnitudes.
The tests are checked too
A mutation campaign breaks each kernel on purpose and confirms the suite notices. A green suite proves the tests ran, not that they can fail.
How it fits together
Nine layers, and the dependency direction is enforced by a build gate rather than by convention.
util 0 · core 1 · graph 2 · backend/api 3 · backend/cpu | backend/vulkan 4
dispatch 5 · plan 5 · api 6 · autograd 7A Python call enters the API layer, builds graph nodes, and returns. When a result is observed the dispatcher walks the graph, picks a backend per node, and submits. On Vulkan that means selecting one of six matmul pipelines, packing push constants, and recording a dispatch — all described in Concepts and the architecture pages.
Explore
Get started
Build it, check your device is usable, put a tensor on the GPU and train MNIST end to end.
Concepts
Lazy graphs, the two backends, dtypes and devices, and why the CPU one is the oracle.
API reference
102 operators and 27 classes, with signatures generated from the installed module so they cannot go stale.
Architecture
Tensors and views, the graph, autograd, both backends, the shader layer and the numerics.
Performance
Where the time actually goes, measured with timestamp queries, and what is being done about it.
Source
Apache-2.0. Issues, design records and the engineering handbook.