<div align="center">
# ⚡ nuvec
**MLX for Rust, written in Rust.** Lazy n-dimensional arrays (tensors), composable transforms and Metal kernels for Apple Silicon, with no bindings to MLX or any C/C++ library.
[](https://github.com/hupe1980/nuvec/actions/workflows/ci.yml)
[](https://hupe1980.github.io/nuvec)
[](#-license)
[Documentation](https://hupe1980.github.io/nuvec/docs/) · [Getting started](https://hupe1980.github.io/nuvec/docs/getting-started/) · [Benchmarks](https://hupe1980.github.io/nuvec/docs/benchmarks/)
</div>
---
nuvec brings [MLX](https://github.com/ml-explore/mlx)'s programming model to Rust: arrays that are evaluated lazily on unified memory, `grad`/`vjp`/`jvp`/`vmap`/`compile` that compose freely, and the Metal kernels that make MLX fast, ported and tested against MLX itself. On top sits an LLM runtime that loads mlx-community checkpoints and generates as fast as mlx-lm.
> 🚧 **Pre-alpha.** The API can change between releases.
## ✨ Highlights
- 🦀 **Pure Rust.** No C++ build, no FFI to MLX. Metal goes through [`objc2-metal`](https://crates.io/crates/objc2-metal). `unsafe` is confined to the Metal runtime, three kernel dispatches in nuvec-core and the CPU matrix product.
- 💤 **Lazy arrays on unified memory.** Operations build a graph that runs when you read a result. CPU and GPU share memory, with no copies between them.
- 🔁 **Composable transforms.** `grad`, `value_and_grad`, `vjp`, `jvp`, `vmap`, `custom_function`, `checkpoint` and `compile` (shaped or shapeless) nest in any order, over complex numbers, tuples, maps and your own structs.
- 🧠 **Models as plain structs.** `#[derive(Module)]` gives MLX's parameter names; optimizers, freezing, LoRA, `save_weights` and a compiled training step (`compile_mut`) work on the model's own type.
- 🔥 **MLX's Metal kernels.** GEMM and GEMV, quantized products (affine 1, 2, 3, 4, 5, 6 and 8 bits; MXFP4, MXFP8, NVFP4), fused attention, RMSNorm, RoPE, sort, scan, mixture-of-experts gather kernels, and your own kernels from source (`fast::MetalKernel`). Metal 4 TensorOps kernels for M5 Neural Accelerators are opt-in (`NUVEC_NAX=on`).
- 🌐 **Several Macs.** MLX's collectives (`all_sum`, `all_gather`, `send`/`recv`, …) as lazy, differentiable array operations over MLX's TCP ring, with tensor-parallel layers.
- 🧪 **Tested against MLX.** Over a thousand cases replay results recorded from MLX, with random layouts, gradients and every dtype.
- 🤖 **LLM runtime.** Llama, Mistral, Qwen 2/3/3.5, gpt-oss and Gemma 4 (dense, hybrid Gated DeltaNet, mixture-of-experts) from mlx-community or Hugging Face checkpoints. Speculative decoding (MTP heads, prompt lookup, SuffixDecoding) and batched serving.
- 💾 **Memory-mapped weights.** Aligned checkpoints (such as those nuvec writes) go to the GPU without a copy. Unaligned ones (mlx-community's) are read once, in parallel, with no more memory than MLX.
- ⚙️ **Built for servers.** Attention plans and score-memory caps, a memory limit that refuses instead of swapping, checkpoint footprints before loading, batch-invariant serving and documented determinism. Used by the [mrun](https://github.com/hupe1980/mrun) inference server.
## 📊 Performance
nuvec-lm against mlx-lm 0.32.2 with the same token ids on an Apple M4 Pro (48 GB), measured on 2026-10-31 with nuvec 0.5.0 (`just bench-parity --tier full`). Decoding is greedy; ratios above 1 favour nuvec.
| Qwen3-0.6B 4-bit | dense | 1.10× / 1.10× / 1.06× | 1.01× / 1.01× | 0.98× |
| Llama 3.2 3B 4-bit | dense | 0.99× / 1.00× / 1.00× | 1.00× / 1.01× | 1.00× |
| Llama 3.1 8B 4-bit | dense | 1.01× / 1.00× / 1.00× | 1.00× / 1.00× | 1.01× |
| Qwen3-8B 4-bit | dense | 1.00× / 1.01× / 1.01× | 1.00× / 1.01× | 1.01× |
| Qwen3.5-9B 4-bit | hybrid | 1.00× / 1.00× / 1.00× | 1.00× / 0.99× | 1.00× |
| Qwen3.8-27B 4-bit | hybrid | 1.00× / 1.01× / 1.05× | 1.01× / 1.00× | 1.00× |
| Qwen3-30B-A3B 4-bit | MoE | 1.00× / 1.00× / 1.00× | 0.99× / 1.00× | 0.98× |
| Qwen3.6-35B-A3B 4-bit | hybrid MoE | 0.99× / 1.00× / 1.00× | 1.01× / 0.98× | 0.98× |
| gpt-oss-20b MXFP4 | MoE, sliding window | 0.99× / 0.99× / 0.98× | 1.04× / 1.03× | 1.00× |
| gpt-oss-20b 8-bit | MoE, sliding window | 1.01× / 1.01× / 1.00× | 1.02× / 1.02× | 1.08× |
| Gemma 4 E4B 4-bit | dense, sliding window | 0.98× / 1.00× / 1.00× | 1.07× / 1.07× | 1.02× |
| Gemma 4 26B-A4B 4-bit | MoE, sliding window | 1.01× / 1.01× / 1.01× | 1.06× / 1.08× | 0.99× |
- **Geometric means:** decode 1.01×, time to first token 1.02× at 2K and 32K tokens, peak memory 1.02× (nuvec uses slightly less), warm load to the first token 1.80×, 1 / 4 / 8 concurrent requests 1.01× / 1.00× / 1.00×.
- **Longest contexts:** Qwen3.6-35B-A3B's first token at 64K tokens is 0.98×. At 261,952 tokens nuvec prefills it in 20.4 minutes; mlx-lm runs out of GPU memory on the same machine.
- **Speculative decoding** with the model's multi-token-prediction head: Qwen3.8-27B 1.66× and Qwen3.5-9B 1.36× faster than plain decoding on natural text, with identical output.
The full record (JSON) and method are on the [benchmarks](https://hupe1980.github.io/nuvec/docs/benchmarks/) page.
## 🚀 Quick start
```sh
cargo add nuvec # arrays, transforms, Metal kernels
cargo add nuvec-nn # layers and optimizers (optional)
cargo add nuvec-lm # the LLM runtime (optional)
```
**Arrays are lazy, and gradients are arrays too:**
```rust
use nuvec::prelude::*;
use nuvec::grad;
fn main() -> Result<()> {
let x = array![1.0, 2.0, 3.0];
let y = (&x * 2.0 + 1.0) * &x; // builds a graph, computes nothing
println!("{:?}", y.to_vec::<f32>()?); // [3.0, 10.0, 21.0], on the GPU
let f = |x: &Array| x.multiply(x).sum(..); // sum(x²)
println!("{:?}", grad(f)(&x)?.to_vec::<f32>()?); // [2.0, 4.0, 6.0]
Ok(())
}
```
**Create, index and chain; errors surface where you read results:**
```rust
use nuvec::prelude::*;
use nuvec::ops as mx;
fn main() -> Result<()> {
let a = array![[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]; // float32, shape [2, 3]
let b = Array::from(vec![10.0f32, 20.0, 30.0]); // from Rust values
let z = mx::zeros([2, 3], DType::F32); // free functions live in `ops`
let cols = a.i((.., 1..3)); // MLX's a[:, 1:3], a view
let z = z.at(0).set(&b); // MLX's z[0] = b, a new array
assert_eq!(cols.sum(..).item::<f32>()?, 16.0);
assert_eq!(z.sum_keepdims(-1).shape(), [2, 1]);
let bad = a.matmul(&a).exp(); // [2, 3] @ [2, 3]: a failed array
let err = bad.check().unwrap_err(); // no panic, no `Result` per op
assert_eq!(err.op_name(), Some("matmul")); // with your file and line
Ok(())
}
```
Operations return arrays, never `Result`. An invalid call gives a *failed array* that records the operation, call site and reason; anything computed from it fails too. The error surfaces at `check()?`, `eval()?`, `item()?`, `to_vec()?`, I/O and transforms.
**Train a model: a struct with `#[derive(Module)]`, a compiled step over `(model, optimizer)`:**
```rust
use nuvec::{Array, Result, array, compile_mut, random, value_and_grad};
use nuvec_nn::optim::{Optimizer, Sgd};
use nuvec_nn::{Linear, Module, losses, relu};
#[derive(Clone, Module)]
struct Mlp { hidden: Linear, out: Linear } // parameters: hidden.weight, hidden.bias, …
fn main() -> Result<()> {
random::seed(0);
let mut state = (Mlp { hidden: Linear::new(2, 16), out: Linear::new(16, 2) }, Sgd::new(0.5));
state.1.init(&state.0);
let step = compile_mut(|(model, opt): &mut (Mlp, Sgd), (x, y): &(Array, Array)| {
let loss = |m: &Mlp| losses::cross_entropy(&m.out.forward(&relu(&m.hidden.forward(x))), y, losses::Reduction::Mean);
let (value, grads) = value_and_grad(loss)(model)?; // grads: Mlp
opt.update(model, &grads)?;
Ok(value)
});
let batch = (array![[1.0, -1.0], [-1.0, 1.0]], array![0u32, 1]);
let first = step(&mut state, &batch)?.item::<f32>()?;
for _ in 0..50 {
let _ = step(&mut state, &batch)?; // traced once, then replayed
}
assert!(step(&mut state, &batch)?.item::<f32>()? < first);
println!("{}", state.0.display()); // Mlp (82 params, 82 trainable) …
nuvec::save(std::env::temp_dir().join("ckpt.safetensors"), &state)?; // model and optimizer
Ok(())
}
```
**Generate text with an mlx-community checkpoint:**
```rust,no_run
use nuvec_lm::{GenerateOptions, generate, load};
fn main() -> nuvec::Result<()> {
let (model, tokenizer) = load("mlx-community/Qwen3-8B-4bit")?; // from the HF cache
let prompt = tokenizer.encode("The capital of France is");
let options = GenerateOptions { max_tokens: 32, ..Default::default() };
let ids: Vec<u32> = generate(&model, &prompt, options)?.collect::<nuvec::Result<_>>()?;
println!("{}", tokenizer.decode(&ids));
Ok(())
}
```
From the command line (after `hf download mlx-community/Qwen3-8B-4bit`):
```sh
cargo run --release -p nuvec-lm --example generate -- \
--model mlx-community/Qwen3-8B-4bit --prompt "Write a haiku about Rust" --chat
```
## 📦 Crates
| [`nuvec`](crates/nuvec) | The facade: arrays, transforms, `fast` kernels, random, I/O and the GPU runtime in one import |
| [`nuvec-core`](crates/nuvec-core) | Lazy graph, primitives, dtypes, evaluation, derivative and batching rules |
| [`nuvec-compile`](crates/nuvec-compile) | `compile` and `compile_mut`: tracing, simplification, element-wise kernel fusion |
| [`nuvec-macros`](crates/nuvec-macros) | `#[derive(Tree)]`, `#[derive(Module)]`, `#[derive(Static)]` (re-exported by `nuvec` and `nuvec-nn`) |
| [`nuvec-gpu`](crates/nuvec-gpu) | Metal runtime: pipelines, buffers, hazard-tracked command encoding, allocator, wired memory, profiler |
| [`nuvec-kernels`](crates/nuvec-kernels) | Metal kernel sources and the precompiled `nuvec.metallib` |
| [`nuvec-cpu`](crates/nuvec-cpu) | Pure-Rust CPU kernels |
| [`nuvec-io`](crates/nuvec-io) | safetensors, npy and IDX; memory-mapped loading |
| [`nuvec-distributed`](crates/nuvec-distributed) | MLX's TCP ring for `nuvec::distributed` (collectives across processes and Macs) |
| [`nuvec-nn`](crates/nuvec-nn) | Modules (`#[derive(Module)]`), layers (Linear, quantized, LoRA, norms, RoPE, dropout, tensor-parallel sharded linear), losses, optimizers (SGD, Adam, AdamW, Lion, Muon, parameter groups) |
| [`nuvec-lm`](crates/nuvec-lm) | LLM runtime: models, tokenizer, KV cache, sampling, speculative decoding, LoRA adapters |
## 🖥️ Platforms
| 🍎 **macOS 15+ on Apple Silicon** | The primary target, with the full Metal backend |
| 🐧 **Linux** | The core and CPU crates build and run; the GPU crates compile as stubs |
## 🛠️ Building from source
| Rust (pinned in `rust-toolchain.toml`) | [rustup](https://rustup.rs) installs it automatically |
| Metal toolchain | `xcodebuild -downloadComponent MetalToolchain` |
| [just](https://just.systems) | `brew install just` |
| [cargo-deny](https://github.com/EmbarkStudios/cargo-deny), Miri, jq | `cargo install cargo-deny --locked`, `rustup +nightly component add miri`, `brew install jq` |
```sh
git clone https://github.com/hupe1980/nuvec && cd nuvec
just ci # CI's gates: fmt, clippy, tests, docs, Miri, cargo-deny, kernel and hygiene checks
```
Without the Metal toolchain, kernels are compiled from embedded sources at runtime (`NUVEC_METALLIB=off` forces this).
### 🧪 Against MLX
Building and testing need no Python: the MLX fixtures are checked in. With [uv](https://docs.astral.sh/uv/) you can re-record them and benchmark against MLX:
| `just bench-parity` | nuvec-lm vs mlx-lm and the previous release, one model per architecture (about 45 minutes); `--tier full` for all models and long contexts (the table above) |
| `just bench-gemm` / `bench-quant` / `bench-fast` | matrix products, quantized products, fused kernels |
| `just bench-compile`, `just bench-load`, `just mnist` | compiled functions, checkpoint loading, MNIST training |
| `just fixtures`, `lm-fixtures`, `nn-fixtures`, `io-fixtures` | re-record the reference results |
`NUVEC_PROFILE=1` prints the GPU time per kernel of any run (see [profiling](https://hupe1980.github.io/nuvec/docs/gpu/#profiling)).
## 🧭 Stability
nuvec is 0.x: a minor release may break the API; the [changelog](CHANGELOG.md) lists every break. Semver covers items reachable through `nuvec::` paths and `nuvec-lm` (pinned in [`crates/nuvec/tests/api_surface.rs`](crates/nuvec/tests/api_surface.rs)); the sub-crates (`nuvec-core`, `nuvec-gpu`, …) make no promise when used directly. MSRV is Rust 1.99; raising it is a minor release.
## 📄 License
Licensed under either of [Apache License, Version 2.0](LICENSE-APACHE) or [MIT license](LICENSE-MIT), at your option. Unless you explicitly state otherwise, any contribution intentionally submitted for inclusion in this work, as defined in the Apache-2.0 license, shall be dual licensed as above, without any additional terms or conditions.
Code ported from MLX and mlx-lm keeps its MIT attribution ([THIRD_PARTY_NOTICES](THIRD_PARTY_NOTICES)). nuvec is not affiliated with Apple.