nbxtensor — the torch-free tensor library for inference

What it is

An inference-first tensor library, torch-free. The graph and the shapes are known before the first run, so memory is planned rather than discovered, every Triton kernel declares its cost, and a repeated sequence of launches is captured and replayed. nbxtensor never imports torch; exchange goes through DLPack. NeuroBrix's Triton engine is its first client.

Backends: NVIDIA CUDA and Apple Metal (through triton-ext). Python 3.10 or later.

Fast: tensor cores back on the V100

Still running V100s? Triton leaves their tensor cores idle; nbxtensor turns them back on. On sm_70 Triton compilestl.dotto plain FMA, with no matrix instruction in the PTX. nbxtensor ships its own mma.m8n8k4 kernels for matrix products, convolution and the fused mixture-of-experts GEMM, and a kernel runs on the matrix unit only where its certificate measured it faster.

Kernel timings on a V100, FMA against the matrix unit inside nbxtensor, as recorded in the 0.0.1 changelog. These timings compare nbxtensor's own kernels with each other.

Faster than PyTorch on many models.

Lightweight

Lightweight: nbxtensor itself is 1 MB; the whole install, Triton included, is about 1 GB — PyTorch's is 5.6 GB.

Measured on a fresh container, Python 3.10.

Devices

Runs wherever Triton runs — NVIDIA, AMD, Intel and Apple Silicon GPUs, on Linux, Windows and macOS.

Found a bug on your hardware? Open an issue— we fix fast.

There is no CPU kernel backend in this version; host work (loading, conversions) runs on every core with numpy.

Certified tuning

Kernel settings measured once per GPU class and shipped — no autotuning at run time for the shapes we cover. We are certifying every model we can, so you never have to autotune.

Certified your GPU? Share your certificates.

Example

from nbxtensor.dtype import itemsize as I

assert I.itemsize("bfloat16") == 2
assert I.itemsize("complex64") == 8
assert I.kind("bfloat16") == "float" and I.kind("uint8") == "uint"
assert I.from_safetensors("BF16") == "bfloat16"
try:
    I.itemsize("fp8_unknown")     # a name the table does not carry is refused by name
except ValueError as e:
    assert "unknown dtype" in str(e)

From the README; every README example is run by the library's test suite.

Known limits

• Inference only: no autograd.
• Slicing has step 1; a strided slice is refused by name.
• No array namespace (__array_namespace__) yet.
• conv_transpose2d runs on the FMA kernel only (no matrix-unit path).
• On Apple silicon, mochi-1-preview and Allegro / Allegro-TI2V do not complete on the M4 Pro yet.

Latest version Feed