Nx.Vulkan

A GPU tensor backend for Nx that runs on anything with a Vulkan driver — including FreeBSD, where CUDA and Metal don't exist.

✓ Linux + NVIDIA RTX 3060 Ti      (proprietary driver)
✓ FreeBSD + NVIDIA GT 750M        (NVIDIA legacy driver)
✓ FreeBSD + NVIDIA GT 650M        (NVIDIA legacy driver)

Why this exists →WHY.md — the f64 conviction, autograd-for-free, reach over peak FLOPS, and one-GPU-to-a-fleet.

Goals

• Cover the FreeBSD gap. Nx's existing GPU backends (EXLA on Linux+CUDA, EMLX on macOS+Metal) don't run on FreeBSD. If you have NVIDIA on FreeBSD, this is the only path.
• Portable Vulkan across vendors. Same SPV blobs run on NVIDIA, AMD RADV, Intel, and via MoltenVK on macOS. The Rust vulkano crate keeps the driver surface pure Rust — no C++ FFI ownership traps.
• Enough coverage that Nx models Just Work. A native op set plus a host-fallback long tail let real workloads (Axon training, eXMC NUTS sampling, Scholar linear regression) run today without op-by-op porting.
• Deterministic across hardware. Cross-Kepler runs of the same IR produce byte-identical posterior chains. Reproducible science across GPU generations is a documented feature.

What works today

• A native compute op set — elementwise binary/unary, reductions (sum / max / min over ANY axis set, including a kept axis in the middle), reshape / squeeze, transpose (any permutation, rank <= 4), matmul (any rank-2 contraction orientation), conv (im2col + GEMM, any layout permutation), FFT, select / compare / cast / broadcast / reverse, and max/min pooling in both directions. Native f32 and f64 — the hot ops (elementwise, matmul, conv, reduce, transpose) dtype-dispatch native f32 shaders as well as f64; f64 is the default accumulator policy, f32 wins on bandwidth-bound ops.
• Whole-graph fusion via Nx.Vulkan.Compiler, an Nx.Defn.Compiler — see The Nx.Defn fusion compiler.
• Batched command submission — dispatches are recorded and submitted as one command buffer with one fence wait instead of a submit-and-block per op, worth 1.45–1.71× on a training step across the fleet. Flushed automatically at every host boundary, so correctness never depends on managing it. NXV_BATCH_MAX=0 restores submit-per-dispatch; flush/0 exists for benchmarks that need to time the work rather than the recording of it.
• Host fallback for the long tail — sort/argsort, SVD/QR/solve/cholesky from Nx.LinAlg (via Nx.Block.LinAlg), rank-5+ shapes, and pow in f64 (GLSL.std.450 has no f64 pow). Slow but correct.
• A fallback counter (Nx.Vulkan.Fallback) — a host fallback returns a bit-identical result, so no assertion on values can detect one. count/1 makes it countable, and the suite asserts zero fallbacks for ops that must stay on-device. A full CNN training step now performs exactly one: pow in f64.
• Nx.Defn.grad autograd, for free — see The autograd insight.
• Axon training step end-to-end, gradient sum agrees to 1e-8 vs BinaryBackend reference.
• eXMC NUTS regime log_p at f64: byte-identical to CPU reference.
• Scholar linear regression: coefficients match BinaryBackend to 2e-6 (SVD via host fallback).
• Long-running stability — 7000+ chain-shader dispatches on Ampere without crash after the 2026-07 primary_buffer_count=128 fix; unbounded on Kepler.
• Pipeline cache on disk, UUID-validated, survives BEAM restarts.

Roadmap and future work: ROADMAP.md.

The Nx.Defn fusion compiler (thrust 3)

EXLA's structural edge over an eager backend is whole-graph compilation: it fuses a chain of ops into one kernel instead of dispatching each separately. Nx.Vulkan.Compiler is an Nx.Defn.Compiler that does the same for the cases it supports — the closest this project comes to closing that gap.

Nx.Defn.jit(&my_fun/2, compiler: Nx.Vulkan.Compiler).(a, b)

It traces a defn to an expression DAG and compiles it to a stage schedule that runs on-device with GPU-resident intermediates and no fallback to the interpreter:

• Elementwise fusion — a same-shape f32/f64 chain becomes one generated GLSL shader, one dispatch (3.62× over eager per-op on a 10-op chain). Broadcasting, CSE-within-a-shader, and scalar constants are baked in.
• Parallel fused reductions — an elementwise chain feeding sum/product/max/min/mean fuses into one workgroup-per-slot tree reduce (f64 accumulator, matches BinaryBackend).
• Multi-stage split at boundaries — a graph with a dot, conv, reduce, or transpose splits into stages: each boundary is a stage, each maximal elementwise region between boundaries is one shader whose inputs may be earlier stages' buffers. reshape/squeeze are zero-copy view boundaries (no dispatch); transpose moves data.
• Multi-output — a defn returning a tuple compiles to one shared schedule; subexpressions common to several outputs are computed once.
• f32 and f64 — the codegen is dtype-parameterised (f64 gated on the device advertising shader_float64; f64 transcendentals host-fall-back rather than silently losing precision).

Whole layers fuse end-to-end: relu(x @ W + b), relu(conv(x, k) + b), a CNN classifier head (conv → flatten → dense), softmax and layernorm reduction patterns, and transposed-weight layers (x @ Wᵀ). Anything unsupported falls back to Nx.Defn.Evaluator, so results are always correct — worst case is "no fusion, same as eager."

Standing

The fusion compiler is the goal this effort set out to reach: a credible #2 compute backend for elixir-nx, with EXLA's whole-graph compilation now present in the one place a Vulkan backend can offer it — on any GPU with a driver, CUDA or not.

• Correctness first. Every fused result is checked exact against Nx.BinaryBackend. The suite — 851 doctests, 415 tests, 0 failures — is green on both a 2012 Kepler (GT 650M, FreeBSD) and a 2021 Ampere (RTX 3060 Ti, Linux), with the f64 fused path active on both. Gradient parity and host-fallback counts are asserted, not assumed: Nx.Defn.grad is compared against BinaryBackend op by op, and a CNN training step is asserted to leave the GPU exactly once.
• Fusion's win is structural — and it has a floor. It removes dispatches and intermediate buffers and keeps the interpreter out of the loop; it does not make kernels faster. So the gain grows with the elementwise work around a boundary — and below some amount of it, fusion is a net loss. Measured on a dense-only MLP, where the graph is almost all dot and there is nothing to amortise tracing and boundary buffers against, Nx.Vulkan.Compiler runs at 0.76× of eager — a 24% regression — and at 0.98× on a conv CNN. Both are bit-identical to the host; they are correct and slower. Fusion is opt-in (compiler: Nx.Vulkan.Compiler), so this is a choice to make per graph shape rather than a default you inherit.
• Every heuristic is fleet-validated, never assumed. Win/loss crossovers are hardware-specific, so they are measured across the fleet (Kepler + Ampere), not the local box. The many-slot reduce is device-class-gated because it wins on weak GPUs and regresses on strong ones. Batched submission was raced the same way and wins on all three hosts with no crossover, which is why it needs no gate and ships on by default — the outcome, not the assumption. Cross-stage CSE was built, raced, and found to never win on either device class (recompute is cheaper than the dispatch it takes to avoid) — so it ships default-off, opt-in via NXV_CSE=1. See bench_results/CSE_SOFTMAX_RACE.md and the write-up, Compute It Twice: When CSE Lost the Race.

Building on a compute kernel of your own? See the vulkan-nx-compute skill for the shader → NIF → Nx playbook and the hard-won parity/dispatch gotchas.

Position vs EXLA and EMLX

The autograd insight

Nx.Defn.grad is a graph transformation that runs at compile time on the Nx.Defn.Expr AST. For every forward op in the graph, it inserts the corresponding backward op expressed in terms of more forward ops. The backend never sees a "backward op" — it just keeps executing forward primitives. Forward op coverage IS gradient coverage when running through Nx.Defn.Evaluator.

That means VulkanoBackend supports gradients for any function expressible in its native ops + host-fallback long tail. No backward callbacks were written. Validated by running a complete Axon training step (Dense → sigmoid → Dense → MSE → Nx.Defn.value_and_grad) on Nx.Vulkan.VulkanoBackend, with gradient sum agreeing to 1e-8 against the BinaryBackend reference.

Benchmarks

Batched dispatch (August 2026)

One value_and_grad step of an MNIST MLP at batch 32, submit-per-dispatch (NXV_BATCH_MAX=0) vs batched. Same graph, same commit, arms differing only in the environment variable:

The loss is identical in every arm on every host — two architectures, two operating systems. Batching changes when work is submitted, never what is computed. Reproduce with examples/mnist_mlp_step_bench.exs; method and the full cap sweep in bench_results/BATCHED_DISPATCH.md.

The three tables below predate batched dispatch and are therefore pessimistic by roughly the factors above. They have not been re-run because their harnesses were scratch projects rather than committed examples — so they are left at their measured values rather than adjusted by arithmetic.

CNN training (August 2026, pre-batching)

One value_and_grad step, batch 32, versus Nx.BinaryBackend. Losses are bit-identical to the host on every row.

The LeNet step was 20.9 s before the backward pass stopped falling back to the host — the same measurement, same box. Read the absolute times rather than the multipliers: a speedup here mostly measures how slow pure-Elixir BinaryBackend is, which varies by host CPU. The absolute GPU times cluster in 25–85 ms across three cards spanning 2012–2021, because at this model size the work is dispatch-bound rather than compute-bound — which is why a 2012 laptop GPU is competitive with a 2021 desktop one.

vs EXLA (August 2026, pre-batching)

The Axon MNIST guide model, one training step at batch 32 on the RTX 3060 Ti — a dense-only MLP, which is the shape most favourable to EXLA and least favourable here:

EXLA is ~20× ahead, and fusion does not close it — it costs 24% (0.76× fused vs eager). On a graph that is almost all dot, there is no elementwise work for fusion to amortise against, so the deficit is dispatch overhead and GEMM quality rather than missing whole-graph compilation.

That diagnosis is what batched submission acts on, and it is why this race is labelled pre-batching: the eager row above has since improved ~1.7× on this box. The race has not been re-run — it needs a working EXLA, which this repo deliberately does not depend on — so no combined figure is claimed here. The remaining lever the measurement points at is GEMM quality.

On a 2×strided-conv CNN the gap is similar — 41.3 ms eager vs 1.45 ms, with fusion neutral at 0.98× — so EXLA leads on both graph shapes tested. The qualitative difference is availability: EXLA cannot be installed on the two FreeBSD Keplers at all, where Vulkan runs the same CNN in 64–78 ms. Full numbers, an XLA gradient-compile failure narrowed to one specific conv configuration, and what it took to get EXLA running: bench_results/MNIST_EXLA_RACE.md.

f32 vs f64 per op

sh scripts/race.sh — f32 speedup over f64, same shapes, all on-GPU:

f32 matmul is slower than f64, and that is the accumulator policy working as designed: f32 matmul defaults to matmul_f32_f64acc.spv, paying an f32→f64 conversion on top of the same f64 MAC rate. f32 wins where it is meant to — bandwidth-bound elementwise and reductions. Switch with Nx.Vulkan.VulkanoBackend.put_f32_matmul_accumulator(:f32) if you want speed over the f64-accumulated reference.

Square matmul (May 2026)

Milliseconds per dispatch, median of 50–200 iterations:

Full bench: examples/full_bench.exs.

Quickstart

As a backend in your project

# mix.exs
def deps do
  [
    {:nx, "~> 0.13"},
    {:nx_vulkan, "~> 0.3"}
  ]
end

Or track main directly:

{:nx_vulkan, git: "https://github.com/borodark/nx_vulkan"}

On 0.2.0 and training on the GPU? Upgrade. 0.2.0 computed correct results, but its backward pass ran largely on the host — a LeNet step took 20.9 s there and 84 ms here. See the CHANGELOG.

# Build a tensor, transfer to GPU, do work
x_bin = Nx.tensor([1.0, 2.0, 3.0, 4.0], type: :f32)
x_vk  = Nx.backend_transfer(x_bin, Nx.Vulkan.VulkanoBackend)

y_vk  = Nx.sigmoid(x_vk)
y_bin = Nx.backend_transfer(y_vk, Nx.BinaryBackend)
IO.inspect(Nx.to_list(y_bin))
# [0.7310585975646973, 0.8807970881462097, 0.9525741338729858, 0.9820137619972229]

Try the Axon training example

git clone https://github.com/borodark/nx_vulkan
cd nx_vulkan
mix deps.get && mix compile
elixir examples/axon_training_loop.exs

Runs a 100-step Dense(4→32, tanh)→Dense(1) regression with manual SGD. Compares loss trajectories on BinaryBackend vs VulkanoBackend. PASS verdict on both Linux + FreeBSD.

Try the full bench

mix run examples/full_bench.exs

Per-op + end-to-end + robustness across every backend Nx can find. Auto-detects EXLA availability. Runs in ~10 minutes on RTX 3060 Ti, ~15 on GT 650M.

Building

Prerequisites

• Erlang/OTP 26+, Elixir 1.17+
• Rust 1.78+
• C++ compiler (only needed for the legacy spirit backend; vulkano is pure Rust)
• Vulkan SDK + glslangValidator:
  • Debian/Ubuntu: apt install libvulkan-dev vulkan-tools glslang-tools
  • FreeBSD: pkg install vulkan-loader vulkan-headers vulkan-tools glslang shaderc

Build

mix deps.get
mix compile

Vulkano compiles in ~30s on Linux, ~3:18 on FreeBSD 15.0 (mostly dependency compilation). The spirit/C++ path compiles in parallel.

Rust toolchain pin

rust-toolchain.toml pins rustc to 1.85. The reason is in the file's comment; bump when upstream rustler emits a corrected rustler-sys signature.

Blog series

• Compute It Twice: When CSE Lost the Race — the Nx.Defn fusion compiler; why cross-stage CSE never won the fleet race
• The Backend That Didn't Need to Know — the C++→vulkano migration; descriptor pool debugging; autograd was free
• The GPU That Doesn't Need CUDA — the FreeBSD Vulkan story (spirit-era)
• A Walkable Path Under the Mountain — eXMC + zed integration

Sibling: zed

zed is the declarative ZFS + Elixir deploy tool that orchestrates BEAM nodes. nx_vulkan is consumed inside deployed BEAM nodes — not as a zed dependency. See specs/nx-vulkan-execution.md in the zed repo for the integration story.

License

Apache 2.0. Same as Spirit and Nx.