Certified kernels instead of autotuning on your machine

The first time you run a Triton-based model, it often stalls. It is not loading weights; it is autotuning: trying block sizes, warp counts and pipeline stages for each kernel and shape, timing them, keeping the fastest. It happens on your machine, every time a new shape shows up, and the result depends on what else your GPU was doing at that moment.

NeuroBrix does it differently. Kernel settings are certified once per GPU class and shipped with the package.

How a setting is certified

1. The launch key. For each model, the kernels it will launch and their shape classes are listed from its graph, without running the model.
2. The float64 oracle. Every candidate setting for a key is run and checked against a float64 reference. A setting that computes the wrong answer is excluded, however fast it is.
3. The timing. The settings that pass are timed on the GPU class, under a witness that refuses the measurement if the GPU's regime drifts while candidates are being compared.
4. The certificate. The winner is stored with its proof — the deviation from the oracle, the time, the compiler that produced it — in the repository.
5. Served. At run time the engine reads the certificate. No sweep.

The compiler is part of the proof. A certificate made under one Triton build is not served under another: it is re-proven first. The configuration is a property of the compiler that produced it.

What it changes for you

• A first run that runs. No minutes of tuning before the first token or the first frame.
• Reproducible behaviour. The same GPU class gets the same settings, so two users comparing results are comparing the same kernels.
• Correctness first. Every served setting has been checked against float64 before it was ever timed.

18,552 settings are certified for Apple GPUs today.

What you see at run time

Every run says, in one line, where its kernel settings came from. This is from a real run of Sana 1600M on an Apple M4 Pro, installed from PyPI:

[autotune] certified directory: 70 key(s) served without a sweep, 0 swept at runtime (kept locally), 0 from the local replay cache

Seventy launch keys, all served from certificates, none tuned on the machine. If a key were missing, the same line would say how many were swept and kept locally, so you always know.

Your GPU, your certificate

A GPU class with no certificate still works: the missing settings are tuned once on your machine and kept locally. You can then certify them with the same tool NeuroBrix uses and share them by pull request, so nobody with your GPU tunes them again:

neurobrix run --model TinyLlama-1.1B-Chat-v1.0 --triton --prompt "Hello" --max-tokens 32
neurobrix autotune certify --profile <your-gpu-profile> --keys <the keys your run recorded> --out ./my-certificates

The full procedure, including how to check that a certified-only run serves everything, is in the nbxtensor README.

Links

• github.com/NeuroBrix/neurobrix
• github.com/NeuroBrix/nbxtensor

The way we build it

This post is one piece of a consistent method:

• One engine, zero per-model code. The runtime knows tensors, graphs and plans; it never knows which model it is running. A fix in the engine is a fix for every model.
• .nbx, one container for every model: graph and weights in one file format.
• Prism decides from facts: placement chosen from a hardware profile, not from heuristics in code.
• Certified, not autotuned: every kernel setting checked against a float64 oracle, timed once per GPU class, shared by pull request.
• Validated through the real entry point: all 49 models on both engines at every release (196 of 196 cells for 0.6.0 on 4x V100), outputs compared byte for byte with judged outputs, runs deterministic from one repetition to the next.

The ecosystem

• NeuroBrix Engine: the inference engine, pip install neurobrix.
• nbxtensor: the tensor library without PyTorch, pip install nbxtensor.
• Model Hub: .nbx containers for 49 models at neurobrix.org/models, more with every release.
• triton-ext: our Triton build for Apple GPUs.
• Studio: a desktop application for Linux, Windows and macOS, in development.

Where we are going

Faster than the frameworks (kernel fusion, graph-level simplification), every processor (AMD and Intel GPUs, many-core CPUs, ARM64, accelerators beyond GPUs, multi-GPU across NVLink, PCIe and networks), community-certified hardware, training in nbxtensor, local agents, a local server and a JSON interface for applications, Studio for everyone, nbxtensor 2.0 in Rust, and .nbx as an open model standard.

Get involved

• Star github.com/NeuroBrix/neurobrix and github.com/NeuroBrix/nbxtensor.
• Try it: pip install neurobrix.
• Certify your GPU class and share it (how).
• Open an issue when something breaks, or to request a model.

Hocine Benkelaya, NeuroBrix