numba-vulkan

A Vulkan compute target for Numba: write a kernel in Python, run it on any GPU with a Vulkan driver.
Background: numba/numba#10116 »

Documentation · Usage · Benchmarks · Roadmap

About The Project

Numba's only maintained GPU target is CUDA, which ties GPU-accelerated Numba code to NVIDIA hardware. Supporting every vendor separately is a lot of work; Vulkan is implemented by nearly all of them.

numba-vulkan registers a vulkan target through Numba's target extension API and compiles numba.cuda-style kernels to SPIR-V compute shaders. The same kernel has been run on an NVIDIA GPU, an Intel integrated GPU and a CPU rasteriser (llvmpipe), with results checked against NumPy.

This is alpha software. Its test suite passes on NVIDIA and Intel GPUs and llvmpipe on Linux, on an Apple M3 Pro through MoltenVK, and on Windows with lavapipe, and the benchmarks run on an Intel Arc GPU under Windows; but its interfaces may still change and it has not yet been proven on real workloads.

(back to top)

Built With

• Numba: bytecode frontend, type inference, lowering and the target extension API
• libclc, LLVM's OpenCL math library, linked into kernels for accurate and double-precision math
• llvmlite 0.50+ (LLVM 22), whose SPIR-V backend emits Vulkan-flavoured SPIR-V
• vulkan: Python bindings for the Vulkan API
• devenv and uv for the development environment

How It Works

1. @nv.jit creates a dispatcher for the vulkan target. Numba's own pipeline (bytecode analysis, type inference, lowering) produces LLVM IR.
2. Arrays are a custom Numba type that carries the descriptor binding of its storage buffer, because shaders have no general pointers. Element access is emitted as a placeholder call taking a binding and an index.
3. Everything is inlined into one entry point and simplified with a few LLVM passes. Only then are the placeholders expanded into buffer accesses.
4. LLVM's SPIR-V backend turns the result into a compute shader. It runs in a child process, so a backend failure raises an exception instead of aborting Python.
5. The runtime uploads the arguments (binding 0: error status, binding 1 + k: array argument k; scalars and array shapes as push constants), dispatches, and copies written buffers back. Device arrays are used in place.

(back to top)

Getting Started

Prerequisites

• libclc's clspv--.bc from LLVM 22 or older, for float64 and accurate float32 math. Wheels bundle it and the devenv shell provides it; only a bare source checkout needs it installed (see the math library docs).
• A Vulkan 1.2 driver. lavapipe (Mesa's CPU implementation) is enough to try it without a GPU.
• Python 3.11 or newer.
• Optional: spirv-val from SPIRV-Tools to validate generated shaders.

Installation

From PyPI, with the Vulkan loader (libvulkan1 on Debian and Ubuntu, vulkan-loader on Fedora and Arch) and a Vulkan driver for your GPU installed:

pip install numba-vulkan        # or: uv add numba-vulkan
python -c "import numba_vulkan as nv; print(nv.list_devices())"

The wheel bundles libclc for accurate and float64 math. On macOS it also needs MoltenVK and an llc of LLVM 22 named by NUMBA_VULKAN_LLC (KI-32); see below.

To work on numba-vulkan itself: with devenv (NixOS or any system with Nix), everything including the Vulkan loader and SPIRV-Tools is provided by the shell:

devenv shell
uv run pytest

Without Nix, install the Vulkan loader from your distribution, then:

uv sync
uv run pytest

On macOS, Vulkan runs on MoltenVK, and llvmlite's SPIR-V backend has to be replaced by an llc of LLVM 22 (KI-32). The devenv shell sets up both (devenv shell -- uv run pytest); the test suite passes on an Apple M3 Pro. See Getting started for the details and a Homebrew alternative.

The test suite runs every test on every Vulkan device it finds.

Documentation

The Sphinx documentation is published at https://arunoruto.github.io/numba-vulkan/ and covers usage, the design, the limitations, the known issues and the benchmarks. To build it locally:

uv sync --group docs
cd docs && uv run sphinx-build -M html ./source ./build

(back to top)

Usage

A kernel looks like a numba.cuda kernel. global_id(axis) plays the role of cuda.grid, and forall launches it over a 1-d to 3-d grid:

import numpy as np
import numba_vulkan as nv

@nv.jit
def saxpy(a, x, y):
    i = nv.global_id(0)
    if i < x.shape[0]:
        y[i] = a * x[i] + y[i]

x = np.arange(1000, dtype=np.float32)
y = np.ones(1000, dtype=np.float32)
saxpy.forall(x.size)(np.float32(2.0), x, y)

NumPy arrays are copied to the device and back on every call. Device arrays stay there, as with numba.cuda:

dx, dy = nv.to_device(x), nv.to_device(y)
saxpy.forall(x.size)(np.float32(2.0), dx, dy)   # no copies
saxpy.forall(x.size)(np.float32(0.5), dx, dy)
y = dy.copy_to_host()

Functions compiled with @nv.jit can be called from kernels, may return values or tuples, and can take arrays. Functions compiled with numba.njit are recompiled for Vulkan when a kernel calls them:

import math
from numba import njit

@njit
def norm(x, y):
    return math.sqrt(x * x + y * y)

@nv.jit
def radius(x, y, out):
    i = nv.global_id(0)
    if i < out.size:
        out[i] = norm(x[i], y[i])

Numba's ufunc decorators work too, without writing a kernel:

from numba import vectorize, guvectorize

@vectorize(["float32(float32, float32)"], target="vulkan")
def hypot(a, b):
    return math.sqrt(a * a + b * b)

@guvectorize(["void(float32[:], float32[:], float32[:])"], "(n),(n)->()",
             target="vulkan")
def dot(a, b, out):
    out[0] = np.dot(a, b)

hypot(x, y)               # element-wise, with broadcasting
dot(matrix, vector)       # one dot product per row

Cooperating kernels use the same building blocks as in numba.cuda: workgroup-shared arrays, barriers, atomics and an explicit launch configuration:

@nv.jit
def block_sums(x, out):
    partial = nv.shared.array(256, np.float32)
    t, i = nv.local_id(0), nv.global_id(0)
    partial[t] = x[i] if i < x.shape[0] else np.float32(0)
    nv.barrier()
    step = 128
    while step > 0:
        if t < step:
            partial[t] += partial[t + step]
        nv.barrier()
        step //= 2
    if t == 0:
        nv.atomic.add(out, 0, partial[0])

block_sums[groups, 256](x, out)    # like kernel[blocks, threads] in CUDA

Because it is a real Numba target, @overload works with it:

from numba.extending import overload

@overload(my_function, target="vulkan")
def ol_my_function(x):
    return lambda x: x * 2

Choosing a device:

nv.list_devices()               # [<0: NVIDIA ... (discrete)>, <1: ...>, ...]
nv.select_device("intel")       # by name, index, or NUMBA_VULKAN_DEVICE
saxpy.forall(n, device="llvmpipe")(a, x, y)

Inspecting what was generated:

from numba import float32
compiled = saxpy.compile((float32, float32[::1], float32[::1]))
compiled.llvm_ir, compiled.spirv, compiled.capabilities

Set NUMBA_VULKAN_VALIDATE=1 to run every generated shader through spirv-val before it reaches a driver.

(back to top)

Benchmarks

Two suites compare numba-vulkan with Numba's CPU target and numba-cuda: the same scalar function compiled for every backend (benchmarks/bench.py), and the same CUDA-style kernels with shared memory, barriers and atomics (benchmarks/kernels.py). Results are collected in the repository, one file per run, and charted in the documentation, including every backend, device and variant and how results change over time. The fastest result per backend, GPUs with data on the device (log scale; further left is faster):

In the first collected run (Intel i9-9900K, NVIDIA TITAN X, Intel UHD 630):

• With data on the device, numba-vulkan takes 25 to 55 % longer than numba-cuda on the same card on most workloads, and is about level on the tiled matrix product (1.85 against 1.75 ms). Repeated runs vary by 10 to 20 %, so small differences are not meaningful.
• With NumPy arrays as arguments, the copies dominate: the memory-bound saxpy takes 7 to 10 ms on every GPU backend and 1.8 ms on one CPU thread.
• A kernel's first call, which compiles it, takes 0.04 to 0.25 s on Vulkan, comparable to numba-cuda.

On an Apple M3 Pro (MoltenVK 1.4.2, kernels translated by llc, KI-32), with data on the device:

• The compute-bound Mandelbrot set takes 7 to 9 ms, about level with the TITAN X; within a full run the GPU shares the chip's power budget with the CPU benchmarks before it, which costs up to 20 %.
• A launch that waits for its result costs about 0.2 ms, so the small workloads mostly measure that: saxpy takes 0.61 ms, about the same as the CPU.
• The tiled matrix product reaches only about 185 GFLOP/s (11.7 ms), the clearest candidate for tuning.

Please add a run from your machine; it takes a few minutes:

uv run python benchmarks/collect.py --user <your GitHub name>

(back to top)

What Works and What Does Not

Works:

• int32/int64/float32/float64/bool scalars and N-d arrays of bool, 8 to 64-bit integers, float32 and float64
• indexing with integers and slices; slices are views, as in NumPy, and can be iterated over, passed to functions, assigned to and reduced (a[i, 1:-1].sum(), for v in a[i], out[i, :3] = 0, np.dot, .T)
• arithmetic, comparisons, bit operations, casts, tuples, complex numbers, if/while/for ... in range(...), early returns, min/max/abs
• the math module in float32 and float64 (including hypot, log1p, erf, gamma, isnan...), accurate to the last digit and consistent across devices, and the ** operator
• NumPy functions on scalars, such as np.sqrt(x[i]) or np.maximum(a, b)
• calling other @nv.jit and @njit functions; @overload(target="vulkan")

Does not work:

• float64 math without libclc. Vulkan's math library is 32-bit only, so sin, exp, pow and friends in float64 come from libclc, LLVM's OpenCL math library. Wheels bundle it; in a source checkout without it such calls raise, unless you opt in to float32 precision with @nv.jit(narrow_math=True).
• Arrays of run-time size in kernels (a[mask], a.copy(), np.zeros(n)), try/except, recursion. Arrays of constant shape, array expressions such as a * 2 + b and reductions along an axis such as a.sum(axis=0) work; on the host, device arrays support masks and index arrays.
• float64 precision on devices without float64. Many mobile GPUs and Apple devices lack float64 or int64. Kernels are narrowed to 32-bit types there, with a warning; float narrowing is tested on an Apple M3 Pro, integer narrowing only on desktop GPUs with the feature switched off.
• Selecting between array elements by reference in ways the optimiser turns into pointer selects is rejected.

Things found along the way:

• NVIDIA's driver returns wrong results for 64-bit signed remainders with negative operands; remainders are computed from the quotient instead.
• NVIDIA's shader compiler segfaults on some invalid SPIR-V, which is why modules are checked before being handed to a driver.
• float32 sin on Intel/Mesa is only accurate to about 1e-4.
• LLVM's SPIR-V backend miscompiles nested aggregate inserts and aborts on several unsupported inputs.
• LLVM's SPIR-V structurizer emits invalid shaders for early returns and short-circuit conditions; control flow is restructured beforehand.
• All three drivers reassociate float arithmetic unless told not to, so every float operation is marked exact.
• LLVM's SPIR-V backend mistakes x = a * b; (x < 0) ? -y : y for GLSL's faceforward and crashes; such comparisons are rewritten.

(back to top)

Known Issues

The full list, with causes, workarounds and pointers to where each would be fixed, is in known issues. The ones most likely to bite:

tests/test_known_issues.py reproduces the coverage gaps as expected failures, so the list stays honest:

uv run pytest tests/test_known_issues.py -rxX

To continue the work, start with the development guide.

(back to top)

Roadmap

• Vulkan target registered with Numba's target extension API
• Kernels, device functions, overloads
• Test suite running on every available Vulkan device
• Sphinx documentation and benchmark suite
• Math functions, NumPy ufuncs on scalars, complex numbers
• Early returns and short-circuit conditions (control-flow restructuring)
• Loops with break/return anywhere, while loops (fuzz-tested)
• Device arrays and buffer reuse, to avoid copying on every call
• Asynchronous launches on device arrays (nv.synchronize())
• Indexing (basic and advanced), views and NumPy ufuncs on device arrays
• Structured (record) arrays
• Autotuning of kernel variants (nv.autotune)
• Device-side timing with events (nv.event)
• Streams overlapping copies and kernels (nv.stream, nv.pinned_array)
• Launch plans: repeated launches skip typing
• Scalars as push constants
• Slices, array views, iteration and reductions, also along an axis
• Array expressions and local arrays
• Kernels narrowed to 32-bit types on devices without float64/int64
• float64 math, through libclc
• Bundle libclc in wheels
• Publish to PyPI
• numba.vectorize / numba.guvectorize with target="vulkan"
• reduce of vectorize functions
• print() in kernels, float16 arrays, error_model="python"
• Other ufunc methods (accumulate, outer)
• Shared memory, atomics, barriers and CUDA-style launch configuration
• Subgroup operations: reductions, scans, votes, ballots, shuffles
• On-disk caching of compiled kernels
• cache=True: compiled kernels reused without recompiling
• Testing on Apple (MoltenVK)
• Testing on AMD and mobile GPUs

(back to top)

Authorship and AI Disclosure

Every line here was generated by a language model in October 2026; the human author steered the work and reviewed it but did not write the code. Many implementation-level choices were made by the model while solving problems it ran into, and should be judged on their merits. Benchmark numbers and hardware observations come from real runs on the author's machine.

This is an independent project and is not endorsed by the Numba maintainers. Numba has an AI tools policy that requires attributing LLM-generated contributions; it applies if any part of this is ever proposed upstream.

(back to top)

Contributing

Issues and pull requests are welcome, in particular reports of how the test suite behaves on other GPUs and drivers:

uv run pytest
uv run python benchmarks/bench.py

Contributions written with AI tools are welcome on the same terms as the project itself: say which tool produced them, and have a human review them.

(back to top)

License

Distributed under the BSD 2-Clause License, the same license as Numba. See LICENSE.

(back to top)

Contact

Mirza Arnaut - @arunoruto

(back to top)

Acknowledgments

• Numba and numba-cuda, whose CUDA target is the template for this one
• Taichi, for showing that Python kernels on Vulkan work
• The LLVM SPIR-V backend developers
• Best-README-Template

(back to top)