Hocine Benkelaya

Posted on Nov 7

nbxtensor: a tensor library for inference that doesn't need PyTorch

#ai #python #machinelearning #opensource

If you want to run a modern model on a GPU from Python, the default answer is "install PyTorch". That is a fine answer for training. For inference it brings a lot along: a large install, a CUDA-first view of the world, and a runtime whose memory behaviour you discover rather than decide.

nbxtensor is NeuroBrix's answer for inference: a tensor library with Triton kernels, planned memory and certified tuning — and no PyTorch.

pip install nbxtensor

It runs wherever Triton runs: NVIDIA, AMD, Intel and Apple GPUs, on Linux, Windows and macOS.

What "no PyTorch" means here

Not "PyTorch is optional". nbxtensor never imports torch — not as a dependency, not as an optional extra, not in a fallback branch. Its kernels, its tests and the models NeuroBrix runs on it all work in an environment where torch is not installed.

The practical effect shows up in the NeuroBrix engine that is built on it: measured on a clean machine, the light install (the Triton engine on nbxtensor) is 1.48 GB; the same engine with PyTorch added is 6.82 GB.

If you do live in a PyTorch codebase, exchange goes through DLPack, so you are not locked in either direction.

A first kernel

This is the example from the README, executed by the test suite on every change:

import numpy as np
from nbxtensor.tensor import NBXTensor
from nbxtensor.ops.linalg_wrappers import mm

rng = np.random.default_rng(0)
a = rng.standard_normal((256, 512)).astype(np.float16)
b = rng.standard_normal((512, 128)).astype(np.float16)

c = mm(NBXTensor.from_numpy(a), NBXTensor.from_numpy(b))   # a Triton kernel on the GPU, no torch

ref = a.astype(np.float32) @ b.astype(np.float32)
assert c.shape == (256, 128)
assert np.allclose(c.numpy().astype(np.float32), ref, rtol=2e-2, atol=2e-1)

NBXTensor holds device memory, a shape and a dtype. The operation dispatches to a Triton kernel. There is nothing else in the path.

Four parts

1. Tensors. Device memory, shapes, strides, views, dtypes and their promotion rules — the plumbing you need to run a graph, without a framework underneath.

2. Triton kernels. Matrix products, attention, convolutions, normalisations, elementwise and reductions, written in Triton so that one source serves every GPU vendor Triton supports. A device is described by a file of facts — tile sizes, shared memory, the matrix unit — not by code.

3. Planned memory. Allocators per device behind one interface, and a planner that sizes buffers before the first run. Every kernel declares what it costs (its outputs, workspace and scratch), so the planner can price a model before running it. In NeuroBrix, that planner is Prism: it decides from your hardware profile whether a model fits on one GPU, is split across several, is streamed, or has its largest operations tiled.

4. Certified tuning. A Triton kernel has settings — block sizes, warps, stages — and the right ones depend on the GPU and the shape. The usual approach is to autotune on the user's machine at first run. nbxtensor instead serves settings that were certified once per GPU class: each candidate is checked against a float64 oracle, the ones that pass are timed, and the winner is stored with its proof. 18,552 settings are certified for Apple GPUs today.

If your GPU class has no certificate yet, the missing settings are tuned once on your machine and kept; you can then certify them and share them by pull request, so the next person with your GPU never tunes at all.

Fast where it counts

Two examples of what that buys:

In NeuroBrix's global benchmark of 2026-11-05 on V100 16 GB GPUs, the Triton engine on nbxtensor was faster than the PyTorch engine on 18 of 49 models, up to 5x on Janus-Pro-7B (212.6 s down to 41.8 s).
On Volta, Triton compiles matrix products to FMA, so the V100's tensor cores sit idle. nbxtensor brings them back with hand-written m8n8k4 kernels, chosen per operation by measurement rather than by rule.

Apple Silicon

PyPI publishes no Triton for macOS. nbxtensor installs NeuroBrix's own Triton build for Apple GPUs (NeuroBrix/triton-ext) in one command:

pip install nbxtensor
python -m nbxtensor install-metal
python -m nbxtensor doctor     # -> backend in force: metal

doctor works on every platform: it names the backend in force and, when a GPU runtime or a Triton backend is missing, the command that installs it.

Where it is used

nbxtensor is the foundation of the Triton engine of NeuroBrix, which runs 49 models in 9 families — language, code, vision-language, omni, speech recognition, speech understanding, text-to-speech, image, video and upscaling — from one .nbx container format and one entry point. Every NeuroBrix release runs all 49 models on both engines through that entry point and compares the outputs byte for byte with judged outputs.

Inference first: training is not in this release.

DEV Community